#include "ITriggerCandidate.hh"
#include "ICDCGeomId.hxx"
#include "ICTHGeomId.hxx"
#include "IFieldManager.hxx"
#include "IG4HitGas.hxx"
#include "IGeomIdManager.hxx"
#include "IOADatabase.hxx"

#include "TGeoManager.h"
#include "TGeoMatrix.h"
#include "TGeoBBox.h"


typedef COMET::IG4VHit Hit;
typedef COMET::IG4HitGas HitGas;
typedef COMET::IG4HitContainer Hits;
typedef COMET::IG4TrajectoryContainer Tracks;


static TGeoMatrix * gMatrix = nullptr;
static double gSTRadius = 0.;
static const char * gTag = "MCSelector";
static int gNbOfCTHSegments = CTH_MAX_SEGMENTS;


const char *& ITriggerCandidate::Tag()
{
        return gTag;
}


bool ITriggerCandidate::UpdateGeometry()
{
        /* Parse the path to the Detector Solenoid node */
        auto id = COMET::GeomId::CTH::Detector();
        auto & geoMgr = COMET::IOADatabase::Get().GeomId();
        auto cthPath = geoMgr.FullNodePath(id.GetPosition());

        std::string pattern = "DetectorSolenoid_0";
        auto index = cthPath.find(pattern);
        if (index == std::string::npos) {
                COMETNamedError(gTag,
                    "Could not find the " << pattern << " geometry");
                return false;
        }
        index += pattern.size();
        auto path = cthPath.substr(0, index);

        /* Get the local transform */
        auto g = COMET::IOADatabase::Get().GetGeometry();
        g->cd(path.c_str());
        auto matrix = g->GetCurrentNode()->GetMatrix();

        /* Get the extent of the stopping target disks */
        auto t = matrix->GetTranslation();
        auto shape = g->FindNode(t[0], t[1], t[2])->GetVolume()->GetShape();
        auto bbox = dynamic_cast<TGeoBBox *>(shape);
        const double u = geoMgr.GetRootUnitLength();
        gSTRadius = bbox->GetDX() * u;

        COMETNamedInfo(gTag, "Geometry updated");
        COMETNamedInfo(gTag, "SD translation : [" << t[0] * u << ", " <<
                             t[1] * u << ", " << t[2] * u << "] mm");
        COMETNamedInfo(gTag, "ST radius      : " << gSTRadius << " mm");

        /* Change the units of the ROOT transform to COMET ones*/
        gMatrix = matrix->MakeClone();
        gMatrix->SetDx(t[0] * u);
        gMatrix->SetDy(t[1] * u);
        gMatrix->SetDz(t[2] * u);

        /* Get the number of CTH segments */
        path = cthPath.substr(0, cthPath.size() - 1);
        int segments;
        for (segments = 0; segments < CTH_MAX_SEGMENTS; segments++) {
                auto s = path + std::to_string(segments);
                if (!g->cd(s.c_str())) break;
        }
        gNbOfCTHSegments = segments;

        return true;
}


ITriggerCandidate::ITriggerCandidate(int tid, int pid) : trackId(tid),
    pdgId(pid), cdcLayerMin(21), cdcLayerMax(-1), crvTracks(-1),
    cdcZMin(DBL_MAX), cdcZMax(-DBL_MAX), cdcMomentumMin(DBL_MAX),
    cdcMomentumMax(-DBL_MAX), cdcMomentumMean(0), cdcT0(DBL_MAX),
    cdcT1(-DBL_MAX), crvT0(DBL_MAX), crvT1(-DBL_MAX), cthT0(DBL_MAX),
    cthT1(-DBL_MAX), cdcFit({ 0, 0, 0, 0, 0, 0})
{
        memset(fTrigger, 0x0, sizeof(fTrigger));
}


ITriggerCandidate::~ITriggerCandidate() 
{
        cthHits.clear();
        cdcHits.clear();
}


