#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <cstdint>

// ROOT Headers
#include <TFile.h>
#include <TH2.h>

void root2mat(const std::string& inputRoot, const std::string& histName, const std::string& outputMat) {
    constexpr int X_SIZE = 4096;
    constexpr int Y_SIZE = 4096;
    constexpr size_t TOTAL_BINS = static_cast<size_t>(X_SIZE) * Y_SIZE;

    // Open ROOT file
    TFile* inFile = TFile::Open(inputRoot.c_str(), "READ");
    if (!inFile || inFile->IsZombie()) {
        std::cerr << "Error: Could not open input ROOT file " << inputRoot << std::endl;
        return;
    }

    // Get 2D histogram
    TH2* h2 = dynamic_cast<TH2*>(inFile->Get(histName.c_str()));
    if (!h2) {
        std::cerr << "Error: Could not find TH2 histogram named '" << histName << "' in file." << std::endl;
        inFile->Close();
        delete inFile;
        return;
    }

    // Determine target format and integer depth based on output extension
    bool is32Bit = false;
    if (outputMat.length() >= 4 && outputMat.substr(outputMat.length() - 4) == ".spn") {
        is32Bit = true;
    } else if (outputMat.length() >= 4 && outputMat.substr(outputMat.length() - 4) == ".mat") {
        is32Bit = false;
    } else {
        std::cerr << "Error: Unsupported output file extension. Must be .mat or .spn" << std::endl;
        inFile->Close();
        delete inFile;
        return;
    }

    std::cout << "Extracting '" << histName << "' into " << X_SIZE << " x " << Y_SIZE 
              << " binary matrix (" << (is32Bit ? "32-bit .spn" : "16-bit .mat") << ")..." << std::endl;

    // Open output binary file
    std::ofstream outFile(outputMat, std::ios::binary);
    if (!outFile.is_open()) {
        std::cerr << "Error: Could not create output file " << outputMat << std::endl;
        inFile->Close();
        delete inFile;
        return;
    }

    // Extract bin contents row-by-row matching column-major order (x * Y_SIZE + y)
    if (is32Bit) {
        std::vector<int32_t> buffer(TOTAL_BINS);
        for (int x = 0; x < X_SIZE; ++x) {
            for (int y = 0; y < Y_SIZE; ++y) {
                // ROOT bins are 1-indexed
                buffer[x * Y_SIZE + y] = static_cast<int32_t>(h2->GetBinContent(x + 1, y + 1));
            }
        }
        outFile.write(reinterpret_cast<const char*>(buffer.data()), TOTAL_BINS * sizeof(int32_t));
    } else {
        std::vector<int16_t> buffer(TOTAL_BINS);
        for (int x = 0; x < X_SIZE; ++x) {
            for (int y = 0; y < Y_SIZE; ++y) {
                // ROOT bins are 1-indexed[cite: 1]
                buffer[x * Y_SIZE + y] = static_cast<int16_t>(h2->GetBinContent(x + 1, y + 1));
            }
        }
        outFile.write(reinterpret_cast<const char*>(buffer.data()), TOTAL_BINS * sizeof(int16_t));
    }

    outFile.close();
    inFile->Close();
    delete inFile;

    std::cout << "Successfully written matrix to " << outputMat << std::endl;
}

int main(int argc, char* argv[]) {
    if (argc < 4) {
        std::cout << "Usage: " << argv[0] << " <input.root> <hist_name> <output.mat/.spn>\n";
        return 1;
    }

    root2mat(argv[1], argv[2], argv[3]);
    return 0;
}