#include "MCascade.h"

#include <math.h>     // fabs, for alpha
#include <iostream.h>

#include <TF1.h>
#include <TH2.h>
#include <TList.h>
#include <TFile.h>
#include <TTree.h>
#include <TTimer.h>
#include <TStyle.h>
#include <TBranch.h>
#include <TCanvas.h>
#include <TRandom3.h>
#include <TStopwatch.h>

#include "MPhoton.h"
#include "MElectron.h"

#include "MH.h"
#include "MBinning.h"

ClassImp(MCascade);

Double_t PrimSpect(Double_t *x, Double_t *k)
{
    return pow(pow(10, x[0]), k[0]);
}

Double_t PhotonSpect(Double_t *x, Double_t *k=NULL)
{
    Double_t Ep = pow(10, x[0]);

    Double_t res = MPhoton::Int2(&Ep, k);
    return res*1e55; //65/k[0];

    //return MPhoton::Planck(&Ep, &k[1]);
}

Double_t Sbar_sigmas(Double_t *x, Double_t *k)
{
    Double_t sbar = pow(10, x[0]);

    Double_t s = 1./(sbar*4);

    Double_t sigma = MPhoton::Sigma_gg(&s);

    return sigma*sbar*1e28;
}

Double_t RandomThetaG(Double_t Eg, Double_t Ep)
{
    Double_t E0 = 511e-6; // [GeV]

    Double_t f = Eg/E0*Ep/E0;

    if (f<1)
        return 0;

    static TF1 func("RndThetaG", Sbar_sigmas, 0, 0, 0);

    func.SetRange(0, log10(f));
    func.SetNpx(50);

    Double_t sbar  = pow(10, func.GetRandom());
    Double_t theta = acos(1.-sbar*2/f);

    return theta;
}

Double_t MCascade::GetEnergy()
{
    static int bin=0;
    Double_t w = log10(fEHi/fELo)/fNumBins;

    Double_t E = fELo*pow(10, gRandom->Uniform(w) + w*(fNumBins-bin-1));

    if (++bin==fNumBins)
        bin=0;

    return E;
}

TFile *MCascade::OpenFile(TString fFilename)
{
    TFile *fFile = new TFile(fFilename, "CREATE", "Intergalactic cascade", 9);

    if (fFile->IsZombie())
    {
        delete fFile;
        return NULL;
    }

    Write("Setup");

    cout << "Trees... " << flush;

    TTree *T1 = new TTree ("Photons",    "Photons from Cascade");
    TTree *T2 = new TTree ("Electrons",  "Electrons in the Cascade");

    cout << "Branches... " << flush;

    MPhoton dummyp;
    void *ptr = &dummyp;
    fBranchGammas = T1->Branch("MPhoton.",   "MPhoton",   &ptr);

    MElectron dummye;
    ptr = &dummye;
    fBranchElectrons = T2->Branch("MElectron.", "MElectron", &ptr);

    return fFile;
}

void MCascade::CloseFile(TFile *fFile)
{
    fFile->Write();
    cout << "Wrote: " << fFile->GetName() << endl;
    delete fFile;
}

void MCascade::ProcessElectron(MElectron &e, TList &fListGammas)
{
    Double_t Ee = e.GetEnergy();

    cout << ":" << flush;

    int test = fNumMaxInvCompton<0 ? -1 : 0;
    int n=0;
    while ((test<0 ? true : (test++<fNumMaxInvCompton)) &&
           (e.GetEnergy() > fRatioInvCompton*Ee))
    {
        n++;

        if (!e.SetNewPositionB(e.GetZ()>fBubbleZ ? fB : 0))
        {
            cout << "!" << flush;
            return;
        }

        MPhoton *p = e.DoInvCompton();

        fBranchElectrons->GetTree()->Fill();

        //cout << "." << flush;
        fListGammas.Add(p);

        if (fabs(e.GetTheta()*3437)>60)  // < 60min
        {
            cout << "T" << flush;
            return;
        }

        if (e.GetEnergy()<Ee*1e-3) // <2e3
        {
            cout << "E" << flush;
            return;
        }

        if (e.GetEnergy()<1e2)
        {
            cout << "x" << flush;
            return;
        }
    }
    cout << n << flush;
}

