#include <cmath>
#include <iostream>

#include "TApplication.h"

#include "Garfield/FundamentalConstants.hh"
#include "Garfield/MediumMagboltz.hh"
#include "Garfield/ViewMedium.hh"

using namespace Garfield;

int main(int argc, char* argv[]) {
  TApplication app("app", &argc, argv);
  
  // - - - - - - - - -
  // Original Gas File
  // - - - - - - - - -
  MediumMagboltz gas01;
  gas01.LoadGasFile("ar_80_co2_20_2T.gas");
  
  std::vector<double> efields;
  std::vector<double> bfields;
  std::vector<double> angles;
  gas01.GetFieldGrid(efields, bfields, angles);
  const auto nE = efields.size();
  const auto nB = bfields.size();
  const auto nA = angles.size();

  // std::cout<<"Before reading :: Efield size = "<<efields.size()<<" Bfield size = "<<bfields.size()<<" Angles size = "<<angles.size()<<std::endl;
  
  std::vector<double> townsend;
  std::vector<double> attachment;
  std::vector<double> driftvelocityE;
  std::vector<double> driftvelocityB;
  std::vector<double> driftvelocityExB;
  std::vector<double> longdiffusion;
  std::vector<double> transdiffusion;
  std::vector<double> lorentzangle;
  
  for (size_t j = 0; j < nB; ++j) {
    for (size_t k = 0; k < nA; ++k) {
      for (size_t i = 0; i < nE; ++i) {
	// std::cout<<" READING :: i,j,k = "<<i<<","<<j<<","<<k<<" - max ["<<nE<<","<<nB<<","<<nA<<"] - Efield = "<<efields[i]<<" V/cm"<<std::endl;
	// drift velocity
        double ve = 0.;
        gas01.GetElectronVelocityE(i, j, k, ve); // cm/ns
        driftvelocityE.push_back(ve);
	gas01.GetElectronVelocityB(i, j, k, ve); // cm/ns
	driftvelocityB.push_back(ve);
	gas01.GetElectronVelocityExB(i, j, k, ve); // cm/ns
	driftvelocityExB.push_back(ve);
	// townsend coefficient
        double alpha = 0.;
        gas01.GetElectronTownsend(i, j, k, alpha);
        alpha = exp(alpha);
	townsend.push_back(alpha);
	// attachment coefficient
	double eta = 0.;
	gas01.GetElectronAttachment(i, j, k, eta);
	eta = exp(eta);
	attachment.push_back(eta);
	// diffusion
	double df = 0.;
	gas01.GetElectronLongitudinalDiffusion(i, j, k, df);
	longdiffusion.push_back(df);
	gas01.GetElectronTransverseDiffusion(i, j, k, df);
	transdiffusion.push_back(df);
	// lorentzangle
	double la = 0.;
	gas01.GetElectronLorentzAngle(i, j, k, la);
	lorentzangle.push_back(la);
      }
    }
  }

  // - - - - - - - - -
  // New Gas File
  // - - - - - - - - -
  MediumMagboltz gas02;
  gas02.SetComposition("ar", 80., "co2", 20.);
  gas02.SetTemperature(293.15); // Kelvin
  gas02.SetPressure(760.); // Torr
  // set same field grid
  gas02.SetFieldGrid(efields, bfields, angles);

  // std::cout<<"Before writing :: Efield size = "<<efields.size()<<" Bfield size = "<<bfields.size()<<" Angles size = "<<angles.size()<<std::endl;

  int counter = 0;
  for (size_t j = 0; j < nB; ++j) {
    for (size_t k = 0; k < nA; ++k) {
      for (size_t i = 0; i < nE; ++i) {
	// std::cout<<" WRITING :: i,j,k = "<<i<<","<<j<<","<<k<<" - max ["<<nE<<","<<nB<<","<<nA<<"] - Efield = "<<efields[i]<<" V/cm - counter = "<<counter<<std::endl;
        gas02.SetElectronVelocityE(i, j, k, driftvelocityE[counter]);
	gas02.SetElectronVelocityB(i, j, k, driftvelocityB[counter]);
	gas02.SetElectronVelocityExB(i, j, k, driftvelocityExB[counter]);
	gas02.SetElectronTownsend(i,j,k,log(townsend[counter]));
	gas02.SetElectronAttachment(i,j,k,log(attachment[counter]));
	gas02.SetElectronLongitudinalDiffusion(i,j,k,longdiffusion[counter]);
	gas02.SetElectronTransverseDiffusion(i,j,k,transdiffusion[counter]);
	gas02.SetElectronLorentzAngle(i,j,k,lorentzangle[counter]);
	++counter;
      }
    }
  }
  gas02.WriteGasFile("ar_80_co2_20_2T_copy.gas");


  // - - - - - - - - - - -
  // Compare two gas files
  // - - - - - - - - - - -

  ViewMedium view;

  TCanvas c1;
  view.SetCanvas(&c1);
  view.SetMedium(&gas01);
  view.PlotElectronVelocity('e');
  std::vector<std::string> labels = {"Ar/CO_{2} 80/20 2T Original", "Ar/CO_{2} 80/20 2T Copy"};
  view.SetLabels(labels);
  view.SetMedium(&gas02);
  view.PlotElectronVelocity('e', true);

  TCanvas c2;
  view.SetCanvas(&c2);
  view.SetMedium(&gas01);
  view.PlotElectronTownsend('e');
  view.SetLabels(labels);
  view.SetMedium(&gas02);
  view.PlotElectronTownsend('e', true);


  
  app.Run(true);
}