/*  Description: This macro calculate MAP for radius correction

    Author : Simeon Lebedev
    E-mail : S.Lebedev@gsi.de
*/

void run_radius_correction ()
{
    TStopwatch timer;
    timer.Start();
    // ----  Load libraries   -------------------------------------------------
    gROOT->LoadMacro("$VMCWORKDIR/gconfig/basiclibs.C");
    basiclibs();
    gSystem->Load("libGeoBase");
    gSystem->Load("libParBase");
    gSystem->Load("libBase");
    gSystem->Load("libCbmBase");
    gSystem->Load("libField");
    gSystem->Load("libGen");
    gSystem->Load("libPassive");
    gSystem->Load("libMvd");
    gSystem->Load("libSts");
    gSystem->Load("libRich");
    gSystem->Load("libTrd");
    gSystem->Load("libTof");
    gSystem->Load("libEcal");
    gSystem->Load("libGlobal");
    gSystem->Load("libKF");
    gSystem->Load("libL1");
    // ------------------------------------------------------------------------
    ///A axis
    TH2D* fh_AxisAXY = new TH2D("fh_AxisAXY","Axis A distribution (x,y)",40,-200,200,50,-250,250);
    TH2D* fh_AxisAXYW = new TH2D("fh_AxisAXYW","Axis A distribution(x,y)",40,-200,200,50,-250,250);

        
    ///B axis
    TH2D* fh_AxisBXY = new TH2D("fh_AxisBXY","Axis B distribution (x,y)",40,-200,200,50,-250,250);
    TH2D* fh_AxisBXYW = new TH2D("fh_AxisBXYW","Axis B distribution(x,y)",40,-200,200,50,-250,250);
   
    ///Set Mean value of A and B axeses
    Double_t fMeanAaxis = 6.198;
    Double_t fMeanBaxis = 5.63;    

    char fileMC[200], fileRec[200];
    /*sprintf(fileMC,"/d/cbm02/slebedev/rich/ellipse_qa/auau.25gev.centr.0000.mc.root");
    cout<<fileMC<<endl;
    TFile *f1 = new TFile(fileMC,"R");
    TTree* t1 = (TTree*)f1->Get("cbmsim");
    TFolder *fd1 = f1->Get("cbmroot");
    TClonesArray* fListStack = (TClonesArray*) fd1->FindObjectAny("MCTrack");
    t1->SetBranchAddress("MCTrack",&fListStack);*/
    sprintf(fileRec, "/d/cbm02/slebedev/rich/ellipse_qa/el.0002.reco.root");
    //sprintf(fileRec,"/d/cbm02/slebedev/rich/ellipse_qa/auau.25gev.centr.0000.reco.root");
    TFile *f = new TFile(fileRec,"R");
    TTree* t = f->Get("cbmsim");
    TFolder *fd = f->Get("cbmout");
    TClonesArray *fRichRings = (TClonesArray*) fd->FindObjectAny("RichRing");
    t->SetBranchAddress(fRichRings->GetName(),&fRichRings);

    Int_t nEvents=t->GetEntries();
    cout<<" nEvents ="<<nEvents<<endl;
    for(Int_t ievent=0;ievent<nEvents;ievent++ ) {
        cout<<"ievent = "<<ievent;
        CbmRichRing *ring=NULL;
        
        //t1->GetEntry(ievent);
        t->GetEntry(ievent);

        Int_t nofRings = fRichRings->GetEntries();
        cout<<"  nofRings = "<<nofRings<<endl;

        for(Int_t iRing=0; iRing < nofRings; iRing++){

            ring = (CbmRichRing*)fRichRings->At(iRing);
            if (!ring) continue;
            
            ///check for primary electrons
            if (ring->GetRecFlag() != 3) continue;
            
            Double_t radius = ring->GetRadius();
            Double_t axisA = ring->GetAaxis();
            Double_t axisB = ring->GetBaxis(); 
            Double_t centerX = ring->GetCenterX();
            Double_t centerY = ring->GetCenterY();
            
            if (axisA > 8. || axisB > 8.) continue;
            if (axisA < 4. || axisB < 4.) continue;
                        
            fh_AxisAXYW->Fill(centerX, centerY, axisA);  
            fh_AxisAXY->Fill(centerX, centerY);   
            
            fh_AxisBXYW->Fill(centerX, centerY, axisB);  
            fh_AxisBXY->Fill(centerX, centerY);                
                       
        } //iRing
    }//iEvent
    //f->Close();
    //f1->Close();
    
    
    gStyle->SetPalette(1,0);
    gStyle->SetHistLineWidth(2);

    c1 = new TCanvas("c1","c1",10,10,800,800);
    c1->Divide(1,2);
    c1->cd(1);
    fh_axisAXY = new TH2D("fh_axisAXY","A distribution (x,y);X, [cm];Y, [cm]",40,-200,200,50,-250,250);
    fh_axisAXY->Divide(fh_AxisAXYW,fh_AxisAXY);    
    fh_axisAXY->SetMinimum(5.0);
    fh_axisAXY->SetMaximum(7.0);           
    fh_axisAXY->Draw("COLZ"); 
     
    c1->cd(2);
    fh_axisBXY = new TH2D("fh_axisBXY","B distribution (x,y);X, [cm];Y, [cm]",40,-200,200,50,-250,250);
    fh_axisBXY->Divide(fh_AxisBXYW,fh_AxisBXY);    
    fh_axisBXY->SetMinimum(5.0);
    fh_axisBXY->SetMaximum(6.5);           
    fh_axisBXY->Draw("COLZ");      

    c2 = new TCanvas("c2","c2",10,10,800,800);
    c2->Divide(1,2);
    fh_mapaxisAXY = new TH2D("fh_mapaxisAXY","dA distribution (x,y);X, [cm];Y, [cm]",40,-200,200,50,-250,250);
    fh_mapaxisBXY = new TH2D("fh_mapaxisBXY","dB distribution (x,y);X, [cm];Y, [cm]",40,-200,200,50,-250,250);  
    
    for (Int_t iX = 1; iX < fh_mapaxisAXY->GetNbinsX() + 1; iX++){
        for (Int_t iY = 1; iY < fh_mapaxisAXY->GetNbinsY() + 1; iY++){
            cout << fh_axisAXY->GetBinContent(iX, iY) << " ";
            if (TMath::Abs( fh_axisAXY->GetBinContent(iX, iY) ) != 0) 
                fh_mapaxisAXY->SetBinContent(iX, iY, fMeanAaxis - fh_axisAXY->GetBinContent(iX, iY) );
                
            if (TMath::Abs( fh_axisBXY->GetBinContent(iX, iY) ) != 0) 
                fh_mapaxisBXY->SetBinContent(iX, iY, fMeanBaxis - fh_axisBXY->GetBinContent(iX, iY) );                      
        }
        cout << endl;
    }      
    fh_mapaxisAXY->SetMinimum(-0.5);
    fh_mapaxisAXY->SetMaximum(0.5);           

    fh_mapaxisBXY->SetMinimum(-0.5);
    fh_mapaxisBXY->SetMaximum(0.5);          

    c2->cd(1);    
    fh_mapaxisAXY->Draw("COLZ");
    c2->cd(2);         
    fh_mapaxisBXY->Draw("COLZ");  
    
///// Check correction procedure
    TH1D* fh_Abefore = new TH1D("fh_Abefore","A before;radius, [cm]", 90, 0., 9.);;
    TH1D* fh_Bbefore= new TH1D("fh_Bbefore","B before;radius, [cm]", 90, 0., 9.);;
    
    TH1D* fh_A = new TH1D("fh_A","A after;radius, [cm]", 90, 0., 9.);;
    TH1D* fh_B = new TH1D("fh_B","B after;radius, [cm]", 90, 0., 9.);;
    
    cout <<"Check correction procedure......" << endl;
    for(Int_t ievent=0;ievent<nEvents;ievent++ ) {
        CbmRichRing *ring=NULL;
        if (ievent % 100 == 0) cout << ievent << "   ";
        //t1->GetEntry(ievent);
        t->GetEntry(ievent);

        Int_t nofRings = fRichRings->GetEntries();

        for(Int_t iRing=0; iRing < nofRings; iRing++){

            ring = (CbmRichRing*)fRichRings->At(iRing);
            if (!ring) continue;
            
            ///check for primary electrons
            if (ring->GetRecFlag() != 3) continue;
            
            Double_t radius = ring->GetRadius();
            Double_t centerX = ring->GetCenterX();
            Double_t centerY = ring->GetCenterY();
            Double_t axisAbefore = ring->GetAaxis();
            Double_t axisBbefore = ring->GetBaxis();             
            fh_Abefore->Fill(axisAbefore);
            fh_Bbefore->Fill(axisBbefore); 
            
            Double_t axisA = ring->GetAaxis() + fh_mapaxisAXY->GetBinContent(fh_mapaxisAXY->FindBin(centerX,centerY));
            Double_t axisB = ring->GetBaxis() + fh_mapaxisBXY->GetBinContent(fh_mapaxisBXY->FindBin(centerX,centerY));               
            fh_A->Fill(axisA);
            fh_B->Fill(axisB); 
        } //iRing
    }//iEvent
    
    c3 = new TCanvas("c3","c3",10,10,800,800);
    c3->Divide(2,2);
    c3->cd(1);
    fh_Abefore->Draw();
    c3->cd(2);
    fh_Bbefore->Draw(); 
    c3->cd(3);
    fh_A->Draw();
    c3->cd(4);
    fh_B->Draw();     

/// Write correction map to the file
    TFile *file = new TFile("radius_correction_map.root", "recreate");
    fh_mapaxisAXY->Write();
    fh_mapaxisBXY->Write();
       
}