void ITriggerCandidate::UpdateCTH(
    const COMET::IGeometryId & id, const Hit * hit)
{
        const auto t = hit->GetPosT();
        if (t < this->cthT0) this->cthT0 = t;
        if (t > this->cthT1) this->cthT1 = t;

        const auto i = COMET::GeomId::CTH::GetModule(id);
        const auto j = COMET::GeomId::CTH::GetScintillator(id);
        const auto k = COMET::GeomId::CTH::GetSegmentId(id);
        this->fTrigger[i][j][k] = 1;
}


bool ITriggerCandidate::FinalizeCTH()
{
        for (int i = 0; i < 2; i++) {
                for (int k = 0; k < gNbOfCTHSegments; k++) {
                        int k1 = (k + 1) % gNbOfCTHSegments;
                        if (fTrigger[i][0][k] && fTrigger[i][0][k1] &&
                            fTrigger[i][1][k] && fTrigger[i][1][k1]) {
                                return true;
                        }
                }
        }
        return false;
}


void ITriggerCandidate::UpdateCDC(const COMET::IGeometryId & id, const Hit * hit)
{
        const auto layer = COMET::GeomId::CDC::GetLayerId(id); 
        if (layer < this->cdcLayerMin) {
                this->cdcLayerMin = layer;
        }
        if (layer > this->cdcLayerMax) {
                this->cdcLayerMax = layer;
        }

        const Double_t master[] = { hit->GetPosX(),
                                    hit->GetPosY(),
                                    hit->GetPosZ(), 1. };
        Double_t local[4];
        gMatrix->MasterToLocal(master, local);
        const auto z = local[2] * 1E-03;
        if (z < this->cdcZMin) {
                this->cdcZMin = z;
        }
        if (z > this->cdcZMax) {
                this->cdcZMax = z;
        }

        const auto t = hit->GetPosT();
        if (t < this->cdcT0) {
                this->cdcT0 = t;
        }
        if (t > this->cdcT1) {
                this->cdcT1 = t;
        }

        const auto p = hit->GetMomentumMag();
        if (p < this->cdcMomentumMin) this->cdcMomentumMin = p;
        if (p > this->cdcMomentumMax) this->cdcMomentumMax = p;
        this->cdcMomentumMean += p;
}