Bool_t MCascade::ProcessGamma(MPhoton &p, Double_t weight, TList &fListElectrons)
{
    Double_t Eg = p.GetEnergy();

    Double_t E0    = 511e-6;
    Double_t z     = p.GetZ()+1;
    Double_t lolim = E0*E0/Eg;
    Double_t inf   = (Eg<1e6 ? 3e-11*z : 3e-12*z);
    if (Eg<5e4)
        inf = 3e-11*z*pow(10, 9.4-log10(Eg)*2);

    TF1 phot("PhotonSpectrum", PhotonSpect, 0, 0, 2);
    phot.SetRange(log10(lolim), log10(inf));
    phot.SetNpx(50);
    phot.SetParameter(0, Eg);
    while (1)
    {
        if (!p.SetNewPosition())
            return kTRUE;

        //
        // Sample phtoton from background and interaction angle
        //
        phot.SetParameter(1, p.GetZ());
        Double_t pe = phot.GetRandom();
        if (pe==0)
        {
            cout << "z" << flush;
            continue;
        }

        Double_t Ep = pow(10, pe);
        Double_t theta = RandomThetaG(Eg, Ep);
        if (theta==0)
        {
            cout << "t" << flush;
            continue;
        }

        if (!fPair.Process(&p, Ep, theta, &fListElectrons))
        {
            // should never happen
            cout << "0" << flush;
            continue;
        }

        return kFALSE;
    }
}

Double_t MCascade::ProcessGammas(TList &fListGammas, TList &fListElectrons, Double_t weight)
{
    Double_t Esum = 0;

    MPhoton *p = NULL;
    fBranchGammas->SetAddress(&p);

    cout << ":" << fListGammas.GetSize() << ":" << flush;

    TIter NextP(&fListGammas);
    while ((p=(MPhoton*)NextP()))
    {
        if (!ProcessGamma(*p, weight, fListElectrons))
        {
            delete fListGammas.Remove(p);
            cout << "." << flush;
            continue;
        }

        fBranchGammas->GetTree()->Fill();
        Esum += p->GetEnergy();
        //cout << "!" << flush;
    }

    return Esum;
}

Double_t MCascade::ProcessElectrons(TList &fListElectrons, TList &fListGammas)
{
    Double_t E = 0;

    cout << ":" << fListElectrons.GetSize() << flush;

    TIter Next(&fListElectrons);
    MElectron *e = NULL;
    fBranchElectrons->SetAddress(&e);
    while ((e=(MElectron*)Next()))
    {
        e->SetIsPrimary(kTRUE);
        fBranchElectrons->GetTree()->Fill();
        e->SetIsPrimary(kFALSE);

        ProcessElectron(*e, fListGammas);

        E += e->GetEnergy();
    }
    fListElectrons.Delete();

    return E;
}

void MCascade::ProcessPrimaryGamma(Double_t E, Double_t weight)
{
    TList fListGammas;
    TList fListElectrons;

    fListGammas.SetOwner();
    fListElectrons.SetOwner();

    MPhoton *gamma=new MPhoton(E, fSrcZ);
    gamma->SetSrcR(fSrcR);
    gamma->InitRandom();
    fListGammas.Add(gamma);

    gamma->SetIsPrimary(kTRUE);
    fBranchGammas->SetAddress(&gamma);
    fBranchGammas->GetTree()->Fill();
    gamma->SetIsPrimary(kFALSE);

    Double_t Esum=0; // sum of all energies
    Double_t Emis=0; // sum of the energies thrown away
    while (1)
    {
        if (fListGammas.GetSize()==0)
            break;

        cout << " |P" << flush;

        Esum += ProcessGammas(fListGammas, fListElectrons, weight);

        if (!fIsBatch)
            fListGammas.ForEach(MPhoton, Fill)(fHist, fDisplayIndex, weight);
        fListGammas.Delete();

        if (fListElectrons.GetSize()==0)
            break;

        cout << " |E" << flush;

        Emis += ProcessElectrons(fListElectrons, fListGammas);
    }
    Esum += Emis;

    cout << " ----> " << Form("%3.1f %3.1e / %3.1f %3.1e", Emis/E, Emis, Esum/E, Esum) << endl;
}

MCascade::MCascade()
{
    if (gRandom)
        delete gRandom;

    TRandom r(0);
    gRandom = new TRandom3(r.GetSeed());

    fHist.SetName("Spectrum");
    fHist.SetXTitle("E [GeV]");
    fHist.SetYTitle(Form("E^{%.1f} Counts", fDisplayIndex));
    fHist.GetXaxis()->SetLabelOffset(-0.015);
    fHist.GetXaxis()->SetTitleOffset(1.1);
    fHist.SetFillStyle(0);
    fHist.SetMarkerStyle(kPlus);
}

MCascade::~MCascade()
{
    delete gRandom;
    gRandom = 0;
}

void MCascade::SetSourceZ(Double_t z)
{
    fSrcZ = z;
    fSrcR = MParticle::RofZ(&z);
}

