#include <TMath.h>

#include <RAT/DB.hh>
#include <RAT/Log.hh>
#include <RAT/QuadraticRoots.hh>
#include <RAT/DU/LightPathStraightFitter.hh>

using namespace RAT;
using namespace RAT::DU;

ClassImp(RAT::DU::LightPathStraightFitter)

Double_t LightPathStraightFitter::fNeckInnerRadius;
Double_t LightPathStraightFitter::fNeckOuterRadius;
Double_t LightPathStraightFitter::fAVInnerRadius;
Double_t LightPathStraightFitter::fAVOuterRadius;
Double_t LightPathStraightFitter::fPMTRadius;
Double_t LightPathStraightFitter::fFillZ;
size_t LightPathStraightFitter::fAVSystemId;

void
LightPathStraightFitter::BeginOfRun()
{
  fNeckInnerRadius = 0.0;
  fNeckOuterRadius = 0.0;
  fAVInnerRadius = 0.0;
  fAVOuterRadius = 0.0;
  fPMTRadius = 0.0;
  fFillZ = 0.0;

  // copy code from LightPathCalculator
  DB *db = DB::Get();

  // Obtain the inner/outer radii of the acrylic vessel from the RAT database
  DBLinkPtr link = db->GetLink( "NATIVE_GEO_DIMENSIONS", "natgeo_dimensions" );
  fAVInnerRadius = link->GetD( "inner_av_radius" );
  double AVthickness = link->GetD( "av_thickness" );
  fAVOuterRadius = fAVInnerRadius + AVthickness;

  // Obtain the inner/outer radii of the acrylic vessel neck from the RAT database
  fNeckInnerRadius = link->GetD( "neck_inner_radius" );
  double NeckThickness = link->GetD( "neck_thickness" );
  fNeckOuterRadius = fNeckInnerRadius + NeckThickness;

  // Obtain the radius of the PMT bucket from the RAT database
  fPMTRadius = db->GetLink( "GREY_DISC_PARAMETERS", "DiscOptics0_0" )->GetD( "disc_radius" );

  // set default z split at top of world
  link = db->GetLink( "GEO", "world" );
  std::vector<Double_t> hs = link->GetDArray("half_size");
  if (hs.size() >= 3) {
    fFillZ = hs[2];
  } else {
    fFillZ = 40000.0; // this is what it was in all geo files in Aug 2019
    warn << "LightPathStraightFitter::BeginOfRun: No world volume found! "
         << "Default value " << fFillZ << " used." << newline;
  }

  // If the geometry is using a partial fill configuration, obtain the fill fraction (fFillZ) as well as the refractive indices
  // for the upper/lower targets
  link = db->GetLink( "GEO", "inner_av" );
  try {
    fFillZ = link->GetD( "split_z" );
    info << "LightPathStraightFitter::BeginOfRun: Non-homogeneous detector medium (partial fill) assumed." << newline;
  } catch( DBNotFoundError& e ) {
    info << "LightPathStraightFitter::BeginOfRun: Homogeneous detector medium (not partial fill) assumed." << newline;
  }

  // find av shift
  fAVSystemId = Point3D::GetSystemId("av");
}

