#include "opatIO.h"
#include <fstream>
#include <iostream>
#include <stdexcept>
#include <cstring>
#include <algorithm>
#include <iomanip>
#include <map>
#include <utility>
#include <cmath>
#include <limits>
#include <deque>

// Function to check system endianness
bool is_big_endian() {
    uint16_t test = 0x1;
    return reinterpret_cast<uint8_t*>(&test)[0] == 0;
}

// Generic function to swap bytes for any type
template <typename T>
T swap_bytes(T value) {
    static_assert(std::is_trivially_copyable<T>::value, "swap_bytes only supports trivial types.");
    T result;
    uint8_t* src = reinterpret_cast<uint8_t*>(&value);
    uint8_t* dest = reinterpret_cast<uint8_t*>(&result);
    for (size_t i = 0; i < sizeof(T); i++) {
        dest[i] = src[sizeof(T) - 1 - i];
    }
    return result;
}


// Constructor
OpatIO::OpatIO() {}

OpatIO::OpatIO(std::string filename) : filename(filename) {
    load();
}

// Destructor
OpatIO::~OpatIO() {
    unload();
}

// Load the OPAT file
void OpatIO::load() {
    if (loaded) return;

    std::ifstream file(filename, std::ios::binary);
    if (!file.is_open()) {
        throw std::runtime_error("Could not open file: " + filename);
    }

    readHeader(file);
    readTableIndex(file);

    loaded = true;
    file.close();
}

// // Unload the OPAT file
void OpatIO::unload() {
    if (!loaded) return;

    tableIndex.clear();
    while (!tableQueue.empty()) {
        tableQueue.pop_front();
    }

    loaded = false;
}

// Read the header from the file
void OpatIO::readHeader(std::ifstream &file) {
    file.read(reinterpret_cast<char*>(&header), sizeof(Header));
    if (file.gcount() != sizeof(Header)) {
        throw std::runtime_error("Error reading header from file: " + filename);
    }

    if (is_big_endian()) {
        header.version = swap_bytes(header.version);
        header.numTables = swap_bytes(header.numTables);
        header.indexOffset = swap_bytes(header.indexOffset);
    }
}

// Read the table index from the file
void OpatIO::readTableIndex(std::ifstream &file) {
    file.seekg(header.indexOffset, std::ios::beg);
    tableIndex.resize(header.numTables);
    file.read(reinterpret_cast<char*>(tableIndex.data()), header.numTables * sizeof(TableIndex));
    if (file.gcount() != static_cast<std::streamsize>(header.numTables * sizeof(TableIndex))) {
        throw std::runtime_error("Error reading table index from file: " + filename);
    }

    buildTableIDToComposition();
}

void OpatIO::buildTableIDToComposition(){
    tableIDToComposition.clear();
    int tableID = 0;
    std::pair<double, double> comp;
    for (const auto &index : tableIndex) {
        comp.first = index.X;
        comp.second = index.Z;
        tableIDToComposition.emplace(tableID, comp);
        tableID++;
    }
    XZLookupEpsilon();
}

void OpatIO::XZLookupEpsilon() {
    /* 
     Get 10% of the minimum spacing between XZ values 
     in the tableIDToComposition map. This can be used
     to set the comparison distance when doing a reverse
     lookup (composition -> tableID)
     */
    std::vector<double> Xvalues, Zvalues;
    double epsilonX, epsilonZ, xgap, zgap;

    // Start these out as larger than they will ever be
    epsilonX = 1;
    epsilonZ = 1;
    for (const auto& pair : tableIDToComposition) {
        Xvalues.push_back(pair.second.first);
        Zvalues.push_back(pair.second.second);
    }

    // Sorting is required for this algorithm.
    std::sort(Xvalues.begin(), Xvalues.end());
    std::sort(Zvalues.begin(), Zvalues.end());

    for (size_t i = 1; i < Xvalues.size(); ++i) {
        xgap = Xvalues[i] - Xvalues[i - 1];
        zgap = Zvalues[i] - Zvalues[i - 1];
        if (xgap > 0 && xgap < epsilonX) {
            epsilonX = xgap;
        }
        if (zgap > 0 && zgap < epsilonZ) {
            epsilonZ = zgap;
        }
    }
    // 0.1 to extract 10% of min distance.
    XZepsilon = {0.1*epsilonX, 0.1*epsilonZ};
}

int OpatIO::lookupTableID(double X, double Z){
    bool XOkay;
    bool ZOkay;
    int tableID = 0;
    for (const auto &tableMap : tableIDToComposition){
        XOkay = std::fabs(tableMap.second.first - X) < XZepsilon.first;
        ZOkay = std::fabs(tableMap.second.second - Z) < XZepsilon.second;
        if (XOkay and ZOkay){
            return tableID;
        }
        tableID++;
    }
    return -1;
}

