#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"

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 la direzione della quantità di moto (es. lungo l'asse Z)
  fParticleGun->SetParticleMomentumDirection(G4ThreeVector(0., 0., -1.));

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

  // Impostiamo la posizione di partenza (origine)
  fParticleGun->SetParticlePosition(G4ThreeVector(0., 0., 10. * m));
}

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

void EmitterPrimary::GeneratePrimaries(G4Event *anEvent) {
  // Questo metodo viene invocato all'inizio di ogni evento.
  // Qui potresti anche aggiungere una randomizzazione della posizione o
  // dell'energia.

  fParticleGun->GeneratePrimaryVertex(anEvent);
}

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

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 &points = m_Mesh->Points();
  const auto &quads = m_Mesh->Quads();

  for (const auto &q : quads) {
    uLib::Vector3f v0 = points[q(0)];
    uLib::Vector3f v1 = points[q(1)];
    uLib::Vector3f v2 = points[q(2)];
    uLib::Vector3f v3 = points[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];
  const auto &points = m_Mesh->Points();
  uLib::Vector3f v0 = points[q(0)];
  uLib::Vector3f v1 = points[q(1)];
  uLib::Vector3f v2 = points[q(2)];
  uLib::Vector3f v3 = points[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);
  }
}

} // namespace Geant
} // namespace uLib