LightPath&
LightPathStraightFitter::Fit(LightPath& startPath) const
{
  Point3D start(fAVSystemId);
  Point3D stop(fAVSystemId);

  // preserve information from starting path
  start = startPath.GetPoint(0); // translates into av coordinate system
  stop = startPath.GetEndPos();
  LightPath::RegionType stopRegionType = startPath.GetRegionType(startPath.GetPointCount() - 1);

  startPath.ClearPoints(0); // clear all the points except the starting position
  TVector3 dir = stop.GetDirectionFrom(start); // length of vector is distance from start to stop
  startPath.AddPoint(start, dir, GetRegionType(start));

  std::vector<Double_t> ts; // progress markers between start and stop
  ts.push_back(0.0);

  // find intersections with AV
  Double_t aa = dir.Mag2();
  Double_t ba = 2.0 * (dir.X()*start.X() + dir.Y()*start.Y() + dir.Z()*start.Z());
  Double_t cai = start.Mag2() - fAVInnerRadius * fAVInnerRadius;
  Double_t cao = start.Mag2() - fAVOuterRadius * fAVOuterRadius;

  Double_t t1, t2;
  if (QuadraticRoots(aa, ba, cao, t1, t2)) {
    if (t1 != t2) { // ignore glancing touch
      if (t1 > 0.0 && t1 < 1.0) {
        Insert(startPath, ts, start.GetNewPoint(t1 * dir), t1, LightPath::AV);
      }
      if (t2 > 0.0 && t2 < 1.0) {
        Insert(startPath, ts, start.GetNewPoint(t2 * dir), t2, LightPath::AV);
      }
    }
  }

  if (QuadraticRoots(aa, ba, cai, t1, t2)) {
    if (t1 != t2) {
      if (t1 > 0.0 && t1 < 1.0) {
        Insert(startPath, ts, start.GetNewPoint(t1 * dir), t1, LightPath::AV);
      }
      if (t2 > 0.0 && t2 < 1.0) {
        Insert(startPath, ts, start.GetNewPoint(t2 * dir), t2, LightPath::AV);
      }
    }
  }

  // intersections with neck
  Double_t an = dir.X()*dir.X() + dir.Y()*dir.Y();
  Double_t bn = 2.0 * (dir.X()*start.X() + dir.Y()*start.Y());
  Double_t cni = start.X()*start.X() + start.Y()*start.Y() - fNeckInnerRadius * fNeckInnerRadius;
  Double_t cno = start.X()*start.X() + start.Y()*start.Y() - fNeckOuterRadius * fNeckOuterRadius;

  if (QuadraticRoots(an, bn, cno, t1, t2)) {
    if (t1 != t2) {
      if (t1 > 0.0 && t1 < 1.0) {
        Point3D np = start.GetNewPoint(t1 * dir);
        if (np.Z() > 0.0 && np.Mag2() > fAVInnerRadius * fAVInnerRadius) {
          Insert(startPath, ts, start.GetNewPoint(t1 * dir), t1, LightPath::Neck);
        }
      }
      if (t2 > 0.0 && t2 < 1.0) {
        Point3D np = start.GetNewPoint(t2 * dir);
        if (np.Z() > 0.0 && np.Mag2() > fAVInnerRadius * fAVInnerRadius) {
          Insert(startPath, ts, start.GetNewPoint(t2 * dir), t2, LightPath::Neck);
        }
      }
    }
  }

  if (QuadraticRoots(an, bn, cni, t1, t2)) {
    if (t1 != t2) {
      if (t1 > 0.0 && t1 < 1.0) {
        Point3D np = start.GetNewPoint(t1 * dir);
        if (np.Z() > 0.0 && np.Mag2() > fAVInnerRadius * fAVInnerRadius) {
          Insert(startPath, ts, start.GetNewPoint(t1 * dir), t1, LightPath::Neck);
        }
      }
      if (t2 > 0.0 && t2 < 1.0) {
        Point3D np = start.GetNewPoint(t2 * dir);
        if (np.Z() > 0.0 && np.Mag2() > fAVInnerRadius * fAVInnerRadius) {
          Insert(startPath, ts, start.GetNewPoint(t2 * dir), t2, LightPath::Neck);
        }
      }
    }
  }

  // intersection with upper/lower split
  if ((start.Z() < fFillZ && stop.Z() > fFillZ) ||
      (start.Z() > fFillZ && stop.Z() < fFillZ)) {
    t1 = (fFillZ - start.Z()) / (stop.Z() - start.Z());
    Insert(startPath, ts, start.GetNewPoint(t1 * dir), t1, LightPath::UpperInnerAV);
  }

  // add stop point
  startPath.AddPoint(stop, dir, stopRegionType);

  // now change the region types to reflect the actual transitions.
  for (UInt_t i = 1; i < startPath.GetPointCount() - 1; ++i) {
    Point3D p = startPath.GetPoint(i);
    p.Accumulate(startPath.GetPoint(i+1));
    p.Divide(2);
    startPath.SetRegionType(i, GetRegionType(p));
  }

  // if two regions have the same type (e.g., from inner neck volume to inner AV),
  // remove the second point
  for (UInt_t i = startPath.GetPointCount() - 1; i > 0; --i) {
    if (startPath.GetRegionType(i) == startPath.GetRegionType(i-1)) {
      startPath.DeletePoint(i);
    }
  }

  startPath.SetFitter(Name());
  return startPath;
}

