#include "GarfieldPhysics.hh"
#include "GarfieldAnalysis.hh"
#include "Garfield/AvalancheMC.hh"
#include "Garfield/AvalancheMicroscopic.hh"

GarfieldPhysics* GarfieldPhysics::fGarfieldPhysics = nullptr;

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo....

GarfieldPhysics* GarfieldPhysics::GetInstance() {
  if (!fGarfieldPhysics) fGarfieldPhysics = new GarfieldPhysics();
  return fGarfieldPhysics;
}

void GarfieldPhysics::Dispose() {
  if (fGarfieldPhysics) {
    delete fGarfieldPhysics;
    fGarfieldPhysics = nullptr;
    G4cout << "GarfieldPhysics: Singleton disposed" << G4endl;
  }
}

GarfieldPhysics::~GarfieldPhysics() {
  delete fTrackHeed;
  fTrackHeed = nullptr;
  delete fComponentAnalyticField;
  fComponentAnalyticField = nullptr;
  delete fSensor;
  fSensor = nullptr;
  delete fMediumMagboltz;
  fMediumMagboltz = nullptr;
  G4cout << "GarfieldPhysics: Instance destructed" << G4endl;
}

std::string GarfieldPhysics::GetIonizationModel() { return fIonizationModel; }

void GarfieldPhysics::SetIonizationModel(std::string model, bool useDefaults) {
  if (model != "PAIPhot" && model != "PAI" && model != "Heed") {
    std::cout << "Unknown ionization model " << model << std::endl;
    std::cout << "Using PAIPhot as default model!" << std::endl;
    model = "PAIPhot";
  }
  fIonizationModel = model;
  G4cout << "GarfieldPhysics: Setting ionization model to " << fIonizationModel << G4endl;

  if (fIonizationModel == "PAIPhot" || fIonizationModel == "PAI") {
    if (useDefaults) {
      this->AddParticleName("e-", 1e-6, 1e-3, "garfield");
      this->AddParticleName("gamma", 1e-6, 1e+8, "garfield");
      this->AddParticleName("e-", 0, 1e+8, "geant4");
      this->AddParticleName("e+", 0, 1e+8, "geant4");
      /*this->AddParticleName("mu-", 0, 1e+8, "geant4");
      this->AddParticleName("mu+", 0, 1e+8, "geant4");*/
      this->AddParticleName("proton", 0, 1e+8, "geant4");
      this->AddParticleName("pi+", 0, 1e+8, "geant4");
      this->AddParticleName("pi-", 0, 1e+8, "geant4");
      this->AddParticleName("alpha", 0, 1e+8, "geant4");
      this->AddParticleName("He3", 0, 1e+8, "geant4");
      this->AddParticleName("GenericIon", 0, 1e+8, "geant4");
      G4cout << "GarfieldPhysics: Set default particles for " << fIonizationModel << G4endl;
    }
  } else if (fIonizationModel == "Heed" && useDefaults) {
    this->AddParticleName("gamma", 1e-6, 1e+8, "garfield");
    this->AddParticleName("e-", 6e-2, 1e+7, "garfield");
    this->AddParticleName("e+", 6e-2, 1e+7, "garfield");
    this->AddParticleName("mu-", 0, 1e+8, "garfield");  
    this->AddParticleName("mu+", 0, 1e+8, "garfield");
    this->AddParticleName("pi-", 2e+1, 1e+8, "garfield");
    this->AddParticleName("pi+", 2e+1, 1e+8, "garfield");
    this->AddParticleName("kaon-", 1e+1, 1e+8, "garfield");
    this->AddParticleName("kaon+", 1e+1, 1e+8, "garfield");
    this->AddParticleName("proton", 9.e+1, 1e+8, "garfield");
    this->AddParticleName("anti_proton", 9.e+1, 1e+8, "garfield");
    this->AddParticleName("deuteron", 2.e+2, 1e+8, "garfield");
    this->AddParticleName("alpha", 4.e+2, 1e+8, "garfield");
    G4cout << "GarfieldPhysics: Set default particles for Heed" << G4endl;
  }
}

