#include "EmitterPrimary.hh"

#include "G4Box.hh"
#include "G4LogicalVolume.hh"
#include "G4LogicalVolumeStore.hh"
#include "G4ParticleDefinition.hh"
#include "G4ParticleGun.hh"
#include "G4ParticleTable.hh"
#include "G4RunManager.hh"
#include "G4SystemOfUnits.hh"
#include "Randomize.hh"

#include "EcoMug.hh"
#include "Math/Cylinder.h"


namespace uLib {
namespace Geant {

EmitterPrimary::EmitterPrimary()
    : G4VUserPrimaryGeneratorAction(), fParticleGun(nullptr) {
  // Creiamo il ParticleGun impostandolo per sparare 1 particella alla volta
  G4int n_particle = 1;
  fParticleGun = new G4ParticleGun(n_particle);

  // Otteniamo la tabella delle particelle di Geant4
  G4ParticleTable *particleTable = G4ParticleTable::GetParticleTable();

  // Cerchiamo il muone negativo (usa "mu+" per l'antimuone)
  G4String particleName = "mu-";
  G4ParticleDefinition *particle = particleTable->FindParticle(particleName);

  // Configuriamo le proprietà iniziali della particella
  fParticleGun->SetParticleDefinition(particle);

  // Impostiamo l'energia cinetica a 1 GeV
  fParticleGun->SetParticleEnergy(1.0 * GeV);

  // Initial position and direction through AffineTransform
  // 10m on Z axis, pointing towards origin
  this->SetPosition(Vector3f(0, 0, 10000.0));
  // Default orientation is identity (pointing along -Z if we rotate the puppet accordingly)
  // But fParticleGun defaults are set here and overridden in GeneratePrimaries
}

EmitterPrimary::~EmitterPrimary() {
  // Importante: liberare la memoria
  delete fParticleGun;
}

void EmitterPrimary::GeneratePrimaries(G4Event *anEvent) {
  // Use position and direction from AffineTransform
  Vector3f pos = this->GetPosition();
  // Assume default direction is along the -Z axis of the local frame
  Vector4f dir4 = this->GetWorldMatrix() * Vector4f(0, 0, -1, 0);
  Vector3f dir = dir4.head<3>().normalized();

  fParticleGun->SetParticlePosition(G4ThreeVector(pos(0), pos(1), pos(2)));
  fParticleGun->SetParticleMomentumDirection(G4ThreeVector(dir(0), dir(1), dir(2)));

  fParticleGun->GeneratePrimaryVertex(anEvent);
}

EmitterPrimary* EmitterPrimary::Clone() const {
  auto* clone = new EmitterPrimary();
  clone->SetMatrix(this->GetMatrix());
  return clone;
}

// -------------------------------------------------------------------------- //
// SkyPlaneEmitterPrimary using EcoMug

SkyPlaneEmitterPrimary::SkyPlaneEmitterPrimary()
    : EmitterPrimary(), m_EcoMug(new EcoMug()), m_Size(1000.0, 1000.0) {
    // Initial configuration for EcoMug in sky mode
    m_EcoMug->SetUseSky();
    m_EcoMug->SetSkySize({m_Size(0)/1000.0, m_Size(1)/1000.0});
}

SkyPlaneEmitterPrimary::~SkyPlaneEmitterPrimary() {
    delete m_EcoMug;
}

void SkyPlaneEmitterPrimary::SetSkySize(const uLib::Vector2f& size) {
    m_Size = size;
    // EcoMug units are in meters (m=1), Geant4 units are in mm.
    m_EcoMug->SetSkySize({m_Size(0)/1000.0, m_Size(1)/1000.0});
    this->Updated();
}

void SkyPlaneEmitterPrimary::SetPlane(const uLib::Vector3f& p0, const uLib::Vector3f& normal) {
    this->SetPosition(p0);
    // Orient the emitter so that local Z is the normal.
    // This is useful for unconventional planes, though EcoMug sky is usually horizontal.
    // If we want a truly 'sky' plane, it usually stays horizontal.
    this->Updated();
}

void SkyPlaneEmitterPrimary::GeneratePrimaries(G4Event* anEvent) {
    if (!m_EcoMug) return;
    m_EcoMug->Generate();

    // EcoMug position is relative to its internal sky center in meters.
    // Our wrapper uses the AffineTransform for the overall positioning.
    std::array<double, 3> pos_m = m_EcoMug->GetGenerationPosition();
    G4ThreeVector local_pos(pos_m[0]*1000.0, pos_m[1]*1000.0, pos_m[2]*1000.0);

    // EcoMug momentum (direction and magnitude in GeV/c)
    double p_mag = m_EcoMug->GetGenerationMomentum();
    double theta = m_EcoMug->GetGenerationTheta();
    double phi = m_EcoMug->GetGenerationPhi();
    
    // EcoMug theta is generated in a way that PI means pointing down (-Z)
    G4ThreeVector local_dir(
        sin(theta) * cos(phi),
        sin(theta) * sin(phi),
        cos(theta)
    );

    // Transform local coordinates to world
    Matrix4f world_mat = this->GetWorldMatrix();
    Vector3f world_pos = (world_mat * Vector4f(local_pos.x(), local_pos.y(), local_pos.z(), 1.0)).head<3>();
    Vector3f world_dir = (world_mat * Vector4f(local_dir.x(), local_dir.y(), local_dir.z(), 0.0)).head<3>().normalized();

    // Set particle charge
    G4ParticleTable *particleTable = G4ParticleTable::GetParticleTable();
    G4String particleName = (m_EcoMug->GetCharge() > 0) ? "mu+" : "mu-";
    fParticleGun->SetParticleDefinition(particleTable->FindParticle(particleName));

    fParticleGun->SetParticlePosition(G4ThreeVector(world_pos(0), world_pos(1), world_pos(2)));
    fParticleGun->SetParticleMomentumDirection(G4ThreeVector(world_dir(0), world_dir(1), world_dir(2)));
    fParticleGun->SetParticleEnergy(p_mag * GeV);
    
    fParticleGun->GeneratePrimaryVertex(anEvent);
}

EmitterPrimary* SkyPlaneEmitterPrimary::Clone() const {
    auto* clone = new SkyPlaneEmitterPrimary();
    clone->SetSkySize(this->m_Size);
    clone->SetMatrix(this->GetMatrix());
    return clone;
}
 
// -------------------------------------------------------------------------- //
// CylinderEmitterPrimary using EcoMug

CylinderEmitterPrimary::CylinderEmitterPrimary()
    : EmitterPrimary(), m_EcoMug(new EcoMug()), m_Radius(1000.0), m_Height(1000.0) {
    m_EcoMug->SetUseCylinder();
    m_EcoMug->SetCylinderRadius(m_Radius/1000.0);
    m_EcoMug->SetCylinderHeight(m_Height/1000.0);
    m_EcoMug->SetCylinderCenterPosition({0.0, 0.0, m_Height/2000.0});
}

CylinderEmitterPrimary::~CylinderEmitterPrimary() {
    delete m_EcoMug;
}

void CylinderEmitterPrimary::SetRadius(float r) {
    m_Radius = r;
    m_EcoMug->SetCylinderRadius(m_Radius/1000.0);
    this->Updated();
}

void CylinderEmitterPrimary::SetHeight(float h) {
    m_Height = h;
    m_EcoMug->SetCylinderHeight(m_Height/1000.0);
    m_EcoMug->SetCylinderCenterPosition({0.0, 0.0, m_Height/2000.0});
    this->Updated();
}

void CylinderEmitterPrimary::GeneratePrimaries(G4Event* anEvent) {
    if (!m_EcoMug) return;
    m_EcoMug->Generate();

    std::array<double, 3> pos_m = m_EcoMug->GetGenerationPosition();
    G4ThreeVector local_pos(pos_m[0]*1000.0, pos_m[1]*1000.0, pos_m[2]*1000.0);

    double p_mag = m_EcoMug->GetGenerationMomentum();
    double theta = m_EcoMug->GetGenerationTheta();
    double phi = m_EcoMug->GetGenerationPhi();
    
    G4ThreeVector local_dir(
        sin(theta) * cos(phi),
        sin(theta) * sin(phi),
        cos(theta)
    );

    Matrix4f world_mat = this->GetWorldMatrix();
    Vector3f world_pos = (world_mat * Vector4f(local_pos.x(), local_pos.y(), local_pos.z(), 1.0)).head<3>();
    Vector3f world_dir = (world_mat * Vector4f(local_dir.x(), local_dir.y(), local_dir.z(), 0.0)).head<3>().normalized();

    G4ParticleTable *particleTable = G4ParticleTable::GetParticleTable();
    G4String particleName = (m_EcoMug->GetCharge() > 0) ? "mu+" : "mu-";
    fParticleGun->SetParticleDefinition(particleTable->FindParticle(particleName));

    fParticleGun->SetParticlePosition(G4ThreeVector(world_pos(0), world_pos(1), world_pos(2)));
    fParticleGun->SetParticleMomentumDirection(G4ThreeVector(world_dir(0), world_dir(1), world_dir(2)));
    fParticleGun->SetParticleEnergy(p_mag * GeV);
    
    fParticleGun->GeneratePrimaryVertex(anEvent);
}

EmitterPrimary* CylinderEmitterPrimary::Clone() const {
    auto* clone = new CylinderEmitterPrimary();
    clone->SetRadius(this->m_Radius);
    clone->SetHeight(this->m_Height);
    clone->SetMatrix(this->GetMatrix());
    return clone;
}

// -------------------------------------------------------------------------- //












// -------------------------------------------------------------------------- //


QuadMeshEmitterPrimary::QuadMeshEmitterPrimary()
    : EmitterPrimary(), m_Mesh(nullptr), m_TotalArea(0.0) {
}

QuadMeshEmitterPrimary::~QuadMeshEmitterPrimary() {
}

void QuadMeshEmitterPrimary::SetMesh(uLib::QuadMesh *mesh) {
  m_Mesh = mesh;
  CalculateAreas();
}

void QuadMeshEmitterPrimary::CalculateAreas() {
  if (!m_Mesh) return;
  m_CumulativeAreas.clear();
  m_TotalArea = 0.0;

  const auto &quads = m_Mesh->Quads();

  for (const auto &q : quads) {
    uLib::Vector3f v0 = m_Mesh->GetPoint(q(0));
    uLib::Vector3f v1 = m_Mesh->GetPoint(q(1));
    uLib::Vector3f v2 = m_Mesh->GetPoint(q(2));
    uLib::Vector3f v3 = m_Mesh->GetPoint(q(3));

    double a1 = 0.5 * (v1 - v0).cross(v2 - v0).norm();
    double a2 = 0.5 * (v2 - v0).cross(v3 - v0).norm();
    m_TotalArea += (a1 + a2);
    m_CumulativeAreas.push_back(m_TotalArea);
  }
}

void QuadMeshEmitterPrimary::GeneratePrimaries(G4Event *anEvent) {
  if (!m_Mesh || m_TotalArea <= 0.0) return;

  // 1. Choose a quad
  double r = G4UniformRand() * m_TotalArea;
  auto it = std::lower_bound(m_CumulativeAreas.begin(), m_CumulativeAreas.end(), r);
  int quadIdx = std::distance(m_CumulativeAreas.begin(), it);

  const auto &q = m_Mesh->Quads()[quadIdx];
  uLib::Vector3f v0 = m_Mesh->GetPoint(q(0));
  uLib::Vector3f v1 = m_Mesh->GetPoint(q(1));
  uLib::Vector3f v2 = m_Mesh->GetPoint(q(2));
  uLib::Vector3f v3 = m_Mesh->GetPoint(q(3));

  // 2. Choose a point on the quad
  double a1 = 0.5 * (v1 - v0).cross(v2 - v0).norm();
  double a2 = 0.5 * (v2 - v0).cross(v3 - v0).norm();
  
  G4ThreeVector pos;
  uLib::Vector3f normal = m_Mesh->GetNormal(quadIdx);
  
  if (G4UniformRand() < a1 / (a1 + a2)) {
    double u = G4UniformRand();
    double v = G4UniformRand();
    if (u + v > 1.0) { u = 1.0 - u; v = 1.0 - v; }
    uLib::Vector3f p = v0 + u * (v1 - v0) + v * (v2 - v0);
    pos.set(p(0), p(1), p(2));
  } else {
    double u = G4UniformRand();
    double v = G4UniformRand();
    if (u + v > 1.0) { u = 1.0 - u; v = 1.0 - v; }
    uLib::Vector3f p = v0 + u * (v2 - v0) + v * (v3 - v0);
    pos.set(p(0), p(1), p(2));
  }

  // 3. Choose a direction (Cosmic Muon: cos^2(theta))
  G4ThreeVector dir;
  bool accepted = false;
  int tries = 0;
  
  while (!accepted && tries < 1000) {
    tries++;
    double cosTheta = std::pow(G4UniformRand(), 1.0/3.0);
    double sinTheta = std::sqrt(1.0 - cosTheta * cosTheta);
    double phi = 2.0 * M_PI * G4UniformRand();
    
    // Incoming from above (+Z towards -Z)
    dir.set(sinTheta * std::cos(phi), sinTheta * std::sin(phi), -cosTheta);
    
    // Filtering: pointing on the same side of the face normal
    if (dir.x() * normal(0) + dir.y() * normal(1) + dir.z() * normal(2) > 0) {
      accepted = true;
    }
  }

  if (accepted) {
    fParticleGun->SetParticlePosition(pos);
    fParticleGun->SetParticleMomentumDirection(dir);
    // Keep energy from base class or set here if needed
    fParticleGun->GeneratePrimaryVertex(anEvent);
  }
}
EmitterPrimary* QuadMeshEmitterPrimary::Clone() const {
  auto* clone = new QuadMeshEmitterPrimary();
  if (m_Mesh) clone->SetMesh(m_Mesh);
  clone->SetMatrix(this->GetMatrix());
  return clone;
}

} // namespace Geant
} // namespace uLib