// slurp all needed data into a single ntuple.

#include "InterSplineFunc.h"
#include "Ntuple.h"
#include "options.h"
#include <iostream>
#include <vector>
#include <fstream>
#include <cmath>

#include <cstring>

#define PAWC 1000000
int pawc_[PAWC];

class Collator {
    // ordering: 1,2,3,4=numu,nue,anumu,anue, 1,2,3=(a)nue,(a)numu,(a)nutau
    InterSplineFunc flux_func[4], xsec_func[4], prob_func[4][3];
    InterSplineFunc ncmatrix[20];
    Ntuple* ntuple;
    double emin, emax, estep;
    double scalefactor;
    double resolution;
public:
    Collator() {
        ntuple=0;
        scalefactor=1.0;
        resolution=0.1;         // hard code for now.
    }
    ~Collator() {
        if (ntuple) delete ntuple;
    }

    bool load_flux(string filename);
    bool load_xsec(vector<string> filenames);
    bool load_prob(vector<string> filenames);
    bool load_ncmatrix(string filename);
    bool make_output_ntuple(string filename);
    bool set_energy_rage(vector<double> erange);
    bool set_scalefactor(double sf);
    bool okay(void);
    bool collate(void);

private:
    void event_smear(double normalization,
                     InterSplineFunc& flux,
                     InterSplineFunc& xsec,
                     InterSplineFunc& prob,
                     InterSplineFunc& result);
};

bool Collator::set_scalefactor(double sf)
{
    scalefactor = sf; 
    return true; 
}

void Collator::event_smear(double normalization,
                           InterSplineFunc& flux,
                           InterSplineFunc& xsec,
                           InterSplineFunc& prob,
                           InterSplineFunc& result)
{
    int nbins = (int)((emax-emin)/estep);

    vector<double> energy, smeared;

    for (int ind=0; ind<nbins; ++ind) {
        energy.push_back(ind*estep);
        smeared.push_back(0.0);
    }

    bool use_prob = prob.Initialized();
    for (double e=estep; e<10.0e9; e += estep) {
        double p = use_prob ? prob(e) : 1.0;
        double nevts = flux(e)*xsec(e)*p;
        double spread = resolution*e;
        double norm = nevts*estep/(sqrt(2*M_PI)*spread);
        double exden = -1.0/(2*spread*spread);
        for (int ind=0; ind<nbins; ++ind) {
            double ediff = ind*estep - e;
            smeared[ind] += norm*exp(exden*ediff*ediff);
        }
    }
    for (int ind=0; ind<nbins; ++ind) smeared[ind] *= normalization;

    result.Init(energy,smeared);
}


bool Collator::okay()
{
    if (!ntuple) {
        cerr  << "Collator: no ouput ntuple.\n";
        return false;
    }
    for (int i=0; i<4; ++i) {
        if (!flux_func[i].Initialized()) { 
            cerr << "Flux " << i << " not initialized\n";
            return false;
        }
        if (!xsec_func[i].Initialized()) {
            cerr << "Xsec " << i << " not initialized\n";
            return false;
        }
        for (int j=0; j<3; ++j) {
            if (!prob_func[i][j].Initialized()) {
                cerr << "Prob " << i << "," << j << " not initialized\n";
                return false;
            }
        }
    }

    for (int i=0; i<20; ++i) {
        if (!ncmatrix[i].Initialized()) {
            cerr << "NC matrix " << i << " not initialized\n";
            return false;
        }
    }
    return true;
}