void
LightPathStraightFitter::Insert(LightPath& path,
                                std::vector<Double_t>& ts,
                                const Point3D& p,
                                Double_t t,
                                LightPath::RegionType rtype)
{
  // find place to insert point (first has t==0.0 by construction)
  for (UInt_t i = 1; i < ts.size(); ++i) {
    if (t < ts[i]) {
      path.InsertPoint(i, p, rtype);
      ts.insert(ts.begin() + i, t);
      return;
    }
  }
  // add to end
  path.AddPoint(p, path.GetInitialLightVec(), rtype);
  ts.push_back(t);
}

LightPath::RegionType
LightPathStraightFitter::GetRegionType(const Point3D& point)
{
  Point3D p(fAVSystemId);
  p = point;
  Double_t r2 = p.Mag2();
  if (r2 < fAVInnerRadius * fAVInnerRadius) {
    if (fFillZ > fAVInnerRadius) return LightPath::InnerAV;
    else if (p.Z() > fFillZ) return LightPath::UpperInnerAV;
    else return LightPath::LowerInnerAV;
  } else if (p.Z() > 0.0) {
    Double_t rho2 = p.X() * p.X() + p.Y() * p.Y();
    if (rho2 < fNeckInnerRadius * fNeckInnerRadius) {
      // same material as inner AV or upper inner AV
      if (p.Z() > fFillZ) return LightPath::UpperInnerAV;
      else return LightPath::InnerAV;
    } else if (rho2 < fNeckOuterRadius * fNeckOuterRadius) {
      return LightPath::Neck;
    }
  }
  if (r2 <= fAVOuterRadius * fAVOuterRadius) {
    return LightPath::AV;
  } else {
    return LightPath::Water;
  }
}

Double_t
LightPathStraightFitter::GetSolidAngle(const LightPath& lightPath,
                                       const TVector3& pmtNorm,
                                       UInt_t nVal,
                                       Double_t& cosThetaAvg) const
{
  UInt_t nv = nVal;
  if (nv == 0) nv = 4;
  if (nv < 3) {
    warn << "LightPathStraightFitter::GetSolidAngle nVal=" << nv << " must be >= 31!" << newline;
    Log::Die("nVal must be >= 3");
  }

  TVector3 zPrime = pmtNorm.Unit();
  TVector3 yPrime = (zPrime.Orthogonal()).Unit();
  TVector3 xPrime = (yPrime.Cross(zPrime)).Unit();

  Double_t cthav = 0.0; // running cos(theta) average
  Double_t sangle = 0.0; // solid angle
  Double_t rotationStep = TMath::TwoPi() / nv;

  std::vector<TVector3> rays;
  for (UInt_t j = 0; j < nv; ++j) {
    Point3D tmp(lightPath.GetEndPos());
    Double_t angle = rotationStep * j;
    tmp.Move(fPMTRadius * (xPrime * TMath::Cos(angle) + yPrime * TMath::Sin(angle)));
    TVector3 tmpray = tmp.GetDirectionFrom(lightPath.GetPoint(0)).Unit();
    rays.push_back(tmpray);
    cthav -= tmpray.Dot(zPrime);
  }

  if (nVal == 0) {
    // Approximate 4-point calculation for nVal == 0
    Double_t angAlpha = rays[0].Angle(rays[2]) / 2.0;
    Double_t angBeta  = rays[1].Angle(rays[3]) / 2.0;
    sangle = TMath::TwoPi() * (1.0 - (TMath::Cos(angAlpha) + TMath::Cos(angBeta)) / 2.0);
    if (sangle <= 0.0 || sangle > TMath::TwoPi()) {
      warn << "LightPathStraightFitter::GetSolidAngle: solid angle="
           << sangle << " is out of range" << newline;
      sangle = -999.0;
    }
  } else {
    // more general polygon formula (including for nVal == 4)
    Double_t angleSum = 0.0;
    TVector3 evToMid = (lightPath.GetEndPos().GetDirectionFrom(lightPath.GetPoint(0))).Unit();
    for (UInt_t k = 0; k < nv; ++k) {
      TVector3 vecAMid = rays[k] - evToMid;
      TVector3 vecBMid = rays[(k+1) % nv] - evToMid;
      Double_t sinAngle = (vecBMid.Cross(vecAMid)).Mag();
      sangle += sinAngle;
      angleSum += sinAngle / (vecBMid.Mag() * vecAMid.Mag());
    }
    sangle *= 0.5 * TMath::TwoPi() / angleSum;
  }

  cosThetaAvg = cthav / nv;
  return sangle;
}