/* This file is part of MAUS: http://micewww.pp.rl.ac.uk/projects/maus
 *
 * MAUS is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * MAUS is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with MAUS.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <math.h>

#include <string>

#include "Geant4/G4Material.hh"

#include "src/common_cpp/Simulation/GeometryNavigator.hh"
#include "src/common_cpp/Utils/Squeak.hh"

#include "src/common_cpp/Recon/Global/MaterialModel.hh"

namespace MAUS {

std::map<std::string, bool> MaterialModel::_enabled_materials;
bool MaterialModel::_default_enablement = true;


void MaterialModel::SetMaterial(const G4Material* material) {
    _material = material;

    _n_e = _material->GetElectronDensity();
    _I = _material->GetIonisation()->GetMeanExcitationEnergy();
    _x_0 = _material->GetIonisation()->GetX0density();
    _x_1 = _material->GetIonisation()->GetX1density();
    _C = _material->GetIonisation()->GetCdensity();
    _a = _material->GetIonisation()->GetAdensity();
    _k = _material->GetIonisation()->GetMdensity();
    _rad_len_ratio = 13.6*13.6/_material->GetRadlen();
    _density = _material->GetDensity()/g*cm3;
}

double MaterialModel::MeandEdx(double E, double m, double charge) {
    // bethe bloch formula, see PDG particle data book
    // assume density effect is 0
    double beta2 = (E*E-m*m)/E/E; // beta^2 = p^2/E^2
    double gamma = E/m;
    double bg2 = beta2*gamma*gamma;
    double mRatio = _m_e/m;
    double T_max = 2.0*_m_e*bg2/(1.0 + 2.0*gamma*mRatio + mRatio*mRatio);

    double coefficient = std::log(2.0*_m_e*bg2*T_max/(_I*_I))/2. - beta2;
    // 0.307075 is cm^2 per mol; _density is g/cm^3
    coefficient *= 0.307075*charge*charge/beta2*_density/cm;

    double dEdx = 0.;
    for (size_t i = 0; i < _material->GetNumberOfElements(); ++i) {
        double frac = _material->GetFractionVector()[i];
        double a_el = _material->GetElement(i)->GetA()/(g/mole);
        double z_el = _material->GetElement(i)->GetZ();
        dEdx += frac*z_el/a_el*coefficient;
    }
    return -dEdx;
}

double MaterialModel::MostProbabledEdx(double E, double m, double charge) {
    double dens_thickness = _thickness*_density/cm; // g/cm^2
    double beta2 = (E*E-m*m)/E/E; // beta^2 = p^2/E^2
    double gamma = E/m;
    double bg2 = beta2*gamma*gamma;
    double meanZOverA = 0;
    for (size_t i = 0; i < _material->GetNumberOfElements(); ++i) {
        meanZOverA += _material->GetFractionVector()[i]*
                      _material->GetElement(i)->GetZ()/
                      (_material->GetElement(i)->GetA()/(g/mole));
    }
    double xi = 0.307075/2.*meanZOverA*dens_thickness/beta2;
    double dE = -xi*(log(2*m*bg2/_I)+log(xi/_I)-beta2); // ignore density effect
    return dE/_thickness;
}

double MaterialModel::dEdx(double E, double m, double charge) {
    if (_material == NULL) {
        return 0.;
    }
    if (_will_use_mean_energy_loss) {
        return MeandEdx(E, m, charge);
    } else {
        return MostProbabledEdx(E, m, charge);
    }
}

double MaterialModel::d2EdxdE(double E, double m, double charge) {
    double delta_energy_pos = dEdx(E+_d2EdxdE_delta_const, m, charge);
    double delta_energy_neg = dEdx(E-_d2EdxdE_delta_const, m, charge);
    double deriv = (delta_energy_pos-delta_energy_neg)/2./_d2EdxdE_delta_const;
    return deriv;
}

double MaterialModel::dtheta2dx(double E, double m, double charge) {
    // moliere scattering formula, see PDG particle data book
    // assume log term is 0
    double one_over_beta_c_p = E/(E*E-m*m); // E/p^2
    double theta_0_sq = _rad_len_ratio*one_over_beta_c_p*one_over_beta_c_p;
    return theta_0_sq;
}

double MaterialModel::estrag2(double E, double m, double charge) {
    // NIM A 532, Neuffer, "Introduction to Ionisation Cooling", 2004
    // Scraped from Ann. Rev. Nucl. Sci. 13, Fano, 1963
    // Note that I ignore the units in the NIM paper as I think it is
    // wrong; I think the correct units are MeV^2/cm, not MeV^2 cm^2/g
    // (Rogers, 2016)
    double gamma2 = E*E/m/m;
    double beta2 = (E*E-m*m)/E/E;
    double estrag2 = _estrag_const*gamma2*(1-beta2/2.)/cm; // missing factor Z/A

    return estrag2;
}



bool MaterialModel::IsEnabledMaterial(std::string material_name) {
    bool found = _enabled_materials.find(material_name) !=
                                                      _enabled_materials.end();
    if (found) {
        return _enabled_materials[material_name];
    } else if (material_name.size() > 2 && material_name.substr(0, 3) == std::string("G4_")) {
        material_name = material_name.substr(3);
    }
    if (_enabled_materials.find(material_name) != _enabled_materials.end()) {
        return _enabled_materials[material_name];
    }
    return _default_enablement;
}

void MaterialModel::EnableMaterial(std::string material_name) {
    _enabled_materials[material_name] = true;
}

void MaterialModel::ClearEnabledMaterials() {
    _enabled_materials = std::map<std::string, bool>();
}

void MaterialModel::SetDefaultMaterialEnablement(bool default_enablement) {
    _default_enablement = default_enablement;
}


void MaterialModel::DisableMaterial(std::string material_name) {
    _enabled_materials[material_name] = false;
}

void MaterialModel::SetWillUseMeanEnergyLoss(bool will_use_mean_energy_loss) {
    _will_use_mean_energy_loss = will_use_mean_energy_loss;
}


bool MaterialModel::GetWillUseMeanEnergyLoss() const {
    return _will_use_mean_energy_loss;
}


std::ostream& MaterialModel::PrintMaterials(std::ostream& out) {
    out << "Materials:\n";
    for (std::map<std::string, bool>::iterator it = _enabled_materials.begin();
         it != _enabled_materials.end(); ++it) {
        out << "    " << it->first << "    " << it->second << "\n";
    }
    return out;
}

} // namespace MAUS