bool Collator::collate(void)
{

    vector<double> data;    // filling this must be in sync with tags def.

    InterSplineFunc smeared[4][2], unosc_smeared[4], dummy;

    // convert flux[nu/year/cm^2]*xsec[10^-38cm^2*targets/water]
    // to interactions/year/kton.  Assumes xsec is cross section 
    // per water molecule.
    double normalization = 6.023e-6;

    // smear events for nu{e,mu}->nu{e,mu}
    for (int nu=0; nu<4; ++nu)
        for (int lep=0; lep<2; ++lep)
            event_smear(normalization,
                        flux_func[nu],
                        xsec_func[lep],
                        prob_func[nu][lep],
                        smeared[nu][lep]);

    // smear unoscillated events
    for (int nu=0; nu<4; ++nu) 
        event_smear(normalization,
                    flux_func[nu],
                    xsec_func[nu],
                    dummy,
                    unosc_smeared[nu]);

    // convert NC matrix to NC events
    vector<double> ncenergy, ncevents;
    for (int ncind=0; ncind<20; ++ncind) {
        double n=0;
        for (int muccind=0; muccind<20; ++muccind) {
            double enumu = muccind*0.5e9;
            n += normalization*flux_func[0](enumu)*xsec_func[0](enumu)*ncmatrix[ncind](enumu);
        }
        ncevents.push_back(n);
        ncenergy.push_back(ncind*0.5e9);
    }
    InterSplineFunc ncevent_func(ncenergy,ncevents);

    // Must keep this filling in sync with tags def
    for (double energy = emin; energy < emax; energy += estep) {
        data.clear();

        // energy
        data.push_back(energy);

        // fluxen
        for (int ind=0; ind<4; ++ind) 
            data.push_back(flux_func[ind](energy));

        // cross sections
        for (int ind=0; ind<4; ++ind) 
            data.push_back(xsec_func[ind](energy));

        // probabilities 4x3
        for (int ini=0; ini<4; ++ini) 
            for (int fin=0; fin<3; ++fin) 
                data.push_back(prob_func[ini][fin](energy));
        
        // smeared, oscillated QE CC events, 
        for (int nu=0; nu<4; ++nu) for (int lep=0; lep<2; ++lep)
            data.push_back(smeared[nu][lep](energy));
        
        // smeared, unoscillated QE CC events,
        for (int nu=0; nu<4; ++nu)
            data.push_back(unosc_smeared[nu](energy));

        // NC backgrounds events
        data.push_back(ncevent_func(energy));

        float d[512];
        for (unsigned int ind=0; ind<data.size(); ++ind) 
            d[ind] = data[ind];
        ntuple->FillRow(d);
    }
    return true;
}

bool Collator::load_xsec(vector<string> filenames)
{
    if (filenames.size() != 4) {
        cerr << "Need 4 xsec files: CC QE for numu, nue, anumu, anue\n";
        cerr << "Only got: " << filenames.size() << endl;
        for (unsigned int ind=0;ind<filenames.size(); ind++) 
            cerr << "\t" << filenames[ind] << endl;
        return false;
    }

    double targets_per_water[4] = { 10.0/18.0, 10.0/18.0, 8.0/18.0, 8.0/18.0 };
    for (int ind=0; ind<4; ++ind) {
        ifstream fstr(filenames[ind].c_str());
        if (!fstr) {
            cerr << "Failed to open file: \"" << filenames[ind] << "\"\n";
            return false;
        }

        vector<double> evec,xsec;

        // Read in xsec.
        double energy=0;
        while (fstr >> energy) {
            evec.push_back(energy);
            double x;
            fstr>>x;
            x *= targets_per_water[ind];
            xsec.push_back(x);
        }
        
        xsec_func[ind].Init(evec,xsec);
    }
    return true;
}
bool Collator::load_prob(vector<string> filenames)
{
    if (filenames.size() != 4) {
        cerr << "Need 4 prob files: X->nu{e,mu,tau} with X=numu,nue,anumu,anue\n";
        return false;
    }

    for (int ind=0; ind<4; ++ind) {
        ifstream fstr(filenames[ind].c_str());
        if (!fstr) {
            cerr << "Failed to open file: \"" << filenames[ind] << "\"\n";
            return false;
        }

        vector<double> evec,probe,probmu,probtau;

        // Read in prob.
        double energy=0;
        while (fstr >> energy) {
            evec.push_back(energy);
            double p;
            fstr>>p;
            probe.push_back(p);
            fstr>>p;
            probmu.push_back(p);
            fstr>>p;
            probtau.push_back(p);
        }
        
        prob_func[ind][0].Init(evec,probe);
        prob_func[ind][1].Init(evec,probmu);
        prob_func[ind][2].Init(evec,probtau);
    }
    return true;
}

