#include <iostream>
#include <chrono>

#include <TROOT.h>
#include <TFile.h>
#include <TCanvas.h>
#include <TPad.h>
#include <TF1.h>
#include <TH2D.h>
#include <TRandom2.h>

#include "genpxr.h"

TRandom2 *randoms;

genPXR *pxr;

void genpxr() {
   pxr = new genPXR(11.);
   //pxr->setCrystal(45, 19.2 * M_PI/180, 0.03);
   pxr->setCrystal(78.3608294720761,0.125687512569129,0.389203402912244);
   double theta(0.568084461299755), phi(2.24579748833124);
   pxr->setFilterLimits(10, 100);
   double ans1 = pxr->dNdOmegadz(theta * 180/M_PI, phi * 180/M_PI, 1);
   std::cout << "dNdOmegadz(theta, phi, verbosity=1)=" << ans1 << std::endl;
   double ans2 = pxr->dNdOmegadz2(theta * 180/M_PI, phi * 180/M_PI, 1);
   std::cout << "dNdOmegadz2(theta, phi, verbosity=1)=" << ans2 << std::endl;
   int hkl[3] = {-1,-1,1};
   std::vector<double> khat = {sin(theta) * cos(phi),
                               sin(theta) * sin(phi),
                               cos(theta)};
   double ome_keV = pxr->omega_keV(hkl, khat);
   double eps_h = pxr->Esuscept(hkl, ome_keV);
   double kmag = ome_keV * pxr->refIndex(ome_keV);
   std::vector<double> h_hkl(pxr->crystalWavenumber(hkl));
   std::vector<double> k_hkl = {kmag * khat[0] + h_hkl[0] * pxr->hbarc_keV_m,
                                kmag * khat[1] + h_hkl[1] * pxr->hbarc_keV_m,
                                kmag * khat[2] + h_hkl[2] * pxr->hbarc_keV_m};
   double kperp2 = pow(k_hkl[0], 2) + pow(k_hkl[1], 2);
   std::cout << "radiation direction theta,phi=" << theta << "," << phi << std::endl;
   std::cout << "omega_keV((-1,-1,1), khat)=" << std::setprecision(15) << ome_keV << "keV" << std::endl;
   std::cout << "Esuscept((-1,-1,1),omega_keV)=" << std::setprecision(15) << eps_h << std::endl;
   std::cout << "k_hkl_perp2=" << kperp2 << std::endl;

   randoms = new TRandom2(0);
}

double FFatomic(double *vars, double  *pars) {
   return pxr->FFatomic(vars[0]);
}

double FFatomic_Lobato(double *vars, double  *pars) {
   return pxr->FFatomic_Lobato(vars[0]);
}

void plot_FFatomic() {
   TF1 *simpleFF = new TF1("simpleFF", FFatomic, 0, 100);
   TF1 *LobatoFF = new TF1("LobatoFF", FFatomic_Lobato, 0, 100);
   simpleFF->SetLineColor(1);
   LobatoFF->SetLineColor(2);
   simpleFF->SetTitle("atomic electron cloud charged form factor for carbon");
   simpleFF->GetXaxis()->SetTitle("q (keV/c)");
   simpleFF->GetYaxis()->SetTitle("|FF|^{2}");
   simpleFF->Draw();
   LobatoFF->Draw("same");
   TCanvas *c1 = (TCanvas*)gROOT->FindObject("c1");
   c1->Update();
}

TH2D *plot_intensity(int ntrials, int nupdate=10000) {
   TCanvas *c1 = new TCanvas("c1", "c1", 700, 600);
   TH2D *hintens = new TH2D("hintens", "PXR intensity map", 1000, -90, 90, 
                                                            1000, -90, 90);
   hintens->Draw("colz");
   c1->SetLeftMargin(0.15);
   c1->SetRightMargin(0.15);
   hintens->SetStats(0);
   hintens->SetContour(100);
   hintens->GetXaxis()->SetTitle("#theta_{x} / degrees");
   hintens->GetYaxis()->SetTitle("#theta_{y} / degrees");
   double cost1 = 0.;
   double cost2 = 1.;
   double phi1 = 0;
   double phi2 = 2*M_PI;
   double jacobian = (cost2 - cost1) * (phi2 - phi1) / ntrials;
   auto begin = std::chrono::high_resolution_clock::now();
   for (int itry=1; itry <= ntrials; ++itry) {
      double theta = acos(randoms->Uniform(cost1, cost2));
      double phi = randoms->Uniform(phi1, phi2);
      double weight = pxr->dNdOmegadz2(theta * 180/M_PI, phi * 180/M_PI);
      hintens->Fill(theta * 180/M_PI * cos(phi), theta * 180/M_PI * sin(phi),
                    weight * jacobian);
      if (itry % nupdate == 0) {
         auto now = std::chrono::high_resolution_clock::now();
         auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(now - begin);
         std::cout << itry << ": " << hintens->Integral() * ntrials/itry
                   << ", elapsed=" << elapsed.count() * 1e-9 << "cpu-seconds"
                   << std::endl;
         char title[80];
         sprintf(title, "PXR intensity map, %d trials", itry);
         hintens->SetTitle(title);
         hintens->Draw("colz");
         if (c1 == 0)
            c1 = (TCanvas*)gROOT->FindObject("c1");
         c1->Update();
      }
   }
   return hintens;
}

void try_rotations(int ntrials) {
   int hkl[3]= {-1,-1,1};
   double best_kperp2 = 1e9;
   for (int itry=1; itry <= ntrials; ++itry) {
      double crystal_phideg = randoms->Uniform(90);
      double crystal_alphax = randoms->Uniform(-0.5, 0.5);
      double crystal_alphay = randoms->Uniform(-0.5, 0.5);
      pxr->setCrystal(crystal_phideg, crystal_alphax, crystal_alphay, 0);
      //pxr->setCrystal(14.3411829625256,0.482482467312366,-0.418739928863943);
      double ray_phi = randoms->Uniform(2*M_PI);
      double ray_theta = acos(randoms->Uniform(0,1));
      //ray_theta=0.0133799436854893; ray_phi=1.19142649151279;
      std::vector<double> khat = {sin(ray_theta) * cos(ray_phi),
                                  sin(ray_theta) * sin(ray_phi),
                                  cos(ray_theta)};
      double ome_keV = pxr->omega_keV(hkl, khat);
      double kmag = ome_keV * pxr->refIndex(ome_keV);
      std::vector<double> h_hkl(pxr->crystalWavenumber(hkl));
      std::vector<double> k_hkl = {kmag * khat[0] + h_hkl[0] * pxr->hbarc_keV_m,
                                   kmag * khat[1] + h_hkl[1] * pxr->hbarc_keV_m,
                                   kmag * khat[2] + h_hkl[2] * pxr->hbarc_keV_m};
      double kperp2 = pow(k_hkl[0], 2) + pow(k_hkl[1], 2);
      if (ome_keV > 10 && kperp2 < best_kperp2) {
         best_kperp2 = kperp2;
         std::cout << "new best_kperp2 found, " << best_kperp2
                   << " with omega=" << ome_keV << "keV"
                   << " at crystal phi,alphax,alphay=" << crystal_phideg
                   << "," << crystal_alphax << "," << crystal_alphay
                   << " direction theta,phi=" << ray_theta << "," << ray_phi
                   << std::endl;
      }
   }
}