#include "PrimaryGeneratorAction.h"

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

G4String particleType;
G4double Z, A;
G4double M = 931.5;
G4double C1, C2, C3, C4, sum;
std::string tmp, line, name;
std::ifstream modelFile("F:/Project/Geant4/DESCSS/assets/model.txt");

PrimaryGeneratorAction<G4GeneralParticleSource>::PrimaryGeneratorAction() {
    fParticleGun = new G4GeneralParticleSource();
}

PrimaryGeneratorAction<G4ParticleGun>::PrimaryGeneratorAction() {
    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 (tmp == "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;
    fParticleGun = new G4ParticleGun();
    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);
}

template <typename T>
PrimaryGeneratorAction<T>::~PrimaryGeneratorAction() {
    delete fParticleGun;
}

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

G4double randomTheta() {
    G4double u = G4UniformRand();
    return std::acos(2 * u - 1);
}

G4ThreeVector randomPos(G4double rho) {
    G4double phi = randomPhi();
    G4double theta = randomTheta();

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

    return G4ThreeVector(x, y, z);
}

G4ThreeVector randomDir() {
    G4double phi = randomPhi();
    G4double theta = randomTheta();

    G4double x = -std::sin(theta) * std::cos(phi);
    G4double y = -std::sin(theta) * std::sin(phi);
    G4double z = -std::cos(theta);

    return G4ThreeVector(x, y, z);
}

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

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 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<G4GeneralParticleSource>::GeneratePrimaries(G4Event* e) {
    fParticleGun->GeneratePrimaryVertex(e);
}

void PrimaryGeneratorAction<G4ParticleGun>::GeneratePrimaries(G4Event* e) {
    G4double E;
    G4double R = 15. * m;
    G4ThreeVector pos = randomPos(R);
    G4ThreeVector dir = randomDir();

    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);
    fParticleGun->SetParticlePosition(pos);
    fParticleGun->SetParticleMomentumDirection(dir);
    fParticleGun->GeneratePrimaryVertex(e);
}