bool Collator::load_flux(string filename)
{
    ifstream fstr(filename.c_str());
    if (!fstr) {
        cerr << "Failed to open file: \"" << filename << "\"\n";
        return false;
    }

    vector<double> evec,flux[4];

    // Read in flux.  
    double energy=0;
    while (fstr >> energy) {
        evec.push_back(energy);
        for (int i=0;i<4;++i) {
            double f;
            fstr>>f;
            f *= scalefactor;
            flux[i].push_back(f);
        }
    }

    for (int i=0; i<4; ++i) flux_func[i].Init(evec,flux[i]);
    return true;
}

bool Collator::load_ncmatrix(string filename)
{
    ifstream fstr(filename.c_str());
    if (!fstr) {
        cerr << "Failed to open file: \"" << filename << "\"\n";
        return false;
    }

    vector<double> data, energy;

    // Read in nc matrix.  assumed 400 numbers with Enumuccqe fasted changing index.
    double x=0;
    for (double e=0; e<10e9; e+=0.5e9) energy.push_back(e);
    for (int i=0; i<20; ++i) {
        data.clear();
        for (int j=0; j<20; ++j) {
            fstr>>x;
            data.push_back(x);
        }
        ncmatrix[i].Init(energy,data);
    }

    return true;
}

bool Collator::set_energy_rage(vector<double> erange)
{
    if (erange.size() != 3) {
        cerr << "Collator: bad size for energy range: " << erange.size() << endl;
        return false;
    }

    emin = erange[0];
    emax = erange[1];
    estep = erange[2];
    return true;
}

bool Collator::make_output_ntuple(string filename)
{
    vector<string> tags;
    tags.push_back("ev");       // neutrino energy in eV

    tags.push_back("fnumu");    // input numu flux 
    tags.push_back("fnue");     // input nue flux
    tags.push_back("fanumu");   // input anti numu flux
    tags.push_back("fanue");    // input anti nue flux

    tags.push_back("xnumucc");  // input numu CC QE cross section
    tags.push_back("xnuecc");   // input nue CC QE cross section
    tags.push_back("xanumucc"); // input antinumu  CC QE cross section
    tags.push_back("xanuecc");  // input antinue  CC QE cross section

    tags.push_back("pue");      // input P(numu->nue) prob
    tags.push_back("puu");      // input P(numu->numu) prob
    tags.push_back("put");      // input P(numu->nutau) prob
    tags.push_back("pee");      // input P(nue->nue) prob
    tags.push_back("peu");      // input P(nue->numu) prob
    tags.push_back("pet");      // input P(nue->nutau) prob

    tags.push_back("paue");     // input P(anumu->anue) prob
    tags.push_back("pauu");     // input P(anumu->anumu) prob
    tags.push_back("paut");     // input P(anumu->anutau) prob
    tags.push_back("paee");     // input P(anue->anue) prob
    tags.push_back("paeu");     // input P(anue->anumu) prob
    tags.push_back("paet");     // input P(anue->anutau) prob

    tags.push_back("senumu");   // #events 10% smeared, oscillated e- (from numu)
    tags.push_back("smunumu");  // #events 10% smeared, oscillated mu- (from numu)
    tags.push_back("senue");    // #events 10% smeared, oscillated e- (from nue)
    tags.push_back("smunue");   // #events 10% smeared, oscillated mu- (from nue)

    tags.push_back("saenumu");   // #events 10% smeared, oscillated e+ (from antinumu) events
    tags.push_back("samunumu");  // #events 10% smeared, oscillated mu+ (from antinumu)
    tags.push_back("saenue");    // #events 10% smeared, oscillated e+ (from antinue) events
    tags.push_back("samunue");   // #events 10% smeared, oscillated mu+ (from antinue)

    tags.push_back("snumu");    // #events 10% smeared, unoscilated mu- from numu
    tags.push_back("snue");    // #events 10% smeared, unoscilated e- from nue
    tags.push_back("sanumu");    // #events 10% smeared, unoscilated mu+ from anumu
    tags.push_back("sanue");    // #events 10% smeared, unoscilated e+ from anue

    tags.push_back("sncnue");   // #events NC electron events based on  smeared, oscillated mu-

    Ntuple::Hlimit(PAWC);
    ntuple = new Ntuple(filename.c_str(),1,"Collator",tags);
    return true;
}

