#include "TH1.h"
#include "TTree.h"
#include "TFile.h"
#include "TCanvas.h"

#include <iostream>

TH1D** plot_angle_flux(TTree *t, double eg, double d_angle, double d_energy, int orbit) {
    if (!t) std::cout << "ERROR: plot_angle_flux : t == nullptr\n";
    int ch1 = (int) ( ((eg-d_energy)*1000+40.424)/10.536 );
    int ch2 = (int) ( ((eg+d_energy)*1000+40.424)/10.536 );

    int bin1 = (int) (30/d_angle);
    int bin2 = 2*bin1;
    int bin3 = 3*bin1;

    TH1D **h = new TH1D*[6];
	h[0] = new TH1D("h0", "0-60",   bin2, 0.0,  60.0);
	h[1] = new TH1D("h1", "60-90",  bin1, 60.0, 90.0);
    h[2] = new TH1D("h2", "90-180", bin3, 90.0, 180.0);
    h[3] = new TH1D("h3", "60-0",   bin2, 0.0,  60.0);
    h[4] = new TH1D("h4", "90-60",  bin1, 60.0, 90.0);
    h[5] = new TH1D("h5", "180-90", bin3, 90.0, 180.0);

	int **ct = new int*[6]; // total bin filled counts for normalizing (10s/count)
    for(int k=0; k<6; k++) {
    	ct[k] = new int[bin3+2];
    	for(int m=0; m<bin3+2; m++) {
    		ct[k][m] = 0;
    	}
    }

	Double_t s_inc;
    Double_t se[976];
	Int_t s_qs;
    Int_t s_orbit;
	t->SetBranchAddress("s_inc", &s_inc);
	t->SetBranchAddress("se", se);
	t->SetBranchAddress("s_qs", &s_qs);
	t->SetBranchAddress("s_orbit", &s_orbit);

	Long64_t n = t->GetEntries();
	double sum; // sum counts of each line between channel ch1 and ch2 for filling
	double jd;
	int index;
	for(Long64_t i=0; i<n-1; i++){
		t->GetEntry(i+1);
        jd = s_inc;
		t->GetEntry(i);
		if(s_qs==0 && s_orbit==orbit) {
		index = -1; // make sure it's initialized
        	for(int j=ch1; j<ch2; j++) {
            	if(se[j]>0) {
            		sum += se[j];
            		if(s_inc<=60) {
            			if(s_inc < jd) {
                            index = 0;
                        } else {
                            index = 3;
                        }
            		} else if(s_inc>60 && s_inc<=90) {
            			if(s_inc < jd) {
                            index = 1;
                        } else {
                            index = 4;
                        }
            		} else if(s_inc>90 && s_inc<=180) {
            			if(s_inc<jd) {
                            index = 2;
            			} else {
                            index = 5;
                        }
            		}
       	 		}
    		}
		if (index == -1) std::cout << "ERROR: plot_angle_flux : index == -1\n";
    		int bin = h[index]->Fill(s_inc, sum);
    		if(bin>0) {
                ct[index][bin]++;
            }
        	sum = 0;
    	}
    	if(s_orbit>orbit) {
            break;
        }
    }
    t->ResetBranchAddresses(); // disconnect from local variables

	for(int k=1; k<=bin3; k++) {
    	for(int m=0; m<6; m++) {
    		if(ct[m][k]>0) {
                h[m]->SetBinContent(k, h[m]->GetBinContent(k)/ct[m][k]);
            }
    	}
    }

	return h;
}

void get_angle_flux_by_energy() {
	int orbit = 756;
	TFile *f = new TFile("sx_2b.root");
	if ((!f) || f->IsZombie()) std::cout << "ERROR: get_angle_flux_by_energy : unusable f\n";
	TTree *t = (TTree*) f->Get("t2");
	if (!t) std::cout << "ERROR: get_angle_flux_by_energy : t == nullptr\n";

	double d_angle = 1.0;
	double d_energy = 0.05;
	TFile *fn = new TFile(Form("data/angle_flux-%.1fdeg-%.2feg-%d.root", d_angle, d_energy, orbit), "recreate");
	if ((!fn) || fn->IsZombie()) std::cout << "ERROR: get_angle_flux_by_energy : unusable fn\n";

	double eg[] = {1.0, 1.2, 1.4, 1.6, 1.8};

	TH1D **h[10];
	for(unsigned int i=0; i<sizeof(eg)/sizeof(eg[0]); i++) {
		fn->mkdir(Form("%.1f", eg[i]));
		fn->cd(Form("%.1f", eg[i]));
		h[i] = plot_angle_flux(t, eg[i], d_angle, d_energy, orbit);
		TCanvas *c2 = new TCanvas(Form("%.1f", eg[i]), " ", 1280, 800);
		c2->Divide(3,2);
		for(int j=0; j<6; j++) {
			c2->cd(j+1);
			h[i][j]->Draw("same");
			h[i][j]->Write();
		}
		c2->Write();
        c2->SaveAs(Form("data/%d_%.1f.pdf", orbit, eg[i]));
        delete h[i];
        // delete c2;
		fn->cd();
	}
	fn->Close();
}