void GarfieldPhysics::AddParticleName(const std::string particleName,
                                      double ekin_min_MeV, double ekin_max_MeV,
                                      std::string program) {
  if (ekin_min_MeV >= ekin_max_MeV) {
    std::cout << "Ekin_min=" << ekin_min_MeV
              << " keV is larger than Ekin_max=" << ekin_max_MeV << " keV"
              << std::endl;
    return;
  }

  if (program == "garfield") {
    G4cout << "GarfieldPhysics: Adding " << particleName << " to Garfield++ range: " 
           << ekin_min_MeV << " MeV to " << ekin_max_MeV << " MeV" << G4endl;
    fMapParticlesEnergyGarfield.insert(std::make_pair(
        particleName, std::make_pair(ekin_min_MeV, ekin_max_MeV)));
  } else {
    G4cout << "GarfieldPhysics: Adding " << particleName << " to Geant4 range: " 
           << ekin_min_MeV << " MeV to " << ekin_max_MeV << " MeV" << G4endl;
    fMapParticlesEnergyGeant4.insert(std::make_pair(
        particleName, std::make_pair(ekin_min_MeV, ekin_max_MeV)));
  }
}

bool GarfieldPhysics::FindParticleName(std::string name, std::string program) {
  if (program == "garfield") {
    auto it = fMapParticlesEnergyGarfield.find(name);
    bool found = (it != fMapParticlesEnergyGarfield.end());
    G4cout << "GarfieldPhysics: FindParticleName(" << name << ", garfield) = " << (found ? "true" : "false") << G4endl;
    return found;
  } else {
    auto it = fMapParticlesEnergyGeant4.find(name);
    bool found = (it != fMapParticlesEnergyGeant4.end());
    G4cout << "GarfieldPhysics: FindParticleName(" << name << ", geant4) = " << (found ? "true" : "false") << G4endl;
    return found;
  }
}

bool GarfieldPhysics::FindParticleNameEnergy(std::string name, double ekin_MeV,
                                             std::string program) {
  if (program == "garfield") {
    auto it = fMapParticlesEnergyGarfield.find(name);
    if (it != fMapParticlesEnergyGarfield.end()) {
      EnergyRange_MeV range = it->second;
      bool inRange = (range.first <= ekin_MeV && range.second >= ekin_MeV);
      G4cout << "GarfieldPhysics: FindParticleNameEnergy(" << name << ", " << ekin_MeV 
             << " MeV, garfield) - Range: " << range.first << " to " << range.second 
             << " MeV, InRange = " << (inRange ? "true" : "false") << G4endl;
      return inRange;
    }
  } else {
    auto it = fMapParticlesEnergyGeant4.find(name);
    if (it != fMapParticlesEnergyGeant4.end()) {
      EnergyRange_MeV range = it->second;
      bool inRange = (range.first <= ekin_MeV && range.second >= ekin_MeV);
      G4cout << "GarfieldPhysics: FindParticleNameEnergy(" << name << ", " << ekin_MeV 
             << " MeV, geant4) - Range: " << range.first << " to " << range.second 
             << " MeV, InRange = " << (inRange ? "true" : "false") << G4endl;
      return inRange;
    }
  }
  G4cout << "GarfieldPhysics: FindParticleNameEnergy(" << name << ", " << ekin_MeV 
         << " MeV, " << program << ") - Not found" << G4endl;
  return false;
}

double GarfieldPhysics::GetMinEnergyMeVParticle(std::string name,
                                                std::string program) {
  if (program == "garfield") {
    auto it = fMapParticlesEnergyGarfield.find(name);
    if (it != fMapParticlesEnergyGarfield.end()) {
      return it->second.first;
    }
  } else {
    auto it = fMapParticlesEnergyGeant4.find(name);
    if (it != fMapParticlesEnergyGeant4.end()) {
      return it->second.first;
    }
  }
  return -1;
}

double GarfieldPhysics::GetMaxEnergyMeVParticle(std::string name,
                                                std::string program) {
  if (program == "garfield") {
    auto it = fMapParticlesEnergyGarfield.find(name);
    if (it != fMapParticlesEnergyGarfield.end()) {
      return it->second.second;
    }
  } else {
    auto it = fMapParticlesEnergyGeant4.find(name);
    if (it != fMapParticlesEnergyGeant4.end()) {
      return it->second.second;
    }
  }
  return -1;
}

