//Renata Kopecna

#include <folder.hh>
#include <TRandom3.h>
#include <event.hh>
#include <helpers.hh>

#include <spdlog.h>
///fold the set of angles according to the chosen folding scheme
void fcnc::folder::fold(event* e){
    //make sure, that the full angular information was saved before into the correct parameters:
    assert(e->costhetal == e->costhetal_fa);
    assert(e->costhetak == e->costhetak_fa);
    assert(e->phi == e->phi_fa);

    double pi = TMath::Pi();
    switch(this->scheme){
    case -1:
        return;
    case 0:
        if(e->phi < 0)e->phi = e->phi + pi;
        break;
    case 1:
        if(e->phi < 0){
            e->phi = -e->phi;
        }
        if(e->costhetal < 0){
            e->phi = pi - e->phi;
            e->costhetal = -e->costhetal;
        }
        break;
    case 2:
        if(e->phi < 0)       e->phi = -e->phi;

        if(e->costhetal < 0) e->costhetal = -e->costhetal;
        break;
    case 3:
        if(e->phi > pi/2.)   e->phi = pi - e->phi;
        if(e->phi < -pi/2.)  e->phi = -pi - e->phi;

        if(e->costhetal < 0) e->costhetal = -e->costhetal;
        break;
    case 4:
        if(e->phi > pi/2.)   e->phi = pi - e->phi;
        if(e->phi < -pi/2.)  e->phi = -pi - e->phi;

        if(e->costhetal < 0){
            e->costhetal = -e->costhetal;
            e->costhetak = -e->costhetak;
        }

        break;
    }
    return;
}

///inverse fold the set of angles according to the chosen folding scheme
void fcnc::folder::invers_fold(const event* e, event *u_phi, event *u_ctl, event *u_full){

    /*
    phi /^\
         |         :
         |   (2)   :   (4)
         |         :
         |---------+--------
         |         :
         |   (1)   :   (3)
         |_________________>
                       cos(theta)

    (1): original event (e)
    (2): unfolded in phi (u_phi)
    (3): unfolded in ctl (u_ctl)
    (4): twice unfolded (u_full)

    */
    *u_phi  = *e; //only inversly folded in dependence of phi
    *u_ctl  = *e; //only inversly folded in dependence of ctl
    *u_full = *e; //full inversly folded

    double pi = TMath::Pi();
    switch(this->scheme){
    case -1:
        return;
    case 0:
        assert(e->phi > 0.0);

        u_phi->phi  = u_phi->phi  - pi;
        u_full->phi = u_full->phi - pi;

        break;
    case 1:
        assert(e->phi > 0.0);
        assert(e->costhetal > 0.0);

        u_ctl->phi  = pi - u_ctl->phi;
        u_full->phi = pi - u_full->phi;
        u_ctl->costhetal  = -u_ctl->costhetal;
        u_full->costhetal = -u_full->costhetal;

        u_phi->phi  = -u_phi->phi;
        u_full->phi = -u_full->phi;

        break;
    case 2:
        assert(e->phi > 0.0);
        assert(e->costhetal > 0.0);

        u_phi->phi  = -u_phi->phi;
        u_full->phi = -u_full->phi;

        u_ctl->costhetal  = -u_ctl->costhetal;
        u_full->costhetal = -u_full->costhetal;

        break;
    case 3:
        assert(e->phi <  pi/2.);
        assert(e->phi > -pi/2.);
        assert(e->costhetal > 0.0);

        if(e->phi > 0.0){
            u_phi->phi  =  pi - u_phi->phi;
            u_full->phi =  pi - u_full->phi;
        }
        else{
            u_phi->phi  = -pi - u_phi->phi;
            u_full->phi = -pi - u_full->phi;
        }

        u_ctl->costhetal  = -u_ctl->costhetal;
        u_full->costhetal = -u_full->costhetal;

        break;
    case 4:
        assert(e->phi <  pi/2.);
        assert(e->phi > -pi/2.);
        assert(e->costhetal > 0.0);

        if(e->phi > 0.0){
            u_phi->phi  =  pi - u_phi->phi;
            u_full->phi =  pi - u_full->phi;
        }
        else{
            u_phi->phi  = -pi - u_phi->phi;
            u_full->phi = -pi - u_full->phi;
        }

        u_ctl->costhetal  = -u_ctl->costhetal;
        u_ctl->costhetak  = -u_ctl->costhetak;
        u_full->costhetal = -u_full->costhetal;
        u_full->costhetak = -u_full->costhetak;

        break;
    }
}

void fcnc::folder::test_inv_folding(){

    fcnc::event e;
    e.m = 5286.0;
    TRandom3 * r = new TRandom3(0);
    for(int i = 0; i < 10000; i++){

        if(i % 100 == 0 && spdlog_info())std::cout << i / 100. << "%" << std::endl;

        //generate a folded event:
        e.costhetal = r->Rndm() * 2. - 1.;
        e.costhetak = r->Rndm() * 2. - 1.;
        e.phi       = (r->Rndm() * 2. - 1.) * TMath::Pi();
        e.costhetal_fa = e.costhetal;
        e.costhetak_fa = e.costhetak;
        e.phi_fa       = e.phi;

        this->fold(&e);

        //get 3 events for (partial-)inverse folding:
        fcnc::event inv_e[3];

        this->invers_fold(&e, &inv_e[0], &inv_e[1], &inv_e[2]);

        for(int j = 0; j < 3; j++){
            inv_e[j].costhetal_fa = inv_e[j].costhetal;
            inv_e[j].costhetak_fa = inv_e[j].costhetak;
            inv_e[j].phi_fa       = inv_e[j].phi;
            this->fold(&inv_e[j]);

            assert(abs(e.phi       - inv_e[j].phi      ) < 1e-7);
            assert(abs(e.costhetal - inv_e[j].costhetal) < 1e-7);
            assert(abs(e.costhetak - inv_e[j].costhetak) < 1e-7);
        }
    }
    spdlog::info("[DONE]\t\tInverse folding #{0:d} works!", this->scheme);
}

int fcnc::folder::get_scheme(){
    return this->scheme;
}

void fcnc::folder::set_scheme(int s){
    this->scheme = s;
}