#include <iostream>
#include <fstream>
#include <iomanip> // For setting output format
#include "TFile.h"
#include "TTree.h"
#include "TF1.h"
#include "TCanvas.h"
#include "TGraph.h"
#include <vector>


void processDataset( const std::string& label, std::vector<double>& channel, std::vector<double>& energy) {
    // Number of calibration points
    const int nPoints = channel.size();

    // Convert vectors to arrays for TGraph
    double* channelArray = channel.data();
    double* energyArray = energy.data();

    // Channel range for fitting
    double channelMin = channelArray[0];
    double channelMax = channelArray[nPoints - 1];

    // Create a TGraph object
    TGraph *graph = new TGraph(nPoints, channelArray, energyArray);
  
    // Set marker style and size for calibration points
    graph->SetMarkerStyle(20); // Round markers
    graph->SetMarkerSize(1.2); // Marker size
  
    // Fit function: cubic plus square root equation with 5 parameters
    TF1 *fitFunc = new TF1("fitFunc", "[0] + [1] * x + [2] * x^2 + [3] * x^3 + [4] * sqrt(x)", channelMin, channelMax);

    // Set initial values for the fit parameters
    fitFunc->SetParameter(0, 1.0); // Parameter for constant term
    fitFunc->SetParameter(1, 1.0); // Parameter for x
    fitFunc->SetParameter(2, 1.0); // Parameter for x^2
    fitFunc->SetParameter(3, 1.0); // Parameter for x^3
    fitFunc->SetParameter(4, 1.0); // Parameter for sqrt(x)

    // Perform the fit
    graph->Fit("fitFunc", "Q");

    // Get fit parameters
    double A = fitFunc->GetParameter(0);
    double B = fitFunc->GetParameter(1);
    double C = fitFunc->GetParameter(2);
    double D = fitFunc->GetParameter(3);
    double E = fitFunc->GetParameter(4);
    
    // Get chi-square value
    double chiSquare = fitFunc->GetChisquare();

    // Create a canvas and draw the graph with fit
    TCanvas *canvas = new TCanvas("canvas", "Calibration Plot", 800, 600);
    graph->SetTitle("Calibration Plot");
    graph->GetXaxis()->SetTitle("Channel");
    graph->GetYaxis()->SetTitle("Energy");
    graph->Draw("AP"); // Draw points with axes
    fitFunc->Draw("same"); // Draw fit curve on top of points
    canvas->Draw();

    // Display fit parameters in the specified order without capital "E" notation
    std::cout << "Fit Parameters (in order):" << std::endl;
    std::cout << "A: " << std::fixed << std::setprecision(6) << A << std::endl;
    std::cout << "B: " << std::fixed << std::setprecision(6) << B << std::endl;
    std::cout << "C: " << std::fixed << std::setprecision(6) << C << std::endl;
    std::cout << "D: " << std::fixed << std::setprecision(6) << D << std::endl;
    std::cout << "E: " << std::fixed << std::setprecision(6) << E << std::endl;
    std::cout << "Chi-Square: " << std::fixed << std::setprecision(6) << chiSquare << std::endl;

    // Save parameters in parameters.dat with default notation
    std::ofstream paramFile("parameters.dat", std::ios_base::app);
    if (paramFile.is_open()) {
        paramFile << std::scientific << std::setprecision(6);
        paramFile << A << " " << B << " " << C << " " << D << " " << E << " " << chiSquare << std::endl;
        paramFile.close();
        std::cout << "Fit parameters saved in parameters.dat." << std::endl;
    } else {
        std::cerr << "Error: Could not open parameters.dat for writing." << std::endl;
    }
}

void calibration_macro() {
    // Ask the user for the input .dat file
    std::string inputFile;
    std::cout << "Enter the name of the input .dat file: ";
    std::cin >> inputFile;

    // Read channel and energy data from the .dat file
    std::ifstream infile(inputFile);
    if (!infile) {
        std::cerr << "Error: Could not open the input file." << std::endl;
        return;
    }

    // Variables to hold channel and energy data for one dataset
    std::vector<double> channel;
    std::vector<double> energy;
    double channelValue, energyValue;
    std::string label;
    std::string prevLabel;

    while (infile >> label >> channelValue >> energyValue) {
        // Check if the label is the same as the previous one
        if (label != prevLabel) {
            // Clear vectors if it's a new dataset
            channel.clear();
            energy.clear();
            prevLabel = label; // Update previous label
        }
        channel.push_back(channelValue);
        energy.push_back(energyValue);

        if (label != prevLabel) {
            // Process the dataset when label changes
            processDataset(label, channel, energy);
        }
    }

    infile.close();
}