bool ITriggerCandidate::FinalizeCDC(double momentumMin, double momentumMax,
    bool checkTrack, bool dumpJson)
{
        const auto n = this->cdcHits.size();
        if (n == 0) return false;

        this->cdcMomentumMean /= n;
        if ((this->cdcMomentumMean < momentumMin) ||
            (this->cdcMomentumMean > momentumMax)) {
                return false;
        }

        if (checkTrack) {
                if ((this->cdcLayerMin != 0) || (this->cdcLayerMax < 4) ||
                    (this->cdcLayerMax >= 19))
                        return false;
        } else if (!dumpJson) {
                return true;
        }

        /* Fit of the CDC trajectory with an helice: compute the intermediate
         * sum factors for the radius fit
         */
        double u, v, w;
        auto ToLocal = [&u, &v, &w] (const Hit * hit) {
                const Double_t master[] = { hit->GetPosX(),
                                            hit->GetPosY(),
                                            hit->GetPosZ(), 1. };
                Double_t local[4];
                gMatrix->MasterToLocal(master, local);

                u = local[0];
                v = local[1];
                w = local[2];
        };

        double x0 = 0, y0 = 0;
        for (auto& i : this->cdcHits) {
                ToLocal(i);

                x0 += u;
                y0 += v;
        }
        const double n_inv = 1. / this->cdcHits.size();
        x0 *= n_inv;
        y0 *= n_inv;

        double suu = 0, suv = 0, svv = 0;
        double suuu = 0, suvv = 0, svvv = 0, svuu = 0;
        for (auto& i : this->cdcHits) {
                ToLocal(i);
                u -= x0;
                v -= y0;

                suu += u * u;
                suv += u * v;
                svv += v * v;
                suuu += u * u * u;
                suvv += u * v * v;
                svvv += v * v * v;
                svuu += v * u * u;
        }

        /* Fit of the projected hits with a circle, according to R. Bullock.
         * (https://dtcenter.org/met/users/docs/write_ups/circle_fit.pdf)
         */
        const double detr = suu * svv - suv * suv;
        if (fabs(detr) < FLT_EPSILON) return false;
        const double b0 = 0.5 * (suuu + suvv);
        const double b1 = 0.5 * (svvv + svuu);
        const double uc = (b0 * svv - b1 * suv) / detr;
        const double vc = (b1 * suu - b0 * suv) / detr;

        x0 += uc;
        y0 += vc;
        const double R0 = sqrt(uc * uc + vc * vc + (suu + svv) * n_inv);

        /* Compute a 1st estimate of the longitudinal slope using the median
         * of discrete derivatives
         */
        std::vector<double> tmp;
        tmp.reserve(n - 1);
        auto & h0 = this->cdcHits[0];
        ToLocal(h0);
        double p0 = atan2(v - y0, u - x0), z0 = w;
        for (size_t i = 1; i < n; i++) {
                auto & h1 = this->cdcHits[i];
                ToLocal(h1);
                const double p1 = atan2(v - y0, u - x0);
                const double dp = p1 - p0;
                const double dz = w - z0;
                p0 = p1, z0 = w;
                if (fabs(dz) > FLT_EPSILON)
                        tmp.push_back(dp / dz);
        }
        const auto ntmp = tmp.size();
        std::sort(tmp.begin(), tmp.end());
        double w0 = (ntmp % 2 == 0) ?
            0.5 * (tmp[ntmp / 2 - 1] + tmp[ntmp / 2]) : tmp[ntmp / 2];

        /* Compute the sum factors for the longitudinal fit */
        ToLocal(h0);
        p0 = atan2(v - y0, u - x0), z0 = w;
        double offset = 0;
        double sp = 0, sw = 0, spw = 0, sww = 0;
        for (auto& i : this->cdcHits) {
                ToLocal(i);
                double p = atan2(v - y0, u - x0) + offset;
                const double p1 = p0 + w0 * (w - z0);

                while (fabs(p1 - p) > 1.5 * M_PI) {
                        const double sgn = (p < p1) ? 1 : -1;
                        offset += sgn * 2 * M_PI;
                        p += sgn * 2 * M_PI;
                }
                z0 = w;
                p0 = p;

                sp += p;
                sw += w;
                spw += p * w;
                sww += w * w;
        }

        /* Fit of the longitudinal phase */
        const double detz = sww * n - sw * sw;
        if (fabs(detz) < FLT_EPSILON) return false;
        w0 = (spw * n - sp * sw) / detz;
        const double phi0 = (sww * sp - spw * sw) / detz;

        /* Compute the momentum in the CDC using the magnetic rigidity */
        const double c0 = 299792458; /* m / s */
        const double mm = 1E-03;  /* in m */
        const double MeV = 1E+06; /* in eV */
        const double B = 1.; /* T, TODO: read from EMField */
        const double momentum0 = c0 * B * (R0 * mm) / MeV *
                                 sqrt(1 + 1 / (R0 * R0 * w0 * w0));

        /* Fill in the fit parameters */
        this->cdcFit = { x0, y0, R0, w0, phi0, momentum0 };

        if (checkTrack) {
                /* Check the radius intersect */
                const double delta = sqrt(x0 * x0 + y0 * y0) - R0;
                if ((delta < -gSTRadius) || (delta > gSTRadius)) {
                        return false;
                }

                /* Check the z intersect */
                double z0 = (M_PI - phi0 + atan2(y0, x0)) / w0;
                if (w0 > 0) {
                        const double zmin = this->cdcZMin * 1E+03;
                        const double dz = 2 * M_PI / w0;
                        if (z0 > zmin) {
                                while (z0 > zmin) z0 -= dz;
                        } else if (z0 + dz < zmin) {
                                while (z0 + dz < zmin) z0 += dz;
                        }
                } else {
                        const double zmax = this->cdcZMax * 1E+03;
                        const double dz = -2 * M_PI / w0;
                        if (z0 < zmax) {
                                while (z0 < zmax) z0 += dz;
                        } else if (z0 - dz > zmax) {
                                while (z0 - dz > zmax) z0 -= dz;
                        }
                }
                if ((z0 < -450) || (z0 > 450)) {
                        return false;
                }

                /* Check the number of turns */
                const double r1 = 530;
                const double r0 = sqrt(x0 * x0 + y0 * y0);
                const double a = (r1 * r1 - r0 * r0 - R0 * R0) / (2 * R0 * r0);
                const double Dz1 = 2 * acos(a) / fabs(w0);
                if ((this->cdcZMax - this->cdcZMin) * 1E+03 < Dz1) {
                        return false;
                } else {
                        return true;
                }
        } else {
                return true;
        }
}