// Get a table from the queue
OPATTable OpatIO::getTableFromQueue(int tableID) {
    for (const auto &table : tableQueue) {
        if (table.first == tableID) {
            return table.second;
        }
    }
    throw std::out_of_range("Table not found!");
}

// Add a table to the queue
void OpatIO::addTableToQueue(int tableID, OPATTable table) {
    if (static_cast<int>(tableQueue.size()) >= maxQDepth) {
        removeTableFromQueue();
    }
    std::pair<int, OPATTable> IDTablePair = {tableID, table};
    tableQueue.push_back(IDTablePair);
}

// Remove a table from the queue
void OpatIO::removeTableFromQueue() {
    if (!tableQueue.empty()) {
        tableQueue.pop_front();
    }
}

// Flush the queue
void OpatIO::flushQueue() {
    while (!tableQueue.empty()) {
        tableQueue.pop_back();
        tableQueue.pop_front();
    }
}

// Get the OPAT version
uint16_t OpatIO::getOPATVersion() {
    return header.version;
}

// Get a table for given X and Z
OPATTable OpatIO::getTable(double X, double Z) {
    int tableID = lookupTableID(X, Z);
    if (tableID == -1) {
        throw std::out_of_range("X Z Pair Not found!");
    }
    try {
        return getTableFromQueue(tableID);
    }
    catch(const std::out_of_range &e) { 
        return getTable(tableID);
    }
}

OPATTable OpatIO::getTable(int tableID) {
    std::ifstream file(filename, std::ios::binary);
    if (!file.is_open()) {
        throw std::runtime_error("Could not open file: " + filename);
    }

    uint64_t byteStart = tableIndex[tableID].byteStart;

    file.seekg(byteStart, std::ios::beg);

    OPATTable table;

    // Step 1: Read N_R and N_T
    file.read(reinterpret_cast<char*>(&table.N_R), sizeof(uint32_t));
    file.read(reinterpret_cast<char*>(&table.N_T), sizeof(uint32_t));
    

    // Resize vectors to hold the correct number of elements
    table.logR.resize(table.N_R);
    table.logT.resize(table.N_T);
    table.logKappa.resize(table.N_R, std::vector<double>(table.N_T));

    // Step 2: Read logR values
    file.read(reinterpret_cast<char*>(table.logR.data()), table.N_R * sizeof(double));

    // Step 3: Read logT values
    file.read(reinterpret_cast<char*>(table.logT.data()), table.N_T * sizeof(double));

    // Step 4: Read logKappa values (flattened row-major order)
    for (size_t i = 0; i < table.N_R; ++i) {

        file.read(reinterpret_cast<char*>(table.logKappa[i].data()), table.N_T * sizeof(double));
    }

    if (!file) {
        throw std::runtime_error("Error reading table from file: " + filename);
    }

    addTableToQueue(tableID, table);
    file.close();
    return table;
}

// Set the maximum queue depth
void OpatIO::setMaxQDepth(int depth) {
    maxQDepth = depth;
}

int OpatIO::getMaxQDepth() {
    return maxQDepth;
}

// Set the filename
void OpatIO::setFilename(std::string filename) {
    if (loaded) {
        throw std::runtime_error("Cannot set filename while file is loaded");
    }
    this->filename = filename;
}

// Check if the file is loaded
bool OpatIO::isLoaded() {
    return loaded;
}

// Print the header
void OpatIO::printHeader() {
    std::cout << "Version: " << header.version << std::endl;
    std::cout << "Number of Tables: " << header.numTables << std::endl;
    std::cout << "Header Size: " << header.headerSize << std::endl;
    std::cout << "Index Offset: " << header.indexOffset << std::endl;
    std::cout << "Creation Date: " << header.creationDate << std::endl;
    std::cout << "Source Info: " << header.sourceInfo << std::endl;
    std::cout << "Comment: " << header.comment << std::endl;
}

// Print the table index
void OpatIO::printTableIndex() {
    if (tableIndex.empty()) {
        std::cout << "No table indexes found." << std::endl;
        return;
    }

    // Print table header
    std::cout << std::left << std::setw(10) << "X"
        << std::setw(10) << "Z"
        << std::setw(15) << "Byte Start"
        << std::setw(15) << "Byte End"
        << "Checksum (SHA-256)" << std::endl;

    std::cout << std::string(80, '=') << std::endl;  // Separator line

    // Print each entry in the table
    for (const auto &index : tableIndex) {
        std::cout << std::fixed << std::setprecision(4)
            << std::setw(10) << index.X
            << std::setw(10) << index.Z
            << std::setw(15) << index.byteStart
            << std::setw(15) << index.byteEnd
            << std::hex; // Switch to hex mode for checksum

        for (int i = 0; i < 8; ++i) {  // Print first 8 bytes of SHA-256 for brevity
            std::cout << std::setw(2) << std::setfill('0') << (int)index.sha256[i];
        }

        std::cout << "..." << std::dec << std::setfill(' ') << std::endl;  // Reset formatting
    }
}

