// generate flux tables reweighted to different points, from gnumi ntuples

#include "cfortran.h"
#include "hbook.h"
#include "options.h"
#include "Ntuple.h"
#include "ntgnumi.h"
#include "reweight.h"

#include <vector>
#include <string>
#include <cstdlib>
#include <fstream>
#include <sstream>
#include <cmath>

static const char* outbase = "";
static int nutype = 0;
static int gnumitype = 0;
static double NPOT = 0.0;
static int nbins = 20;
static float minbin = 0.0;
static float maxbin = 10.0;
static const char* save_hists = 0;

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

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

static void parse_args(int& argc, const char** argv, 
                       vector<double>& angle, vector<double>& baseline)
{
    const char * optv[] = {
        "a:angles <\"space separated list of angles (in deg) to use\">",
        "b:baselines <\"space separated list of baselines (in cm) to use\">",
        "t:type <type of neutrino 1,2,3 for e,mu,tau, negative for antinu>",
        "n:npot <total number of POT in all the input files>",
        "o:outbase <base name for output files, if dir, need trailing '/'>",
        "h:histfile <file for histograms>",
        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);
    while ((optchar = options(optitr,optarg))) {
        switch (optchar) {
        case 'a':
            set_vector(angle,optarg);
            break;
        case 'b':
            set_vector(baseline,optarg);
            break;
        case 't':
            nutype = atoi(optarg);
            break;
        case 'o':
            outbase = optarg;
            break;
        case 'n':
            NPOT = atof(optarg);
            break;
        case 'h':
            save_hists = optarg;
            break;
        case Options::POSITIONAL:
            newargv.push_back(optarg);
            break;
        default:
            usage(options,"Unknown option");
            break;
        }
    }

    if (!nutype || !angle.size() || !baseline.size() || !NPOT) {
        usage(options,"Need to set nutype,angles,baselines, and NPOT");
        exit (1);
    }

    switch(nutype) {
    case 1: gnumitype = 53; break;
    case -1: gnumitype = 52; break;
    case 2: gnumitype = 56; break;
    case -2: gnumitype = 55; break;
    }

    argc =  newargv.size();

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

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

}

const char* get_nutype_as_string(int n)
{
    switch (n) {
    case  1: return "nue";       break;
    case -1: return "antinue";   break;
    case  2: return "numu";      break;
    case -2: return "antinumu";  break;
    case  3: return "nutau";     break;
    case -3: return "antinutau"; break;
    }
    return "unknown";
}

void book_hists(vector<double>& angle, vector<double>& baseline)
{
    int asiz = angle.size(), bsiz = baseline.size();
    for (int a = 0; a < asiz; ++a) {
        for (int b = 0; b < bsiz; ++b) {
            int hid = (a+1)*100 + (b+1);
            stringstream ss;
            ss << get_nutype_as_string(nutype) << " ";
            ss << "flux, angle=" << angle[a] << ", baseline=" << baseline[b];
//            cerr << "Filling hist: " << ss.str() << endl;
            char* buf = strdup(ss.str().c_str());
            HBOOK1(hid,buf,nbins,minbin,maxbin,0.0);
        }
    }
}

void get_energy_weight(gnumi_t* gnt,double angle, double baseline,
                       double& weight, double& energy)
{
    double fardet[3] = { 0.0, baseline*angle*M_PI/180.0, baseline };
    double vtx[3] = {gnt->vx, gnt->vy, gnt->vz};
    double p[3] = 
        {gnt->ndxdz*gnt->npz, gnt->ndydz*gnt->npz, gnt->npz};
    int parent_type = (int)gnt->ptype;
    double parent_momentum[3] =
        { gnt->pdpx, gnt->pdpy, gnt->pdpz };
    double mu_parent_momentum[3] = 
        { gnt->muparpx, gnt->muparpy, gnt->muparpz };
    
    reweight_neutrino(fardet,vtx,(int)gnt->ntype,p,
                      parent_type, parent_momentum, mu_parent_momentum,
                      energy, weight);
    weight *= gnt->nimpwt;
    
}
void fill_hists(Ntuple& nt, double normalize, 
                vector<double>& angle, vector<double>& baseline)
{
    int asiz = angle.size(), bsiz = baseline.size();
    int nsiz = nt.GetNentries();
    double weight, energy;
    for (int ind = 0; ind < nsiz; ++ind) {
        gnumi_t* gnt = reinterpret_cast<gnumi_t*>(nt[ind]);

        // skip all but numu's
        if (((int)gnt->ntype) != gnumitype) continue;

        for (int a = 0; a < asiz; ++a) {
            for (int b = 0; b < bsiz; ++b) {
                int hid = (a+1)*100 + (b+1);


                get_energy_weight(gnt,angle[a],baseline[b],weight,energy);
                HF1(hid,energy,weight*normalize);
            }
        }
    }
    char name[1024] = "gnumi-flux.hst";
    char T[1024] = "T";
    HRPUT(0,name,T);
}

void dump_to_file(float y[], double angle, double baseline)
{
    stringstream ss;
    ss << outbase << get_nutype_as_string(nutype) 
       << "_angle=" << angle
       << "_baseline=" << baseline
       << ".vec";
    ofstream fstr(ss.str().c_str());
    if (!fstr) {
        cerr << "failed to open " << ss.str() << " for writing\n";
        return;
    }
    
    float binsiz = (maxbin - minbin)/(1.0*nbins);
    for (int i=0; i<nbins; ++i)
        fstr << 1.0e9*binsiz*(i+0.5) << " " << y[i] << endl;
    
}

void dump_hists(vector<double>& angle, vector<double>& baseline)
{
    int asiz = angle.size(), bsiz = baseline.size();
    for (int a = 0; a < asiz; ++a) {
        for (int b = 0; b < bsiz; ++b) {
            int hid = (a+1)*100 + (b+1);
            
            float cont[nbins];
            char blank[10] = "";
            HUNPAK(hid,cont,blank,0);
            dump_to_file(cont,angle[a],baseline[b]);
        }
    }    
    if (save_hists) {
        char tmp1[128] = "Fluxes", N[128]="N",space[128]=" ";
        char fname[128];
        sprintf(fname,"%s",save_hists);
        int one=1, zero=0, tentwentyfour=1024;
        HROPEN(one,tmp1,fname,N,tentwentyfour,zero);
        HROUT(zero,zero,space);
        HREND(tmp1);
    }
        
}

#define PAWC_SIZE 1000000
float pawc_[PAWC_SIZE];

int main (int argc, const char** argv)
{
    vector<double> angle, baseline;
    parse_args(argc,argv,angle,baseline);

    HLIMIT(PAWC_SIZE);

    book_hists(angle,baseline);

    // Gnumi gives total neutrinos per m^2
    // normalize to neutrinos/POT/cm^2
    // the M_PI is some gnumi kruft.
    double normalize = (1.0*NPOT)*M_PI*1.0e4;
    cerr << "NPOT=" << NPOT << ", norm = " << normalize << endl;
    normalize = 1.0/normalize;

    for (int fnum = 1; fnum < argc; ++fnum) {
        cerr << argv[fnum] << endl;
        Ntuple nt(argv[fnum],10);
//        int nent = nt.GetNentries();

        fill_hists(nt,normalize,angle,baseline);
        
    }
    dump_hists(angle,baseline);
    return 0;
} // end of main()