bool ITriggerCandidate::DumpJson(const std::string & path, const char * filename,
    COMET::ICOMETEvent & event)
{
        auto fid = std::fopen(path.c_str(), "a");
        if (!fid) {
                COMETNamedError(gTag,
                    "Error: could not open " << path);
                return false;
        }

        auto DumpHits = [&fid] (const Hits & hits) {
                for (auto& hit : hits) {
                        std::fprintf(fid, "[%.1f, %.1f, %.1f, %.5f],",
                            hit->GetPosX(), hit->GetPosY(), hit->GetPosZ(),
                            hit->GetPosT());
                }
        };

        auto field = COMET::IFieldManager::GetObject();

        auto DumpCDCHits = [&field, &fid] (const Hits & hits) {
                for (auto& hit : hits) {
                        auto h = static_cast<HitGas * const>(hit);
                        const double xyz[] = {h->GetPosX(), h->GetPosY(),
                                              h->GetPosZ()};
                        double B[6] = {0, 0, 0};
                        field->GetFieldValue(xyz, B);
                        std::fprintf(fid, "[%.3f, %.3f, %.3f, %.5f,"
                                          " %.5E, %.5E, %.5E,"
                                          " %.5E, %.5E, %.5E,"
                                          " %.5E],",
                            h->GetPosX(), h->GetPosY(), h->GetPosZ(),
                            h->GetPosT(), h->GetMomX(), h->GetMomY(),
                            h->GetMomZ(), B[0], B[1], B[2],
                            h->GetEnergyDeposit());
                }
        };

        auto && context = event.GetContext();
        auto trajectory = this->cthHits.front()->GetTrajectory();
        std::fprintf(fid, "{\"file\": \"%s\", \"run\": %d, \"subrun\": %d, "
            "\"event\": %d, \"tid\": %d, \"pid\": %d, \"momentum\": %.5E, "
            "\"hits\": {\"CTH\": [",
            filename, context.GetRun(), context.GetSubRun(),
            context.GetEvent(), this->trackId, this->pdgId,
            trajectory->GetInitialMomentum().P());

        DumpHits(this->cthHits);
        std::fputs("], \"CDC\": [", fid);
        DumpCDCHits(this->cdcHits);

        auto && vertex = trajectory->GetInitialPosition();
        auto & geoMgr = COMET::IOADatabase::Get().GeomId();
        auto && vertex_path = geoMgr.FullNodePath(
            vertex.X(), vertex.Y(), vertex.Z());
        std::fprintf(fid, "]}, \"track\": {\"vertex\": \"%s\", \"points\": [",
            vertex_path.c_str());
        auto points = trajectory->GetTrajectoryPoints();
        for (auto& point : points) {
                auto && r = point.GetPosition();
                std::fprintf(fid, "[%.1f, %.1f, %.1f],", r.X(), r.Y(), r.Z());
        }

        if (this->cdcFit.R0 > 0.) {
                std::fprintf(fid,
                    "]}, \"fit\": [%.3f, %.3f, %.3f, %.9E, %.9E, %.9E], ",
                    this->cdcFit.x0, this->cdcFit.y0, this->cdcFit.R0,
                    this->cdcFit.w0, this->cdcFit.phi0, this->cdcFit.momentum);
        } else {
                std::fputs("]}, ", fid);
        }

        auto tracks = event.Get<Tracks>("truth/G4Trajectories");
        auto && track = (*tracks)[1];
        auto && r = track.GetInitialPosition();
        auto && p = track.GetInitialMomentum();

        std::fprintf(fid, "\"primary\": {"
            "\"pid\": %d, \"position:\": [%.1f, %.1f, %.1f], "
            "\"momentum\": [%.5E, %.5E, %.5E]}}\n",
            track.GetPDGEncoding(), r.X(), r.Y(), r.Z(), p.X(), p.Y(),
            p.Z());

        std::fclose(fid);
        return true;
}