void GarfieldPhysics::InitializePhysics() {
  G4cout << "GarfieldPhysics: Initializing physics" << G4endl;
  fMediumMagboltz = new Garfield::MediumMagboltz();
  fMediumMagboltz->SetComposition("ar", 90., "ch4", 10.);
  fMediumMagboltz->SetTemperature(295.15);
  fMediumMagboltz->SetPressure(814.);
  fMediumMagboltz->SetMaxElectronEnergy(100000.);
  fMediumMagboltz->Initialise(true);

  const double rPenning = 0.57;
  const double lambdaPenning = 0.;
  fMediumMagboltz->EnablePenningTransfer(rPenning, lambdaPenning, "ar");

  G4String gasFile = "p10_modified.gas";
  if (fMediumMagboltz->LoadGasFile(gasFile)) {
    G4cout << "GarfieldPhysics: Successfully loaded gas file " << gasFile << G4endl;
  } else {
    G4cout << "GarfieldPhysics: Failed to load gas file " << gasFile << " - Using calculated properties" << G4endl;
  }

  fComponentAnalyticField = new Garfield::ComponentAnalyticField();
  fComponentAnalyticField->SetMedium(fMediumMagboltz);

  constexpr double rWire = 0.005;  // 50 µm = 0.005 cm
  constexpr double rTube = 4.77;   // Approx. 9.54 cm box
  constexpr double vWire = 3000.;  // Real detector voltage
  constexpr double vTube = 0.;

  fComponentAnalyticField->AddWire(0., 0., 2 * rWire, vWire, "w");
  fComponentAnalyticField->AddTube(rTube, vTube, 0, "t");

  fSensor = new Garfield::Sensor();
  fSensor->AddComponent(fComponentAnalyticField);

  fTrackHeed = new Garfield::TrackHeed();
  fTrackHeed->SetSensor(fSensor);
  fTrackHeed->EnableDeltaElectronTransport();
  G4cout << "GarfieldPhysics: Initialized with rWire = " << rWire << " cm, vWire = " << vWire << " V" << G4endl;
}

