#include <Randomize.hh>

#include <RAT/Quad.hh>
#include <RAT/DB.hh>
#include <RAT/Log.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/DU/PMTInfo.hh>
using namespace RAT;
using namespace RAT::Methods;

#include <TH1D.h>
using namespace ROOT;

#include <CLHEP/Units/SystemOfUnits.h>
#include <math.h>
using namespace CLHEP;
using namespace std;

#ifndef __GNUC__
    #define __restrict__
#endif

static const unsigned int TABCUTOFF = 20;

//
// Every way of choosing numbers from 0 to nhit-1, where nhit is the
// table index. Sizes of these tables:
//   0,    0,    0,    0,    1,    5,   15,   35,   70,  126,
// 210,  330,  495,  715, 1001, 1365, 1820, 2380, 3060, 3876
//
static vector< vector<unsigned int> > combotables[TABCUTOFF];

// We cache of all the PMT positions in the detector here so as not to
// have to look them up repeatedly.
static vector< vector<double> > quadpmtpos;

void Quad::BeginOfRun( DS::Run& )
{
    DB * const db = DB::Get();
    string matname;
    try{
        matname = db->GetLink("GEO","inner_av")->GetS("material");
    }
    catch(DBNotFoundError & ){
        try{
            // If we are in a partial-fill geometry, use the upper
            // material
            matname = db->GetLink("GEO","inner_av")->GetS("material_top");
        }
        catch(DBNotFoundError & ){
            Log::Die("Quad: inner_av material and material_top "
            "undefined. Try something like:\n"
            "/rat/db/set GEO[inner_av] material "
            "\"te_labppo_scintillator\" or "
            "/rat/db/set GEO[inner_av] material_top "
            "\"te_labppo_scintillator\"");
        }
    }

    DBLinkPtr quadDB;
    try{
        (quadDB = db->GetLink("QUAD_FIT", matname))->GetS("index");
    }
    catch(DBNotFoundError& e ){
        warn << "Quad::BeginOfRun: No full (QUAD) definition for "
                + matname + " using defaults instead.\n";
        quadDB = db->GetLink( "QUAD_FIT" );
    }

    try {
        // Overall effective speed of light in mm/ns
        fLightSpeed = quadDB->GetD("light_speed");

        // Limit number of successfully generated points
        fNumQuadPoints = quadDB->GetI("num_points");

        // Limit total number of tries at generating a point
        fLimitQuadPoints = quadDB->GetI("limit_points");

        // Whether to calculate the reconstructed time
        fCalcTime = quadDB->GetI("calc_time");

        // Whether to calculate the errors on time and position
        fCalcErrors = quadDB->GetI("calc_errors");
    }
    catch(const DBNotFoundError & e){
        Log::Die(Form("Quad does not have a definition (or a full "
        "one) for material\n\"%s\" and/or associated settings. It's "
        "missing the field \"%s\".", matname.c_str(), e.field.c_str()));
    }
}

// Squared magnitude of the 3-vector "vec".
static inline double mag2(const double vec[3])
{
    return vec[0]*vec[0] + vec[1]*vec[1] + vec[2]*vec[2];
}

// Subtracts the 3-vectors a-b
static inline void vecsub(double * const __restrict__ ans,
                          const double * const __restrict__ a,
                          const double * const __restrict__ b)
{
    ans[0] = a[0] - b[0];
    ans[1] = a[1] - b[1];
    ans[2] = a[2] - b[2];
}

// Dot product a.b for 3-vectors
static inline double vecdot(const double * const __restrict__ a,
                            const double * const __restrict__ b)
{
    return a[0]*b[0] + a[1]*b[1] + a[2]*b[2];
}

// Invert a 3x3 matrix, "m", returning the answer as "ans". This
// attempts to be highly optimized, minimizing the total number of
// operations by using the explicit solution and avoiding expensive
// divisions as much as possible.
static inline void matinvert(double (* const __restrict__ ans)[3],
                             const double (* const __restrict__ m)[3])
{
    const double idet = 1/(
            + m[0][2]*(m[1][0]*m[2][1] - m[1][1]*m[2][0])
            + m[0][1]*(m[1][2]*m[2][0] - m[1][0]*m[2][2])
            + m[0][0]*(m[1][1]*m[2][2] - m[1][2]*m[2][1]));

    ans[0][0] = (-m[1][2]*m[2][1] + m[1][1]*m[2][2])*idet;
    ans[0][1] = ( m[0][2]*m[2][1] - m[0][1]*m[2][2])*idet;
    ans[0][2] = (-m[0][2]*m[1][1] + m[0][1]*m[1][2])*idet;
    ans[1][0] = ( m[1][2]*m[2][0] - m[1][0]*m[2][2])*idet;
    ans[1][1] = (-m[0][2]*m[2][0] + m[0][0]*m[2][2])*idet;
    ans[1][2] = ( m[0][2]*m[1][0] - m[0][0]*m[1][2])*idet;
    ans[2][0] = (-m[1][1]*m[2][0] + m[1][0]*m[2][1])*idet;
    ans[2][1] = ( m[0][1]*m[2][0] - m[0][0]*m[2][1])*idet;
    ans[2][2] = (-m[0][1]*m[1][0] + m[0][0]*m[1][1])*idet;
}

