#include "Geant/Solid.h"
#include "HEP/Geant/GeantEvent.h"
#include "HEP/Geant/Scene.h"
#include "HEP/Geant/EmitterPrimary.hh"
#include "Math/ContainerBox.h"
#include "Math/Dense.h"
#include "Math/Units.h"
#include "HEP/Detectors/DetectorChamber.h"

#include <Geant4/G4SystemOfUnits.hh>

#include <iostream>

using namespace uLib;

int main(int argc, char** argv) {

    int nEvents = 100;
    if (argc > 1) {
        nEvents = std::stoi(argv[1]);
    }

    // 1. Setup Geant4 Scene
    Geant::Scene scene;
    scene.ConstructWorldBox(Vector3f(30_m, 30_m, 30_m), "G4_AIR");

    ContainerBox *iron_box = new ContainerBox();
    iron_box->Scale(Vector3f(18_m, 10_cm, 18_m));
    iron_box->SetPosition(Vector3f(-9_m, -5_cm, -9_m));
    Geant::BoxSolid *iron_cube = new Geant::BoxSolid("IronCube", iron_box);
    Geant::Material *iron_mat = new Geant::Material("G4_Fe");
    Geant::LogicalVolume *iron_lv = new Geant::LogicalVolume("IronCube_lv");
    iron_lv->SetSolid(iron_cube);
    iron_lv->SetMaterial(iron_mat);
    iron_lv->Update();
    scene.AddVolume(new Geant::PhysicalVolume("IronCube", iron_lv));

    // Top Detector Chamber (along Y axis)
    DetectorChamber* top_chamber_box = new DetectorChamber();
    top_chamber_box->Scale(Vector3f(20_m, 40_cm, 20_m));
    top_chamber_box->Rotate(90_deg, Vector3f(1, 0, 0));
    top_chamber_box->SetPosition(Vector3f(-10_m, 12_m, -10_m));
    Geant::BoxSolid* top_chamber = new Geant::BoxSolid("TopChamber", top_chamber_box);
    SmartPointer<Geant::Material> air_mat(new Geant::Material("G4_AIR"));
    Geant::LogicalVolume* top_chamber_lv = new Geant::LogicalVolume("TopChamber_lv");
    top_chamber_lv->SetSolid(top_chamber);
    top_chamber_lv->SetMaterial(air_mat);
    top_chamber_lv->Update();
    scene.AddVolume(new Geant::PhysicalVolume("TopChamber", top_chamber_lv));

    // Bottom Detector Chamber (along Y axis)
    DetectorChamber* bottom_chamber_box = new DetectorChamber();
    bottom_chamber_box->Scale(Vector3f(20_m, 40_cm, 20_m));
    bottom_chamber_box->Rotate(90_deg, Vector3f(1, 0, 0));
    bottom_chamber_box->SetPosition(Vector3f(-10_m, -12_m, -10_m));
    Geant::BoxSolid* bottom_chamber = new Geant::BoxSolid("BottomChamber", bottom_chamber_box);
    Geant::LogicalVolume* bottom_chamber_lv = new Geant::LogicalVolume("BottomChamber_lv");
    bottom_chamber_lv->SetSolid(bottom_chamber);
    bottom_chamber_lv->SetMaterial(air_mat);
    bottom_chamber_lv->Update();
    scene.AddVolume(new Geant::PhysicalVolume("BottomChamber", bottom_chamber_lv));

    // Setup SkyPlaneEmitterPrimary
    Geant::SkyPlaneEmitterPrimary* emitter = new Geant::SkyPlaneEmitterPrimary();
    emitter->SetPosition(Vector3f(0, 14.9_m, 0));
    emitter->Rotate(-90_deg, Vector3f(1, 0, 0));
    emitter->SetSkySize(Vector2f(20_m, 20_m)); 
    
    scene.SetEmitter(emitter);
    scene.SetVerbosity(1);
    // scene.Initialize(); // Removed to avoid premature initialization


    
    std::cout << "Starting simulation of " << nEvents << " events..." << std::endl;
    Vector<Geant::GeantEvent> results;
    scene.RunSimulation(nEvents, results);

    std::cout << "Simulation finished. Collected " << results.size() << " events." << std::endl;

    // Sample output to verify data collection
    if (!results.empty()) {
        std::cout << "Summary: " << std::endl;
        std::cout << "  Total events generated: " << results.size() << std::endl;
        size_t total_steps = 0;
        for (const auto& event : results) {
            total_steps += event.Path().size();
        }
        std::cout << "  Total simulation steps: " << total_steps << std::endl;
        std::cout << "  Average steps per event: " << static_cast<double>(total_steps) / results.size() << std::endl;
    }

    std::cout << "\nStarting Detector Simulation of " << nEvents << " events..." << std::endl;
    Vector<MuonEvent> detectorResults;
    scene.RunDetectorSimulation(nEvents, detectorResults);



    std::cout << "Detector Simulation finished." << std::endl;
    size_t hit_count = 0;
    for (const auto& ev : detectorResults) {
        if (!std::isnan(ev.LineIn().origin.x())) {
            hit_count++;
        }
    }
    std::cout << "  Muons crossing at least one detector: " << hit_count << std::endl;
    if (nEvents > 0) {
        std::cout << "  Efficiency: " << (100.0 * hit_count / nEvents) << "%" << std::endl;
    }

    return 0;
}