void GarfieldPhysics::DoIt(std::string particleName, double ekin_MeV,
                           double time, double x_cm, double y_cm, double z_cm,
                           double dx, double dy, double dz) {
  G4cout << "GarfieldPhysics: DoIt called for " << particleName 
         << " with Ekin = " << ekin_MeV << " MeV at (" 
         << x_cm << ", " << y_cm << ", " << z_cm << ") "
         << "Direction: (" << dx << ", " << dy << ", " << dz << ")" << G4endl;

  G4AnalysisManager* analysisManager = G4AnalysisManager::Instance();
  fEnergyDeposit = 0;
  fSecondaryParticles.clear();
  fAvalancheSize = 0;
  nsum = 0;

  Garfield::AvalancheMC drift;
  drift.SetSensor(fSensor);
  drift.SetDistanceSteps(1.e-4);
  Garfield::AvalancheMicroscopic avalanche;
  avalanche.SetSensor(fSensor);

  constexpr double rWire = 0.005;  // 50 µm = 0.005 cm
  constexpr double rTube = 4.77;   // 9.54 cm box approximated
  constexpr double lTube = 300.;   // 6 m length / 2
  constexpr double prcHeight = 9.54; // Inner height in cm

  double eKin_eV = ekin_MeV * 1e+6;

  if (fIonizationModel != "Heed" || particleName == "gamma") {
    int nc = 0;
    if (particleName == "gamma") {
      fTrackHeed->TransportPhoton(x_cm, y_cm, z_cm, time, eKin_eV, dx, dy, dz, nc);
    } else {
      fTrackHeed->TransportDeltaElectron(x_cm, y_cm, z_cm, time, eKin_eV, dx, dy, dz, nc);
      fEnergyDeposit = eKin_eV;
    }
    G4cout << "GarfieldPhysics: " << nc << " clusters from " 
           << (particleName == "gamma" ? "photon" : "delta electron") << G4endl;

    for (int cl = 0; cl < nc; cl++) {
      double xe, ye, ze, te;
      double ee, dxe, dye, dze;
      fTrackHeed->GetElectron(cl, xe, ye, ze, te, ee, dxe, dye, dze);
      if (fabs(ze) > lTube || sqrt(xe * xe + ye * ye) > rTube || fabs(ye) > prcHeight / 2) {
        G4cout << "GarfieldPhysics: Electron " << cl << " at (" << xe << ", " << ye << ", " << ze << ") out of PRC bounds" << G4endl;
        continue;
      }
      nsum++;
      if (particleName == "gamma") fEnergyDeposit += fTrackHeed->GetW();
      G4cout << "GarfieldPhysics: Electron " << cl << " at (" << xe << ", " << ye << ", " << ze << ") energy = " << ee / 1e6 << " MeV" << G4endl;

      drift.DriftElectron(xe, ye, ze, te);
      double xe1, ye1, ze1, te1;
      double xe2, ye2, ze2, te2;
      int status;
      drift.GetElectronEndpoint(0, xe1, ye1, ze1, te1, xe2, ye2, ze2, te2, status);
      G4cout << "GarfieldPhysics: Electron " << cl << " drifted from (" << xe1 << ", " << ye1 << ", " << ze1 << ") to (" << xe2 << ", " << ye2 << ", " << ze2 << ")" << G4endl;

     
      if (0 < xe2 && xe2 < rWire) xe2 += 2 * rWire;
      else if (0 > xe2 && xe2 > -rWire) xe2 -= 2 * rWire;
      if (0 < ye2 && ye2 < rWire) ye2 += 2 * rWire;
      else if (0 > ye2 && ye2 > -rWire) ye2 -= 2 * rWire;

      double e2 = 0.1;
      avalanche.AvalancheElectron(xe2, ye2, ze2, te2, e2, 0, 0, 0);
      int ne = 0, ni = 0;
      avalanche.GetAvalancheSize(ne, ni);
      fAvalancheSize += ne;
      G4cout << "GarfieldPhysics: Avalanche size = " << ne << " for electron " << cl << " at (" << xe2 << ", " << ye2 << ", " << ze2 << ")" << G4endl;
    }
  } else {
    fTrackHeed->SetParticle(particleName);
    fTrackHeed->SetKineticEnergy(eKin_eV);
    double trackLength = prcHeight; // Use full height for muon path
    fTrackHeed->NewTrack(x_cm, y_cm, z_cm, time, dx * trackLength, dy * trackLength, dz * trackLength);

    int clusterCount = 0;
    for (const auto& cluster : fTrackHeed->GetClusters()) {
      if (fabs(cluster.z) > lTube || sqrt(cluster.x * cluster.x + cluster.y * cluster.y) >= rTube || fabs(cluster.y) > prcHeight / 2) {
        G4cout << "GarfieldPhysics: Cluster at (" << cluster.x << ", " << cluster.y << ", " << cluster.z << ") out of PRC bounds" << G4endl;
        continue;
      }
      clusterCount++;
      nsum += cluster.electrons.size();
      fEnergyDeposit += cluster.energy;
      G4cout << "GarfieldPhysics: Cluster " << clusterCount << " at (" << cluster.x << ", " << cluster.y << ", " << cluster.z << ") energy = " << cluster.energy / 1e6 << " MeV, electrons = " << cluster.electrons.size() << G4endl;

      for (const auto& electron : cluster.electrons) {
        if (fabs(electron.z) > lTube || sqrt(electron.x * electron.x + electron.y * electron.y) >= rTube || fabs(electron.y) > prcHeight / 2) {
          G4cout << "GarfieldPhysics: Electron at (" << electron.x << ", " << electron.y << ", " << electron.z << ") out of PRC bounds" << G4endl;
          continue;
        }

        drift.DriftElectron(electron.x, electron.y, electron.z, electron.t);
        double xe1, ye1, ze1, te1;
        double xe2, ye2, ze2, te2;
        int status;
        drift.GetElectronEndpoint(0, xe1, ye1, ze1, te1, xe2, ye2, ze2, te2, status);
        G4cout << "GarfieldPhysics: Electron drifted from (" << xe1 << ", " << ye1 << ", " << ze1 << ") to (" << xe2 << ", " << ye2 << ", " << ze2 << ")" << G4endl;

        // Match original: Adjust endpoint to avoid wire
        if (0 < xe2 && xe2 < rWire) xe2 += 2 * rWire;
        else if (0 > xe2 && xe2 > -rWire) xe2 -= 2 * rWire;
        if (0 < ye2 && ye2 < rWire) ye2 += 2 * rWire;
        else if (0 > ye2 && ye2 > -rWire) ye2 -= 2 * rWire;

        double e2 = 0.1;
        avalanche.AvalancheElectron(xe2, ye2, ze2, te2, e2, 0, 0, 0);
        int ne = 0, ni = 0;
        avalanche.GetAvalancheSize(ne, ni);
        fAvalancheSize += ne;
        G4cout << "GarfieldPhysics: Avalanche size = " << ne << " at (" << xe2 << ", " << ye2 << ", " << ze2 << ")" << G4endl;
      }
    }
    G4cout << "GarfieldPhysics: " << clusterCount << " clusters processed" << G4endl;
  }

  fGain = nsum > 0 ? fAvalancheSize / nsum : 0;
  G4cout << "GarfieldPhysics: Total Energy Deposit (Gas) = " 
         << fEnergyDeposit / 1e6 << " MeV, nsum = " << nsum 
         << ", AvalancheSize = " << fAvalancheSize << ", Gain = " << fGain << G4endl;
}