// Choose four distinct random numbers less than nhit
static vector<unsigned int> ChooseRandomHits(const unsigned int nhit)
{
    vector<unsigned int> pmti(4);

    for(int j = 0; j < 4; j++){
        while(true){
            pmti[j] = CLHEP::RandFlat::shootInt(nhit);
            bool same = false;
            for(int k = 0; k < j; k++)
                if(pmti[j] == pmti[k])
                    same = true;
            if(!same) break;
        }
    }

    return pmti;
}

// Create a table of all the ways to pick 4 numbers out of n
static vector< vector<unsigned int> > buildtable(const unsigned int n)
{
    if(n >= TABCUTOFF)
        Log::Die("Quad: tried to make a table bigger than expected!\n");

    vector< vector<unsigned int> > table;

    vector<unsigned int> entry(4);
    for(entry[0] = 0;          entry[0] < n; entry[0]++)
        for(entry[1] = entry[0]+1; entry[1] < n; entry[1]++)
            for(entry[2] = entry[1]+1; entry[2] < n; entry[2]++)
                for(entry[3] = entry[2]+1; entry[3] < n; entry[3]++)
                    table.push_back(entry);

    return table;
}

// Return the iterationth way of choosing 4 numbers less than nhit
static vector<unsigned int>
ChooseOrderedHits(const unsigned int nhit, const unsigned int iteration)
{
    if(combotables[nhit].size() > iteration)
        return combotables[nhit][iteration];
    Log::Die(Form("Quad: out of bounds (%d %d) in ChooseOrderedHits. "
                  "This is a bug.", nhit, iteration));
}

// Choose hits, either from a table if the number of hits is small,
//   or randomly if it is large.
static vector<unsigned int> ChooseHits(const unsigned int nhit,
                                       const unsigned int iteration)
{
    if(nhit < TABCUTOFF) return ChooseOrderedHits(nhit, iteration);
    else                 return ChooseRandomHits(nhit);
}

