#include "GnumiNuSource.h"
#include "Ntuple.h"
#include "ntgnumi.h"
#include "reweight.h"

GnumiNuSource::GnumiNuSource(Ntuple& rwn, double* fardet)
    : fNtuple(rwn)
      ,fIndex(0)
      ,fFardet(0)
{
    if (fardet) {
        fFardet = new double[3];
        for (int i = 0; i < 3; ++i) fFardet[i] = fardet[i];
    }
}

GnumiNuSource::~GnumiNuSource()
{
    delete [] fFardet;
}

bool GnumiNuSource::operator()(int& nutype, double& energy, double& weight,
                                double dir[3], double pos[3])
{
    if (fIndex >= fNtuple.GetNentries()) return false;

    gnumi_t* gnt = reinterpret_cast<gnumi_t*>(fNtuple[fIndex]);
    
    switch ((int)gnt->ntype) {
    case 55:                    // anti-numu
        nutype = -2;
        break;
    case 56:                    // numu
        nutype = +2;
        break;
    case 52:                    // anti-nue
        nutype = -1;
        break;
    case 53:                    // nue
        nutype = +1;
        break;
    default:
        nutype = 0;
        break;
    }

    if (fFardet) {
        double vtx[3] = {gnt->vx, gnt->vy, gnt->vz};
        if (pos) for(int i=0; i<3; ++i) pos[i] = vtx[i];
        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(fFardet,vtx,(int)gnt->ntype,p,
                          parent_type, parent_momentum, mu_parent_momentum,
                          energy, weight);
        weight *= gnt->nimpwt;
        if (dir) {
            double tot=0;
            for (int i=0; i<3; ++i) {
                dir[i] = pos[i] - fFardet[i];
                tot += dir[i]*dir[i];
            }
            tot=sqrt(tot);
            for (int i=0; i<3; ++i) dir[i] = dir[i]/tot;
        }
    }
    else {

        energy = gnt->nenergy;
        if (dir) {
            dir[0] = gnt->ndxdz*gnt->npz/energy;
            dir[1] = gnt->ndydz*gnt->npz/energy;
            dir[2] = gnt->npz/energy;
        }
        if (pos) {
            pos[0] = gnt->vx;
            pos[1] = gnt->vy;
            pos[2] = gnt->vz;
        }
        weight = gnt->nimpwt;
    }

    ++fIndex;
    return true;
}