bool ITriggerCandidate::DumpPrimary(
    const std::string & path, COMET::ICOMETEvent & event)
{
        auto fid = std::fopen(path.c_str(), "a");
        if (!fid) {
                COMETNamedError(gTag,
                    "Error: could not open " << path);
                return false;
        }

        enum class PrintMode { CODE, CHARGE };
        auto PrintParticle = [&fid] (
            COMET::IG4Trajectory & track, PrintMode mode) -> double {
                const double GeV = 1E-03; /* MeV to GeV */
                const double mm = 1E-03; /* mm to m */

                auto && r = track.GetInitialPosition();
                auto && p = track.GetInitialMomentum();
                auto && particle = *track.GetParticle();
                const auto && mass = particle.Mass();
                const double kinetic = p.E() * GeV - mass;
                const double u_norm = 1. / sqrt(
                    p.X() * p.X() + p.Y() * p.Y() + p.Z() * p.Z());
                const double u[] = { p.X() * u_norm, p.Y() * u_norm,
                                     p.Z() * u_norm };

                int codeOrCharge = (mode == PrintMode::CHARGE) ?
                    (int)(particle.Charge() / 3) :
                    particle.PdgCode();

                if (mode == PrintMode::CODE) {
                        std::fprintf(fid, "%8d ", track.GetTrackId());
                }

                std::fprintf(fid,
                    "%3d %.3E %13.6f %13.6f %13.6f %14.7E %14.7E %14.7E",
                    codeOrCharge, kinetic, r.X() * mm,
                    r.Y() * mm, r.Z() * mm, u[0], u[1], u[2]);

                return r.T();
        };

        /* Primary info */
        auto tracks = event.Get<Tracks>("truth/G4Trajectories");
        auto && primary = (*tracks)[1];
        const double t0 = PrintParticle(primary, PrintMode::CHARGE);

        /* Run info */
        auto && context = event.GetContext();
        std::fprintf(fid, "   %7d %5d %7d",
            context.GetRun(), context.GetSubRun(), context.GetEvent());

        /* CRV info */
        const double ns = 1E-09; /* ns to s */
        std::fprintf(fid, "   %5d %14.7E %14.7E",
            this->crvTracks, (this->crvT0 - t0) * ns, (this->crvT1 - t0) * ns);

        /* Candidate info */
        auto trackPtr = &(*tracks)[this->trackId];

        std::fputs("   ", fid);
        PrintParticle(*trackPtr, PrintMode::CODE);

        int generation, parent;
        if (this->trackId == 1) {
                generation = 0;
        } else {
                for (generation = 1; (parent = trackPtr->GetParentId()) != 1;
                    generation++, trackPtr = &(*tracks)[parent]) {};
                if (generation == 1) trackPtr = &primary;
        }

        const double GeV = 1E-03; /* MeV to GeV */
        const double mm = 1E-03; /* mm to m */
        std::fprintf(fid, "   %14.7E %14.7E %14.7E   "
                          "%14.7E %14.7E %14.7E %14.7E   %2d",
            this->cdcMomentumMin * GeV, this->cdcMomentumMax * GeV,
            this->cdcMomentumMean * GeV, (this->cthT0 - t0) * ns,
            (this->cthT1 - t0) * ns, (this->cdcT0 - t0) * ns,
            (this->cdcT1 - t0) * ns, generation);
        std::fprintf(fid, "   %2d %2d %8.4f %8.4f",
            this->cdcLayerMin, this->cdcLayerMax, this->cdcZMin, this->cdcZMax);
        std::fprintf(fid, "   %8.4f %8.4f %8.4f %.9E %.9E %.9E ",
            this->cdcFit.x0 * mm, this->cdcFit.y0 * mm, this->cdcFit.R0 * mm,
            this->cdcFit.w0, this->cdcFit.phi0, this->cdcFit.momentum * GeV);
        PrintParticle(*trackPtr, PrintMode::CODE);
        std::fputs("\n", fid);

        std::fclose(fid);
        return true;
}