void MCascade::SetSourceRZ(Double_t r) // [kpc]
{
    fSrcZ = MParticle::ZofR(&r);
    fSrcR = r;
}

void MCascade::SetEnergyBins(Int_t n, Double_t lo, Double_t hi)
{
    fNumBins = n;  // number of bins produced in energy spectrum

    fELo = lo;     // lower limit of displayed spectrum
    fEHi = hi;     // upper limit of spectrum (cutoff)
}

void MCascade::SetBradius(Double_t r) // [Mpc]
{
    Double_t bubbler = fSrcR-1e3*r;
    fBubbleZ = MParticle::ZofR(&bubbler);
}

void MCascade::Run(TString filename, Bool_t draw)
{
    fIsBatch = gROOT->IsBatch() ? kFALSE : draw;

    // ------------------------------

    cout << "Output File '" << filename << "'... " << flush;

    TFile *file=OpenFile(filename);
    if (!file)
        return;

    // ------------------------------

    cout << endl;

    cout << "R = " << fSrcR << "kpc" << endl;
    cout << "Z = " << fSrcZ << endl;

    cout << "Setting up: Histograms... " << flush;

    fHist.Reset();
    TH1D histsrc;

    MBinning bins;
    bins.SetEdgesLog(fNumBins, fELo, fEHi);

    MH::SetBinning(&fHist,   &bins);
    MH::SetBinning(&histsrc, &bins);

    TCanvas *c=NULL;

    if (!fIsBatch)
    {
        fHist.SetMinimum(pow(fELo, fSpectralIndex+fDisplayIndex)/100);
        histsrc.SetMinimum(pow(fELo, fSpectralIndex+fDisplayIndex)/100);

        gStyle->SetOptStat(10);

        //
        // Don't change the order!!!
        //
        histsrc.SetFillStyle(0);
        histsrc.SetMarkerStyle(kMultiply);
        histsrc.SetMarkerColor(kRed);
        histsrc.SetLineColor(kRed);

        c=MH::MakeDefCanvas("Cascade", "Cascade");

        c->SetGrid();
        c->SetLogx();
        c->SetLogy();

        fHist.Draw("P");
        histsrc.Draw("Psame");
        histsrc.Draw("same");
        fHist.Draw("same");
    }

    // ------------------------------

    cout << "Timers... " << flush;

    TTimer timer("gSystem->ProcessEvents();", 333, kFALSE);
    if (!fIsBatch)
        timer.TurnOn();

    TStopwatch clock;
    clock.Start();

    cout << "Done. " << endl;

    Int_t n=0;
    Double_t starttime = TStopwatch::GetRealTime(); // s
    while (TStopwatch::GetRealTime()<starttime+fRuntime)
        for (int i=0; i<fNumBins; i++)
        {
            n++;

            Double_t E = GetEnergy();
            Double_t weight = pow(E, fSpectralIndex);

            if (!fIsBatch)
                histsrc.Fill(E, pow(E, fDisplayIndex) * weight);

            cout << "--> " << n << ". " << Form("%d: %3.1e", (int)(starttime+fRuntime-TStopwatch::GetRealTime()), E) << flush;

            ProcessPrimaryGamma(E, weight);

            if (fIsBatch)
                continue;

            fHist.SetTitle(Form("E^{%.1f}  z=%f  T=%d'%d\"  N=%d", fSpectralIndex, fSrcZ,
                                (int)fRuntime/60, (int)fRuntime%60, n));

            timer.Stop();
            c->Update();
            timer.Start(250);
        }

    cout << endl;

    clock.Stop();
    clock.Print();

    timer.Stop();

    CloseFile(file);

    cout << "Created " << n << " gammas (" << n/fNumBins << "rows) from source with E^" << fSpectralIndex << endl;
    cout << "Processing time: " << Form("%.1f", (TStopwatch::GetRealTime()-starttime)/n) << " sec/gamma  (";
    cout << Form("%.1f", (TStopwatch::GetRealTime()-starttime)/n*fNumBins/60) << " min/row)" << endl;

    // ------------------------------

    if (!fIsBatch)
    {
        c->Clear();

        fHist.SetTitle(Form("E^{%.1f}  z=%f  T=%d'%d\"  N=%d", fSpectralIndex, fSrcZ, (int)fRuntime/60, (int)fRuntime%60, n));

        TH1 &h1 = *fHist.DrawCopy("P");
        TH1 &h2 = *histsrc.DrawCopy("Psame");
        h2.Draw("Csame");
        h1.Draw("Csame");
    }
}