static void usage(Options& options, const char* msg)
{
    cerr << "Error: " << msg << endl;
    options.usage(cerr,"");
}



static string stripws(string s)
{
    unsigned int sp = 0;
    while (s[sp] == ' ') ++sp;
    if (sp == s.size()-1) return "";
    return s.substr(sp,s.size());
}

vector<string> str2vec(string s)
{
    unsigned int begin = 0;
    unsigned int comma = s.find(" ");
    vector<string> v;
    while (begin != string::npos) {
        string ss = s.substr(begin,comma-begin).c_str();
        v.push_back(stripws(ss));
        begin = comma;
        comma = s.find(" ",begin+1);
    }
    return v;
}

vector<double> str2doublevec(string s)
{
    unsigned int begin = 0;
    unsigned int comma = s.find(" ");
    vector<double> v;
    while (begin != string::npos) {
        v.push_back(atof(s.substr(begin,comma-begin).c_str()));
        begin = comma;
        comma = s.find(" ",begin+1);
    }
    return v;
}

static bool parse_args(int& argc, const char** argv, Collator& collator)
{
    const char * optv[] = {
        "s:scalefactor <arb scale factor multiplied to flux, must come before -f option>",
        "f:flux <flux table file. Columns: GeV,numu,nue,anumu,anue nu/cm^2/year scaled to 1 km>",
        "x:xsec <\"numuccqe nueccqe anumuccqe anueccqe\" Columns: GeV,xsec 10^-38cm^2/target>",
        "p:prob <\"numu nue anumu anue\" Columns: eV,nuxnue,nuxnumu,nuxnutau>",
        "n:ncmatrix <NC matrix relating numuccqe to NC elike>",
        "o:output <output ntuple file>",
        "e:energy <\"Emin Emax Estep\" E(eV) in [Emin, Emin+Estep, ..., Emax)>",
        0
    };
    
    Options options(*argv, optv);
    OptArgvIter optitr(argc-1,argv+1);
    const char* optarg;

    vector<const char*> newargv;
    newargv.push_back(argv[0]);

    char optchar;
    options.ctrls(Options::PARSE_POS);
    bool scalefactorset=false;

    while ((optchar = options(optitr,optarg))) {
        switch (optchar) {
        case 'f':
            assert(scalefactorset);
            if (!collator.load_flux(optarg)) {
                usage(options,"Bad flux file");
                return false;
            }
            break;

        case 'x':
            if (!collator.load_xsec(str2vec(optarg))) {
                usage(options,"Bad xsec files");
                return false;
            }
            break;

        case 'p':
            if (!collator.load_prob(str2vec(optarg))) {
                usage(options,"Bad prob files");
                return false;
            }
            break;

        case 'n':
            if (!collator.load_ncmatrix(optarg)) {
                usage(options,"Bad NC matrix file");
                return false;
            }
            break;

        case 'o':
            if (!collator.make_output_ntuple(optarg)) {
                usage(options,"Bad ntuple filename");
                return false;
            }
            break;

        case 'e':
            if (!collator.set_energy_rage(str2doublevec(optarg))) {
                usage(options,"Bad energy range");
                return false;
            }
            break;

        case 's':
            if (!collator.set_scalefactor(atof(optarg))) {
                usage(options,"Bad scalefactor");
                return false;
            }
            scalefactorset=true;
            break;

        case Options::POSITIONAL:
            newargv.push_back(optarg);
            break;

        default:
            usage(options,"Unknown option");
            break;
        }
    }

    if (!collator.okay()) {
        usage(options,"Collator not configured, missing option\n");
        return false;
    }

    argc =  newargv.size();

//    if (argc == 1) usage(options,"No input files");

    for (int i = 0; i < argc; ++i) {
        argv[i] = newargv[i];
    }
    return true;
}

int main (int argc, const char *argv[])
{
    Collator collator;
    if (!parse_args(argc,argv,collator)) return 1;

    collator.collate();


    return 0;
} // end of main()