#include "WireTrackingManager.hh"

#include "ClassPW.h"

#include "G4Event.hh"
#include "G4RunManager.hh"
#include "G4Step.hh"
#include "G4StepPoint.hh"
#include "G4Track.hh"
#include "G4SystemOfUnits.hh"

#include <algorithm>
#include <cmath>

namespace {
constexpr double kAnodeRadius = 37.0 * mm;
constexpr double kCathodeRadius = 43.0 * mm;

bool CrossesRadius(double preR, double postR, double targetR)
{
  return (preR - targetR) * (postR - targetR) <= 0.0 && std::abs(postR - preR) > 1e-12;
}

TVector3 ToTVector3(const G4ThreeVector& value)
{
  return TVector3(value.x(), value.y(), value.z());
}

bool SolveRadiusIntersection(const G4ThreeVector& origin,
                             const G4ThreeVector& direction,
                             double radius,
                             double& distance)
{
  const double a = direction.x() * direction.x() + direction.y() * direction.y();
  if (a <= 0.0) {
    return false;
  }

  const double b = 2.0 * (origin.x() * direction.x() + origin.y() * direction.y());
  const double c = origin.x() * origin.x() + origin.y() * origin.y() - radius * radius;
  const double discriminant = b * b - 4.0 * a * c;
  if (discriminant < 0.0) {
    return false;
  }

  const double sqrtDiscriminant = std::sqrt(discriminant);
  const double t1 = (-b - sqrtDiscriminant) / (2.0 * a);
  const double t2 = (-b + sqrtDiscriminant) / (2.0 * a);

  distance = -1.0;
  if (t1 >= 0.0 && t2 >= 0.0) {
    distance = std::min(t1, t2);
  } else if (t1 >= 0.0) {
    distance = t1;
  } else if (t2 >= 0.0) {
    distance = t2;
  }

  return distance >= 0.0;
}
}

WireTrackingManager& WireTrackingManager::Instance()
{
  static WireTrackingManager instance;
  return instance;
}

WireTrackingManager::WireTrackingManager()
: fCurrentEventId(-1)
{}

WireTrackingManager::~WireTrackingManager() {}

void WireTrackingManager::EnsureEvent(int eventId)
{
  if (eventId != fCurrentEventId) {
    fCurrentEventId = eventId;
    fTrackInfo.clear();
  }
}

void WireTrackingManager::UpdateForStep(const G4Step* step)
{
  if (step == nullptr) {
    return;
  }

  const G4Track* track = step->GetTrack();
  if (track == nullptr) {
    return;
  }

  const G4Event* event = G4RunManager::GetRunManager()->GetCurrentEvent();
  if (event == nullptr) {
    return;
  }
  EnsureEvent(event->GetEventID());

  const G4StepPoint* prePoint = step->GetPreStepPoint();
  const G4StepPoint* postPoint = step->GetPostStepPoint();
  if (prePoint == nullptr || postPoint == nullptr) {
    return;
  }

  if (prePoint->GetPhysicalVolume() == nullptr || prePoint->GetPhysicalVolume()->GetName() != "Target") {
    return;
  }

  if (track->GetDefinition()->GetParticleType() == "nucleus") {
    return;
  }

  const G4ThreeVector& prePos = prePoint->GetPosition();
  const G4ThreeVector& postPos = postPoint->GetPosition();
  const double preR = prePos.perp();
  const double postR = postPos.perp();

  if (!CrossesRadius(preR, postR, kAnodeRadius) && !CrossesRadius(preR, postR, kCathodeRadius)) {
    return;
  }

  PW pw;
  pw.ConstructGeo();
  pw.FindWireID(ToTVector3(track->GetVertexPosition()), ToTVector3(track->GetMomentumDirection()), false);

  auto& info = fTrackInfo[track->GetTrackID()];

  auto interpolateEnergy = [&](double targetR) {
    const double fraction = (targetR - preR) / (postR - preR);
    return prePoint->GetKineticEnergy() + fraction * (postPoint->GetKineticEnergy() - prePoint->GetKineticEnergy());
  };

  if (!info.sawAnode && CrossesRadius(preR, postR, kAnodeRadius)) {
    info.sawAnode = true;
    info.anodeId = pw.GetNearestID().first;
    info.anodeEnergy = interpolateEnergy(kAnodeRadius);
  }

  if (!info.sawCathode && CrossesRadius(preR, postR, kCathodeRadius)) {
    info.sawCathode = true;
    info.cathodeId = pw.GetNearestID().second;
    info.cathodeEnergy = interpolateEnergy(kCathodeRadius);
  }

  if (info.sawAnode && info.sawCathode) {
    info.deltaE = std::abs(info.anodeEnergy - info.cathodeEnergy);
  }
}

bool WireTrackingManager::GetInfo(int eventId, int trackId, WireCrossingInfo& info)
{
  EnsureEvent(eventId);
  auto iterator = fTrackInfo.find(trackId);
  if (iterator == fTrackInfo.end()) {
    return false;
  }
  info = iterator->second;
  return true;
}

WireCrossingInfo WireTrackingManager::InferFromTrackGeometry(const G4Track* track, const G4StepPoint* hitPoint)
{
  WireCrossingInfo info;
  if (track == nullptr || hitPoint == nullptr) {
    return info;
  }

  const G4ThreeVector origin = track->GetVertexPosition();
  G4ThreeVector direction = hitPoint->GetPosition() - origin;
  const double totalPath = direction.mag();
  if (totalPath <= 0.0) {
    return info;
  }
  direction = direction.unit();

  PW pw;
  pw.ConstructGeo();
  pw.FindWireID(ToTVector3(origin), ToTVector3(direction), false);
  info.anodeId = pw.GetNearestID().first;
  info.cathodeId = pw.GetNearestID().second;

  const double vertexEnergy = track->GetVertexKineticEnergy();
  const double hitEnergy = hitPoint->GetKineticEnergy();
  double anodeDistance = 0.0;
  double cathodeDistance = 0.0;

  if (SolveRadiusIntersection(origin, direction, 37.0 * mm, anodeDistance)) {
    const double fraction = std::clamp(anodeDistance / totalPath, 0.0, 1.0);
    info.anodeEnergy = vertexEnergy + fraction * (hitEnergy - vertexEnergy);
    info.sawAnode = true;
  }

  if (SolveRadiusIntersection(origin, direction, 43.0 * mm, cathodeDistance)) {
    const double fraction = std::clamp(cathodeDistance / totalPath, 0.0, 1.0);
    info.cathodeEnergy = vertexEnergy + fraction * (hitEnergy - vertexEnergy);
    info.sawCathode = true;
  }

  if (info.sawAnode && info.sawCathode) {
    info.deltaE = std::abs(info.anodeEnergy - info.cathodeEnergy);
  }

  return info;
}