void OpatIO::printTable(OPATTable table, uint32_t truncateDigits) {
    int printTo;
    bool truncate = false;
    if (table.N_R > truncateDigits) {
        printTo = truncateDigits;
        truncate = true;
    } else {
        printTo = table.N_R-1;
    }
    std::cout << "LogR (size: " << table.logR.size() << "): [";
    for (int i = 0; i < printTo; ++i) {
        std::cout << table.logR.at(i) << ", ";
    }
    if (truncate) {
        std::cout << "..., ";
        for (int i = truncateDigits; i > 1; --i) {
            std::cout << table.logR.at(table.logR.size() - i) << ", ";
        }
    }
    std::cout << table.logR.back() << "]" << std::endl;

    if (table.N_T > truncateDigits) {
        printTo = truncateDigits;
        truncate = true;
    } else {
        printTo = table.N_T-1;
    }
    std::cout << "LogT (size: " << table.logT.size() << "): [";
    for (int i = 0; i < printTo; ++i) {
        std::cout << table.logT.at(i) << ", ";
    }
    if (truncate) {
        std::cout << "..., ";
        for (int i = truncateDigits; i > 1; --i) {
            std::cout << table.logT.at(table.logT.size() - i) << ", ";
        }
    }
    std::cout << table.logT.back() << "]" << std::endl;

    bool truncateRow = false;
    bool truncateCol = false;
    int printToRow, printToCol;
    if (table.N_T > truncateDigits) {
        printToRow = truncateDigits;
        truncateRow = true;
    } else {
        printToRow = table.N_T-1;
    }
    if (table.N_R > truncateDigits) {
        printToCol = truncateDigits;
        truncateCol = true;
    } else {
        printToCol = table.N_R-1;
    }

    std::cout << "LogKappa (size: " << table.N_R << " x " << table.N_T << "): \n[";
    for (int rowIndex = 0; rowIndex < printToRow; rowIndex++) {
        std::cout << "[";
        for (int colIndex = 0; colIndex < printToCol; colIndex++) {
            std::cout << table.logKappa.at(rowIndex).at(colIndex) << ", ";
        }
        if (truncateRow) {
            std::cout << "..., ";
            for (int i = truncateDigits; i > 1; i--) {
                std::cout << table.logKappa.at(rowIndex).at(table.logKappa.at(rowIndex).size() - i) << ", ";
            }
        }
        std::cout << table.logKappa.at(rowIndex).back() << "],\n";
    } 
    if (truncateCol) {
        std::cout << ".\n.\n.\n";
        for (int rowIndex = truncateDigits; rowIndex > 1; rowIndex--) {
            std::cout << "[";
            for (int colIndex = 0; colIndex < printToCol; colIndex++) {
                std::cout << table.logKappa.at(rowIndex).at(colIndex) << ", ";
            }
            if (truncateRow) {
                std::cout << "..., ";
                for (int i = truncateDigits; i > 1; i--) {
                    std::cout << table.logKappa.at(rowIndex).at(table.logKappa.at(rowIndex).size() - i) << ", ";
                }
            }
            std::cout << table.logKappa.at(rowIndex).back() << "],\n";
        } 
        std::cout << "[";
        for (int colIndex = 0; colIndex < printToCol; colIndex++) {
            std::cout << table.logKappa.back().at(colIndex) << ", ";
        }
        if (truncateRow) {
            std::cout << "..., ";
            for (int i = truncateDigits; i > 1; i--) {
                std::cout << table.logKappa.back().at(table.logKappa.back().size() - i) << ", ";
            }
        }
        std::cout << table.logKappa.back().back() << "]";
    }
    std::cout << "]" << std::endl;
}

void OpatIO::printTable(double X, double Z, uint32_t truncateDigits) {
    int tableID = lookupTableID(X, Z);
    OPATTable table = getTable(tableID);
    printTable(table, truncateDigits);
}

// Get the table index
std::vector<TableIndex> OpatIO::getTableIndex() {
    return tableIndex;
}

// Get the header
Header OpatIO::getHeader() {
    return header;
}

// // Get the closest X tables
// std::vector<OPATTable> OpatIO::getClosestXTables(double X, double ZExact, int numTables) {
//     std::vector<OPATTable> closestTables;
//     // Implement logic to find closest X tables
//     return closestTables;
// }

// // Get the closest Z tables
// std::vector<OPATTable> OpatIO::getClosestZTables(double XExact, double Z, int numTables) {
//     std::vector<OPATTable> closestTables;
//     // Implement logic to find closest Z tables
//     return closestTables;
// }

// // Get the closest tables
// std::vector<OPATTable> OpatIO::getClosestTables(double X, double Z, int numTables) {
//     std::vector<OPATTable> closestTables;
//     // Implement logic to find closest tables
//     return closestTables;
// }