#include <cstdlib>
#include <iostream>
#include <fstream>
#include <vector>
#include <math.h>
#include <random>

#include <TCanvas.h>
#include <TROOT.h>
#include <TApplication.h>
#include <TH1D.h>
#include <TGraph.h>
#include <TTree.h>
#include <TFile.h>
#include <TSystem.h>

#include "Garfield/ComponentFieldMap.hh"
#include "Garfield/ComponentGrid.hh"

#include "Garfield/MediumMagboltz.hh"
#include "Garfield/Medium.hh"
#include "Garfield/MediumConductor.hh"

#include "Garfield/Sensor.hh"
#include "Garfield/AvalancheMicroscopic.hh"
#include "Garfield/AvalancheMC.hh"
#include "Garfield/DriftLineRKF.hh"

#include "Garfield/ViewSignal.hh"
#include "Garfield/ViewField.hh"
#include "Garfield/ViewDrift.hh"

#include "Garfield/Random.hh"

#include "Garfield/ComponentComsol.hh"

using namespace Garfield;

int main(int argc, char * argv[]) {

    // This program is to simulate drifts, avalanches and x-ray events in a micromegas chamber

    //Directory of COMSOL field files
    std::string fileloc = "fields_new/";

    // Specific field files
    std::string meshfile = "mesh_smaller_new.mphtxt";
    std::string field = "field_smaller_new700.txt";
    std::string materials = "materials1.dat";
    std::string gas_file = "build/Ar_90_co2_10_0.0T_293.15K.gas";

    // Directory of output CSV file
    std::string out_file = "Output700/MicromegasMain.csv";
    std::string out_signal = "Output700/SignalCurrent.csv";
    std::string out_charge = "Output700/SignalCharge.csv";

    // Specify number of readout strips in single axis
    int numReadoutstrips = 5;


    // Assign directories 
    meshfile = fileloc + meshfile;
    field = fileloc + field;
    materials = fileloc + materials;

    TApplication app("app", &argc, argv);

    // Import field map

    std::cout<<"Beginning to read COMSOL"<<"\n";

    ComponentComsol fm(meshfile, materials, field, "cm");
    fm.EnableConvergenceWarnings(false);

    std::cout<<"Done reading COMSOL"<<"\n";

    // Number of events
    int numEvents = 1;

    // Begin to implement weighting fields
    // Create vector with labels of each weighting fields 
    std::vector<std::string> w_labels;
    // Define x weighting fields
    
    for (int i = 0; i < numReadoutstrips; ++i) {
        w_labels.push_back("x" + std::to_string(i+1));
        continue;
    }
    // Define y weighting fields
   
    for (int i = 0; i < numReadoutstrips; ++i) {
        w_labels.push_back("y" + std::to_string(i+1));
        continue;

    }
    
    // std::cout << "Done loading weighting fields, testing weighting field vector labels \n";

    for (int i = 0; i < w_labels.size(); ++i ) {
        std::cout << "Index: " << i << " is weighting label: " << w_labels[i] << "\n"; 
    }

    // // Note for future, this can take a while, if you're not using weighting fields comment this out
    // // load weighting potentials
	
    for (int i = 0; i < w_labels.size(); ++i) {
		std::cout<<"Reading weighting field "<<w_labels[i]<<"... \n";
    //     // Indicate name of file containing weighting field
		std::string w_field = fileloc+"field-"+w_labels[i]+".txt";
		
		fm.SetWeightingField(w_field, w_labels[i]);
	}

    // Add mesh as electrode
    fm.SetWeightingField(fileloc+"field-mesh.txt", "mesh");





    // Import gas medium
    MediumMagboltz gas;
    gas.LoadGasFile(gas_file);
    const std::string path = std::getenv("GARFIELD_INSTALL");
  	gas.LoadIonMobility(path + "/share/Garfield/Data/IonMobility_Ar+_Ar.txt");
  	// gas.LoadIonMobility(path + "/share/Garfield/Data/IonMobility_CO2+_CO2.txt");
    gas.EnableDrift();

    gas.Initialise(true);

    //assign materials to their respective domains
	const unsigned int nMaterials = fm.GetNumberOfMaterials();
  	for (unsigned int i = 0; i < nMaterials; ++i) {
    	const double eps = fm.GetPermittivity(i);
    	if (eps == 1.) fm.SetMedium(i, &gas);
    }

    // Create sensors for each electrode
    Sensor sensor;
    sensor.AddComponent(&fm);
    // Add electrodes to sensor
    for (int i = 0; i < w_labels.size(); ++i) {
        sensor.AddElectrode(&fm, w_labels[i]);
    }

    sensor.AddElectrode(&fm, "mesh");

    // Setup other stuff for signals
    // Want to avalanche ions
    AvalancheMC drift;
    drift.SetSensor(&sensor);
    drift.SetDistanceSteps(4.e-4);

 

    // Calculating Avalanches
    AvalancheMicroscopic aval;
    aval.SetSensor(&sensor);
    // aval.UseWeightingPotential(true);

    // Begin plotting drift lines
    // ViewDrift driftView;
    // TCanvas* cDrift = new TCanvas("cDrift", "", 600, 600);
    // driftView.SetCanvas(cDrift);
    // aval.EnablePlotting(&driftView);


    // Begin drift line plotting


    // To speed up calculations comment out plotting the ions
    ViewDrift* driftView = nullptr;
	TCanvas* cDrift = nullptr;
	cDrift = new TCanvas("cDrift", "Drift", 600, 600);
	driftView = new ViewDrift();
	driftView->SetCanvas(cDrift);
	aval.EnablePlotting(driftView);
    drift.EnablePlotting(driftView);
    cDrift -> SetLeftMargin(0.15);
    cDrift -> SetRightMargin(0.09);
    cDrift -> SetBottomMargin(0.15);

    // Begin signal plotting
    
    ViewSignal* signalView = nullptr;
    TCanvas* cSignal = nullptr;
    cSignal = new TCanvas("cSignal", "", 600, 600);
    signalView = new ViewSignal();
    signalView -> SetCanvas(cSignal);
    signalView -> SetSensor(&sensor);
    cSignal -> SetLeftMargin(0.15);
    cSignal -> SetRightMargin(0.09);
    cSignal -> SetBottomMargin(0.15);
    // signalView -> SetLabelY("induced current x3");
    // cSignal -> SaveAs("SIGNAL.png");

    // Plot field
    ViewField fieldView;
    // Specify field
    fieldView.SetComponent(&fm);
    // view in x-z plane
    fieldView.SetPlane(0, -1, 0, 0, 0, 0);
    fieldView.SetArea(0, 0.10, 0.5, 0.7);
    TCanvas* cf = new TCanvas("cf", "", 600, 600);
    fieldView.SetCanvas(cf);
    fieldView.SetNumberOfContours(20);
    fieldView.PlotContour();
    // Save plot as png in current directory
    cf->SaveAs("Output700/PotentialProfile.png");

    // Begin Charge plotting
    TCanvas* c2 = nullptr;
    c2 = new TCanvas("Charge", "", 600, 600);

    // Plot Electron signal
    // TCanvas* c3 = nullptr;
    // c3 = new TCanvas("Electron", "", 600, 600);


    // Plot ion signal
    // TCanvas* c4 = nullptr;
    // c4 = new TCanvas("ion", "", 600, 600);

    


    // Set time binning for signal
    // Note: tmin and tmax provide the times at which you want to record the signal betwee and ntimebins tells you how many bins
    // you want during this interval e.g tmin = 0, tmax = 100, nTimeBins = 100 is the same as recording the signal between 
    // t = 0 and 100 ns with 1 bin per ns
    // For this reason you would probably want to start recording the signal when the electron passes the mesh and actually produces
    // a signal
    
    const unsigned int nTimeBins = 10000;
    const double tmin = 0;
    const double tmax = 1000;
    const double tstep = (tmax - tmin) / nTimeBins;
    sensor.SetTimeWindow(tmin, tstep, nTimeBins);

    // Begin plotting signal
   
    // ViewSignal signalView;
    // TCanvas* cSignal = new TCanvas("cSignal", "", 600, 600);
    // signalView.SetCanvas(cSignal);
    // signalView.SetSensor(&sensor);


    // Assign values types
    int ne, ni;


    // Assign range to randomly generate numbers
    double range_from = 0.15;
    double range_to = 0.25;

    // Create a random number generator for generating large sample sizes
    std::default_random_engine generator;
    std::uniform_real_distribution<double> distribution(range_from, range_to);

    // Produce csv file to write data to
    std::ofstream ava_file(out_file);
    std::ofstream ava_current(out_signal);
    std::ofstream ava_charge(out_charge);

    // Build header for csv file
    ava_file << "Testing csv file \n";
    ava_current << "Testing recording signal current data for single event \n";
    ava_charge << "Testing recording signal charge data for single event \n";


    // Begin processing events

    for (int eventInter = 0; eventInter < numEvents; ++eventInter) {
        // sensor.ClearSignal(); // reset signal to 0 after each event


        std::cout << "------------ Processing Event: " << eventInter << "/" << numEvents - 1 << " ----------------\n";

        ava_file << eventInter << ",";
        ava_charge << "Electrode, Current, Time \n";


        // Assign initial values

        // If random values


        // If specific point
        // double xi = distribution(generator);
        double xi = 0.25;
        // double yi = distribution(generator);
        double yi = 0.25;
        double zi = 0.25;
        double ti = 0;
        double ei = 0;

        ava_file << "xi: " << xi << "," << "yi: " << yi << "," << "zi: " << zi << ",";

        // Send the electron on an avalanche (goodbye electron)

        // To check its actually doing something for higher field strengths
        double xa = 0;
        double ya = 0;
        double za = 0;
        double ta = 0;
        const int ip = 0;
        const unsigned int ie = 0;


        // std::cout << "xi = " << xi << " yi = " << yi << "\n";
        aval.AvalancheElectron(xi, yi, zi, ti, ei, 0, 0, -1);
        // The above is with the initial electron travelling down (direction given by last 3 coordinates)
        aval.GetElectronDriftLinePoint(xa, ya, za, ta, ip, ie);

        std::cout << xa << "," << ya << "," << za << "," << ta << "\n";

        aval.GetAvalancheSize(ne, ni);
        std::cout << "ne = " << ne << " ni = " << ni << "\n";

        ava_current << "TOTAL GAIN = " << "," << ne << "\n";
        ava_charge << "TOTAL GAIN = " << "," << ne << "\n";



        // Loop over all electrons in avalanche

        for (const auto& electron : aval.GetElectrons()) {
            std::cout << "---------------Drifting Ion ---------------- \n"; 

            // Simulate an ion drift line from the starting point of the electron

            const auto& p0 = electron.path[0];

        //     // Print location of ion for debugging purposes
            std::cout << "x = " << p0.x << " y = " << p0.y << " z = " << p0.z << "\n";

            // Drift Ions
            drift.DriftIon(p0.x, p0.y, p0.z, p0.t);
            std::cout << "Done \n";

        }
        std::cout << "==================== Done drifting ions ========================= \n";

        std::cout << "ne = " << ne << " ni = " << ni << "\n";

        ava_file << "ne: " << ne << "," << "ni: " << ni << "\n";


        // Plot
    //     driftView.SetPlaneXZ();
    //     driftView.SetArea(0.15, 0.14, 0.3, 0.3);
    //     driftView.Plot();
        driftView -> SetPlaneXZ();
	    driftView -> SetArea(0.15, 0, 0.14, 0.29, 0.5, 0.26);
		driftView -> Plot(true);
        std::cout << "Done Plotting \n";
        cDrift -> SaveAs("Output700/SignalDrift.png");



        // Plotting signal related things
        // signalView -> PlotSignal("x3");
        // cSignal -> SaveAs("SIGNAL.png");
        // signalView2 -> PlotSignal("x2");
        // cSignal2 -> SaveAs("SIGNAL2.png");

        // Create signal plots for all weighting strips

        for (int i = 0; i < w_labels.size(); ++i){
            // signalView -> SetLabelY("induced current "+w_labels[i]+"  [fC/ns]");
            signalView->PlotSignal(w_labels[i]);
            std::string tmp = "Output700/"+w_labels[i]+"_Current.png";
            cSignal->SaveAs(tmp.c_str());
        }

        // Plot mesh current

        signalView->PlotSignal("mesh");
        cSignal->SaveAs("Output700/meshsignal_current.png");

        

        
        // Loop to write current data
        for (int i = 0; i < nTimeBins; ++i) {
            const double t = tmin + i*tstep;
            ava_current << "x1 = " << "," << sensor.GetSignal("x1", i) << ",";
            ava_current << "x2 = " << "," << sensor.GetSignal("x2", i) << ",";
            ava_current << "x3 = " << "," << sensor.GetSignal("x3", i) << ",";
            ava_current << "x4 = " << "," << sensor.GetSignal("x4", i) << ",";
            ava_current << "x5 = " << "," << sensor.GetSignal("x5", i) << ",";
            ava_current << "y1 = " << "," << sensor.GetSignal("y1", i) << ",";
            ava_current << "y2 = " << "," << sensor.GetSignal("y2", i) << ",";
            ava_current << "y3 = " << "," << sensor.GetSignal("y3", i) << ",";
            ava_current << "y4 = " << "," << sensor.GetSignal("y4", i) << ",";
            ava_current << "y5 = " << "," << sensor.GetSignal("y5", i) << ",";
            ava_current << "Time = " << "," << t << "\n";
        }

        // Plot and write data about total charge
        sensor.IntegrateSignals();
        signalView -> SetCanvas(c2);
        for (int i = 0; i < w_labels.size(); ++i){
            // signalView -> SetLabelY("induced current "+w_labels[i]+"  [fC/ns]");
            signalView -> PlotSignal(w_labels[i]);
            std::string tmp = "Output/"+w_labels[i]+"_Charge.png";
            c2 -> SaveAs(tmp.c_str());
        }

        signalView->PlotSignal("mesh");
        c2->SaveAs("Output700/meshsignal_charge.png");

        // Write to charge file
        ava_charge << "TOTAL CHARGE = " << "," << sensor.GetSignal("x3", nTimeBins-1) << "\n";
        for (int i = 0; i < nTimeBins; ++i) {
            const double t = tmin + i*tstep;
            ava_charge << "x1 = " << "," << sensor.GetSignal("x1", i) << ",";
            ava_charge << "x2 = " << "," << sensor.GetSignal("x2", i) << ",";
            ava_charge << "x3 = " << "," << sensor.GetSignal("x3", i) << ",";
            ava_charge << "x4 = " << "," << sensor.GetSignal("x4", i) << ",";
            ava_charge << "x5 = " << "," << sensor.GetSignal("x5", i) << ",";
            ava_charge << "y1 = " << "," << sensor.GetSignal("y1", i) << ",";
            ava_charge << "y2 = " << "," << sensor.GetSignal("y2", i) << ",";
            ava_charge << "y3 = " << "," << sensor.GetSignal("y3", i) << ",";
            ava_charge << "y4 = " << "," << sensor.GetSignal("y4", i) << ",";
            ava_charge << "y5 = " << "," << sensor.GetSignal("y5", i) << ",";
            ava_charge << "Time = " << "," << t << "\n";
        }
    std::cout << "-------------------DONE-----------------------\n";

        





    }
    ava_file.close();
    ava_charge.close();
    ava_current.close();

    app.Run(true);
    return 0;

}