#include "DetectorConstruction.h"
#include "Matrix.h"
#include "PrimaryGeneratorAction.h"

#include "G4IonTable.hh"
#include "G4ParticleDefinition.hh"
#include "G4ParticleTable.hh"
#include "G4PhysicalConstants.hh"
#include "G4RunManager.hh"
#include "G4SystemOfUnits.hh"
#include "Randomize.hh"

G4double M = 931.5;
G4double C1, C2, C3, C4, sum;

PrimaryGeneratorAction::PrimaryGeneratorAction() { fParticleGun = new G4ParticleGun(); }

PrimaryGeneratorAction::~PrimaryGeneratorAction() { delete fParticleGun; }

void PrimaryGeneratorAction::DefineParticle(G4String particleType) {
    G4double Z, A;
    std::string tmp, line, name;
    std::ifstream modelFile("assets/model.txt");

    while (std::getline(modelFile, line)) {
        std::stringstream ss(line);
        std::getline(ss, name, ',');
        if (name == particleType) {
            std::getline(ss, tmp, ','), Z = std::stof(tmp);
            std::getline(ss, tmp, ','), A = std::stof(tmp);
            std::getline(ss, tmp, ','), C1 = std::stof(tmp);
            std::getline(ss, tmp, ','), C2 = std::stof(tmp);
            std::getline(ss, tmp, ','), C3 = std::stof(tmp);
            if (name == "TE")
                sum = C3;
            else if (name == "TP")
                std::getline(ss, tmp, ','), sum = std::stof(tmp);
            else {
                std::getline(ss, tmp, ','), C4 = std::stof(tmp);
                std::getline(ss, tmp, ','), sum = std::stof(tmp);
            }
            break;
        }
    }

    G4ParticleDefinition* ion;
    if (Z == -1)
        ion = G4ParticleTable::GetParticleTable()->FindParticle("e-");
    else if (Z == 1)
        ion = G4ParticleTable::GetParticleTable()->FindParticle("proton");
    else if (Z == 2)
        ion = G4ParticleTable::GetParticleTable()->FindParticle("alpha");
    else
        ion = G4IonTable::GetIonTable()->GetIon(Z, A, 0.);

    fParticleGun->SetParticleCharge(Z * eplus);
    fParticleGun->SetParticleDefinition(ion);
}

G4double randomPhi() {
    G4double v = G4UniformRand();
    return 2 * pi * v;
}

G4double randomTheta(G4double minU = -1, G4double maxU = 1) {
    G4double u = G4UniformRand();
    u = u * (maxU - minU) + minU;
    return acos(u);
}

G4ThreeVector randomPos(G4double rho = 12. * m) {
    G4double phi = randomPhi();
    G4double theta = randomTheta();

    G4double x = rho * sin(theta) * cos(phi);
    G4double y = rho * sin(theta) * sin(phi);
    G4double z = rho * cos(theta);

    return G4ThreeVector(x, y, z);
}

G4ThreeVector randomDir(G4ThreeVector pos, G4double maxTheta = 30. * deg) {
    G4double phi = randomPhi();
    G4double theta = randomTheta(cos(maxTheta), 1);

    G4double u = sin(theta) * cos(phi);
    G4double v = sin(theta) * sin(phi);
    G4double w = cos(theta);
    G4double mD[3][1] = {{u}, {v}, {w}};
    matrix<G4double> D(3, 1);
    D.Input(&mD[0][0]);

    G4double x0 = pos.getX();
    G4double y0 = pos.getY();
    G4double z0 = pos.getZ();
    G4double rho = sqrt(x0 * x0 + y0 * y0 + z0 * z0);
    G4double mO[3][1] = {{x0}, {y0}, {z0}};
    matrix<G4double> O(3, 1);
    O.Input(&mO[0][0]);

    G4double a = 1;
    G4double b = (z0 + rho - y0) / x0;
    G4double c = (z0 - rho + y0) / x0;

    G4double mR[3][3] = {{1 + a * a - b * b - c * c, -2 * c - 2 * a * b, -2 * b + 2 * a * c},
                         {2 * c - 2 * a * b, 1 - a * a + b * b - c * c, -2 * a - 2 * b * c},
                         {2 * b + 2 * a * c, 2 * a - 2 * b * c, 1 - a * a - b * b + c * c}};
    matrix<G4double> R(3, 3);
    R.Input(&mR[0][0]);
    R = R / (1 + a * a + b * b + c * c);

    matrix<G4double> M(3, 1);
    M = R * D;

    return G4ThreeVector(M(1, 1), M(2, 1), M(3, 1));
}

G4double ITM(G4double Emin, G4double Emax, G4double (*f)(G4double)) {
    G4double y = G4UniformRand();
    return (*f)(y);
}

G4double ARM(G4double Emin, G4double Emax, G4double (*f)(G4double)) {
    G4double x, y;
    while (true) {
        x = G4UniformRand() * (Emax - Emin) + Emin;
        y = G4UniformRand();
        if (y <= (*f)(x)) break;
    }
    return x;
}

G4double pdfTElectronInv(G4double y) {
    G4double C = C2 * C1 * exp(-0.1 / C2) / sum;
    return -C2 * log((C - y) * sum / C1 / C2);
}

G4double pdfTProtonInv(G4double y) {
    G4double C = C1 * pow(50 - C2, 1 - C3) / sum / (1 - C3);
    return pow(sum * (C + y) * (1 - C3) / C1, 1 / (1 - C3)) + C2;
}

G4double pdfGCRProton(G4double E) {
    G4double pc = sqrt(2 * M * E + E * E);
    G4double beta = pc / sqrt(M * M + pc * pc);
    return C1 * pow(beta, C2 - 1) * pow(pc / (pc + C3), C4) / pc / sum;
}

G4double pdfGCRIon(G4double E) { return C1 * exp(-C2 * E) * (1 - exp(-C3 * E + C4)) / sum; }

void PrimaryGeneratorAction::GeneratePrimaries(G4Event* e) {
    G4double E;
    G4ThreeVector pos = randomPos();
    G4ThreeVector dir = randomDir(pos);

    DetectorConstruction* detConstruction =
        (DetectorConstruction*)(G4RunManager::GetRunManager()->GetUserDetectorConstruction());
    if (particleType != detConstruction->GetParticleType()) {
        particleType = detConstruction->GetParticleType();
        DefineParticle(particleType);
    }

    if (particleType == "TE")
        E = ITM(0.1, 10, pdfTElectronInv);
    else if (particleType == "TP")
        E = ITM(50, 1000, pdfTProtonInv);
    else if (particleType == "GCR_H")
        E = ARM(220, 1e5, pdfGCRProton);
    else
        E = ARM(220, 1e5, pdfGCRIon);

    fParticleGun->SetParticleEnergy(E * MeV);
    fParticleGun->SetParticlePosition(pos);
    fParticleGun->SetParticleMomentumDirection(dir);
    fParticleGun->GeneratePrimaryVertex(e);
}