DS::FitResult Quad::GetBestFit()
{
    const DU::PMTInfo& pmtInfo = DU::Utility::Get()->GetPMTInfo();
    if(quadpmtpos.empty()){
        info << "Quad filling pmt positions and combinatoric tables\n";

        for(unsigned int n = 4; n < TABCUTOFF; n++)
            combotables[n] = buildtable(n);

        if(!fRun) Log::Die("Quad: Run is NULL, can't do anything!\n");

        for(size_t t=0; t< pmtInfo.GetCount(); t++){
            const TVector3 p = pmtInfo.GetPosition(t);
            vector<double> vp;
            vp.push_back(p[0]);
            vp.push_back(p[1]);
            vp.push_back(p[2]);
            quadpmtpos.push_back(vp);
        }
    }

    // PMT Selector requires a seed vertex.
    // However, Quad doesn't so create a dummy vertex and use this.
    DS::FitVertex dummyVertex;
    dummyVertex.SetPosition(TVector3(0,0,0),0);
    dummyVertex.SetTime(0,0);
    SelectPMTData( dummyVertex );
    dummyVertex.SetPositionValid(0);
    fFitResult.SetVertex( 0, dummyVertex );

    const size_t PMThits = fSelectedPMTData.size();
    if(PMThits < 4) return fFitResult;

    // The positions of a single quadruplet of hits.
    double pmtPos[4][3];

    const unsigned int npoints_to_try = PMThits < TABCUTOFF?
        combotables[PMThits].size(): fLimitQuadPoints;

    // These arrays hold the cloud of quad points.
    std::vector<double> rx, ry, rz, rt;
    rx.reserve(fNumQuadPoints);
    ry.reserve(fNumQuadPoints);
    rz.reserve(fNumQuadPoints);
    if(fCalcTime) rt.reserve(fNumQuadPoints);

    // This is the main loop of the algorithm that calculates the cloud
    // of quad points.
    for(unsigned int endCount=0; endCount < npoints_to_try; endCount++){
        double minTime = 1e9;
        double t[4] = {0};

        const vector<unsigned int> pmti = ChooseHits(PMThits, endCount);

        for(int j = 0; j < 4; j++){
            const FitterPMT & pmt = fSelectedPMTData[pmti[j]];
            t[j] = fLightSpeed*pmt.GetTime();
            if(t[j] < minTime) minTime = t[j];

            for(int c = 0; c < 3; c++)
                pmtPos[j][c] = quadpmtpos[pmt.GetID()][c];
        }

        double rsq[4];
        double N[3], K[3], g[3], h[3], bv[3];
        double M[3][3], iM[3][3];

        // Now do the calculation
        for(int j = 0; j < 4; j++)
            rsq[j] = mag2(pmtPos[j]) - t[j]*t[j];

        for(int k = 0; k < 3; k++){
            M[k][0] = pmtPos[k+1][0] - pmtPos[0][0];
            M[k][1] = pmtPos[k+1][1] - pmtPos[0][1];
            M[k][2] = pmtPos[k+1][2] - pmtPos[0][2];
            N[k] = t[k+1] - t[0];
            K[k] = (rsq[k+1] - rsq[0])/2;
        }

        matinvert(iM, M);
        for(int w = 0; w < 3; w++){
            g[w] = iM[w][0]*K[0] + iM[w][1]*K[1] + iM[w][2]*K[2];
            h[w] = iM[w][0]*N[0] + iM[w][1]*N[1] + iM[w][2]*N[2];
        }

        vecsub(bv, pmtPos[0], g);
        const double a = mag2(h)-1;
        const double b = -2*(vecdot(bv, h) - t[0]);
        const double c = mag2(bv) - t[0]*t[0];

        const double s = b*b - 4*a*c;

        if(s > 0){
            // Only one root is used. The other represents light
            // propagating backwards in time and is unphysical. I think.
            const double tau = (-b + sqrt(s))/(2*a);
            if(tau < minTime){
                double v[3];
                for(int nt = 0; nt < 3; nt++)
                    v[nt] = g[nt] + h[nt]*tau;

                // Maximum radius (mm) for which we will accept a point.
                // This is meant to remove wild fits, not directly
                // restrict them to the physically allowed region.
                // Indeed, if we cut this off at the PMT sphere surface,
                // it would introduce a large bias when we took the
                // median later.  So it is a little larger than that.
                const double rlimit = 9000;
                if(mag2(v) < rlimit*rlimit){
                    if(fCalcTime) rt.push_back(tau/fLightSpeed);
                    rx.push_back(v[0]);
                    ry.push_back(v[1]);
                    rz.push_back(v[2]);
                    if(rx.size() >= fNumQuadPoints) break; // got enough
                }
            }
        }
    }

    TVector3 fitPos;
    double fitTime = 0, fitPosErrorX=0,
           fitPosErrorY=0, fitPosErrorZ=0, fitTimeError = 0;
    const unsigned int size = rx.size();
    if(size){
        if(fCalcErrors){
            sort(rx.begin(), rx.end());
            sort(ry.begin(), ry.end());
            sort(rz.begin(), rz.end());
            if(fCalcTime) sort(rt.begin(), rt.end());

            // Find errors by calculating the width of the quad cloud
            fitPosErrorX = (rx[3*size/4]-rx[size/4])/sqrt(size);
            fitPosErrorY = (ry[3*size/4]-ry[size/4])/sqrt(size);
            fitPosErrorZ = (rz[3*size/4]-rz[size/4])/sqrt(size);
            if(fCalcTime)
                fitTimeError = (rt[3*size/4]-rt[size/4])/sqrt(size);
        }
        else{
            // If we only need the median, we can use nth_element,
            // which is O(N) rather than sort, which is O(NlogN)
            nth_element(rx.begin(), rx.begin()+size/2, rx.end());
            nth_element(ry.begin(), ry.begin()+size/2, ry.end());
            nth_element(rz.begin(), rz.begin()+size/2, rz.end());
            if(fCalcTime)
                nth_element(rt.begin(), rt.begin()+size/2, rt.end());
        }

        // Get median position and time
        if(fCalcTime) fitTime = rt[size/2];
        fitPos.SetXYZ(rx[size/2], ry[size/2], rz[size/2]);
    }
    else
        throw MethodFailError("Quad ERROR: no valid points in cloud");

    DS::FitVertex fitVertex;
    fitVertex.SetPosition( fitPos );
    fitVertex.SetPositionErrors( ROOT::Math::XYZVectorF(
        fitPosErrorX, fitPosErrorY, fitPosErrorZ) );
    fitVertex.SetTime( fitTime );
    fitVertex.SetTimeErrors( fitTimeError );
    if(!fCalcTime)
        fitVertex.SetTimeValid(false);
    fFitResult.SetFOM("numPoints", size );
    fFitResult.SetVertex( 0, fitVertex );
    return fFitResult;
}