/**
 * @file
 * PREM Earth model. Adopted from DaMaSCUS https://github.com/temken/DaMaSCUS.
 * 
 * Only non-relativistic DM has been tested for cross-checking purposes.
 * It will be re-implemented in future versions.
 */

#pragma once

#include <iostream>
#include <stdexcept>
#include "Const.hpp"
#include "Base.hpp"
#include "SIDM.hpp"

namespace darkprop {

template<typename Vector3=Vector3d<double>, typename Value=double>
class PREMEarth : public Earth<Vector3, Value>
{
private:
    Value radius;

    //                          0           1           2           3           4           5           6           7           8           9
    // const Value PREM_Radii[10]={1221.5*unit::km, 3480*unit::km, 5701*unit::km, 5771*unit::km, 5971*unit::km, 6151*unit::km, 6346.6*unit::km, 6356*unit::km, 6368*unit::km, constant::rEarth};
    const std::array<Value, 10> PREM_Radii {1221.5*unit::km, 3480*unit::km, 5701*unit::km, 5771*unit::km, 5971*unit::km, 6151*unit::km, 6346.6*unit::km, 6356*unit::km, 6368*unit::km, constant::rEarth};
    //Density Coefficients, we omit the ocean layer with density 1.02, since experiments are located under rocks, not water.
    //                  0           1           2           3           4           5           6           7       8       9       10
    //                  I.C.        O.C.        L.Mantle    Trans. 1    Trans. 2    Trans. 3    LVZ&LID     Crust1  Crust2  Ocean   Space
    const Value a[11] ={13.0885,    12.5815,    7.9565,     5.3197,     11.2494,    7.1089,     2.6910,     2.9,    2.6,    2.6,    0};
    const Value b[11] ={0,          -1.2638,    -6.4761,    -1.4836,    -8.0298,    -3.8045,    0.6924,     0,      0,      0,      0};
    const Value c[11] ={-8.8381,    -3.6426,    5.5283,     0,          0,          0,          0,          0,      0,      0,      0};
    const Value d[11] ={0,          -5.5281,    -3.0807,    0,          0,          0,          0,          0,      0,      0,      0};

    const std::vector<Target> elements_core {
        Target({"Fe", 26, 56, 56*constant::mu}),  // Iron
        Target({"Si", 14, 28, 28*constant::mu}),  // Silicon
        Target({"Ni", 28, 58, 58*constant::mu}),  // Nickel
        Target({"S", 16, 32, 32*constant::mu}),   // Sulfur
        Target({"Cr", 24, 52, 52*constant::mu}),  // Chromium
        Target({"Mn", 25, 55, 55*constant::mu}),  // Manganese
        Target({"P", 15, 31, 31*constant::mu}),   // Phosphorus
        Target({"C", 6, 12, 12*constant::mu}),    // Carbon
        Target({"H", 1, 1, constant::mp}),        // Hydrogen
    };
    static constexpr std::array<Value, 9> elements_core_fraction {
        0.855, 0.06, 0.052, 0.019, 0.009, 0.003, 0.002, 0.002, 0.0006
    };
    const std::vector<Target> elements_mantle {
        Target({"O", 8, 16, 16*constant::mu}),    // Oxygen
        Target({"Mg", 12, 24, 24*constant::mu}),  // Magnesium
        Target({"Si", 14, 28, 28*constant::mu}),  // Silicon
        Target({"Fe", 26, 56, 56*constant::mu}),  // Iron
        Target({"Ca", 20, 40, 40*constant::mu}),  // Calcium
        Target({"Al", 13, 27, 27*constant::mu}),  // Aluminium
        Target({"Na", 11, 23, 23*constant::mu}),  // Natrium
        Target({"Cr", 24, 52, 52*constant::mu}),  // Chromium
        Target({"Ni", 28, 58, 58*constant::mu}),  // Nickel
        Target({"Mn", 25, 55, 55*constant::mu}),  // Manganese
        Target({"S", 16, 32, 32*constant::mu}),   // Sulfur
        Target({"C", 6, 12, 12*constant::mu}),    // Carbon
        Target({"H", 1, 1, constant::mp}),        // Hydrogen
        Target({"P", 15, 31, 31*constant::mu})    // Phosphorus
    };
    static constexpr std::array<Value, 14> elements_mantle_fraction {
        0.440, 0.228, 0.21, 0.0626, 0.0253, 0.0235, 0.0027, 0.0026, 0.002,
        0.001, 0.0003, 0.0001, 0.0001, 0.00009
    };
    std::array<Value, 9> scatter_probability_core;
    std::array<Value, 14> scatter_probability_mantle;
    std::array<Value, 2> g_PREM {0, 0};

    Value g_Factor(int layer) {
        if (layer <= 1) return g_PREM[0];
        else if (layer <= 9) return g_PREM[1];
        else return 0.0;
    }
    Value tExit(const Vector3& position, const Vector3& vel);
    Value ProbabilityExponent(const Vector3& position,
                              const Vector3& vel,
                              Value L);
    // Value ScatterProbability(const Vector3& position,
    //                          const Vector3& vel,
    //                          Value L) {
    //     return 1. - std::exp(-ProbabilityExponent(position, vel, L));
    // }
    Value NewtonMethod(const Vector3& position, const Vector3& vel, Value xi,
                       Value Lambda_Total, Value NewtonStart=500*unit::km) {
        Value r0 = position.norm();
        Value v = vel.norm();
        Value cosalpha = position.dot(vel) / r0 / v;
        int i = PREMLayer(r0);
        Value delta = 1e-10;
        Value L = NewtonStart;
        Value f, dfdL, Ltilde;
        while (std::abs(Lambda_Total + ProbabilityExponent(position, vel, L)
                   + std::log(1 - xi)) > delta) {
            Ltilde = std::sqrt(L*L + 2*L*r0*cosalpha + r0*r0);
            f = Lambda_Total + ProbabilityExponent(position, vel, L)
                + std::log(1 - xi);
            dfdL = g_Factor(i) * (a[i] + Ltilde / radius * b[i]
                    + std::pow(Ltilde, 2) / (radius * radius) * c[i]
                    + std::pow(Ltilde, 3) / std::pow(radius, 3) * d[i])
                        * unit::gram / std::pow(unit::cm, 3);
            L -= f / dfdL;
        }
        return L;
    }
    Value fRoot(Value L, Value lambda_total, const Vector3&  position,
                const Vector3& vel, Value xi) {
        return lambda_total + ProbabilityExponent(position, vel, L)
               + std::log(1 - xi);
    }
    Value bisection(const Vector3& position, const Vector3& vel, Value xi,
                    Value lambda_total, Value start=500*unit::km) {
        Value delta = 1e-12;
        Value f = fRoot(0.0, lambda_total, position, vel, xi);
        Value fmid = fRoot(start, lambda_total, position, vel, xi);
        if (f >= 0.0 || fmid <= 0.0) {
            std::cerr << "f = " << f << " fmid = " << fmid << std::endl;
            throw std::runtime_error("Root must be bracketed for bisection.");
        }
        Value dL = start;
        Value L = 0;
        Value xmid;
        while (dL >= delta * start && fmid != 0.0) {
            dL *= 0.5;
            xmid = L + dL;
            fmid = fRoot(xmid, lambda_total, position, vel, xi);
            if (fmid <= 0.0) {
                L = xmid;
            }
        }
        return L;
    }

public:
    Value rfinal;
    bool cross_section_energy_dependent;
    PREMEarth(Value t_radius=constant::rEarth,
              Value t_rfinal=unit::m,
              bool t_energy_dependent=false)
        : radius {t_radius}, rfinal {t_rfinal},
          cross_section_energy_dependent {t_energy_dependent} {
        this->targets = elements_mantle;
    }
    ~PREMEarth() {}
    int PREMLayer(Value r) const {
        //The 10 mechanical Layers
        // return std::lower_bound(PREM_Radii.begin(),
        //                         PREM_Radii.end(), r) - PREM_Radii.begin();
        for (int i = 9; i >= 0; i--) {
            if (r > PREM_Radii[i]) {
                return i + 1;
            }
        }
        //Inner Core:
        return 0;
    }
    Value PREM_Mass_Density(Value r) {
        int i = PREMLayer(r);
        Value x = r / radius;
        return (a[i]
                + b[i] * x
                + c[i] * std::pow(x, 2)
                + d[i] * std::pow(x, 3)) * unit::gram * std::pow(unit::cm, -3);
    }
    void setPREM(const Particle<Vector3, Value>& p);
    void setCache(const Particle<Vector3, Value>& p) override { setPREM(p); };
    const std::vector<Target>& getTargets() override { return this->targets; }
    void propagate(Particle<Vector3, Value>& particle,
                   RandomNumber<Value>& rn) override;
    Target sampleTarget(const Particle<Vector3, Value>& particle,
                        RandomNumber<Value>& rn) const override;
};


template<typename Vector3, typename Value>
void PREMEarth<Vector3, Value>::setPREM(const Particle<Vector3, Value> & p)
{
    g_PREM[0] = 0.0;
    g_PREM[1] = 0.0;

    //core
    for (std::size_t i = 0; i < elements_core.size(); ++i) {
        scatter_probability_core[i] = elements_core_fraction[i]
            / elements_core[i].m * p.totalCrossSection(elements_core[i]);
        g_PREM[0] += scatter_probability_core[i];
    }
    for (Value& prob : scatter_probability_core) {
        prob /= g_PREM[0];
    }

    //mantle
    for (std::size_t i = 0; i < elements_mantle.size(); ++i) {
        scatter_probability_mantle[i] = elements_mantle_fraction[i]
            / elements_mantle[i].m * p.totalCrossSection(elements_mantle[i]);
        g_PREM[1] += scatter_probability_mantle[i];
    }
    for (Value& prob : scatter_probability_mantle) {
        prob /= g_PREM[1];
    }
}

template<typename Vector3, typename Value>
Target PREMEarth<Vector3, Value>::sampleTarget(
        const Particle<Vector3, Value>& p,
        RandomNumber<Value>& rn) const
{
    // Note setPREM is not called here
    int layer = PREMLayer(p.r.norm());
    Value xi = rn.uniform_xi();
    Value sum = 0.0;
    // The two compositional layers
    // Mantle
    if (layer > 1 && layer < 10) {
        for (std::size_t i = 0; i < elements_mantle.size(); ++i) {
            sum += scatter_probability_mantle[i];
            if (sum > xi) {
                return elements_mantle[i];
            }
        }
    }
    // Core
    else if (layer < 2) {
        for (std::size_t i = 0; i < elements_core.size(); ++i) {
            sum += scatter_probability_core[i];
            if (sum > xi) {
                return elements_core[i];
            }
        }
    }
    throw std::runtime_error("no target");
}

// When does a particle leave its current layer?
template<typename Vector3, typename Value>
Value PREMEarth<Vector3, Value>::tExit(const Vector3& position,
                                       const Vector3& vel)
{
        Value r = position.norm();
        Value v = vel.norm();
        Value rv = position.dot(vel);
        int layer = PREMLayer(r);
        // Check if the particle is well inside the earth.
        if (layer == 10) {
            std::cerr << "Error in tExit(): "
                      << "Position argument lies outside the earth."
                      << std::endl;
            exit(1);
        }
        // If not we have to distinguish two cases.
        Value R[2];
        // Case A: Most often position lies directly on the border between
        // two layers.
        if(r == PREM_Radii[layer]) {
            if (layer == 9) {
            // Position directly at the Earth surface:
                // Leaving the Earth
                if(rv/r/v >= 0) return 1 * unit::km/v;
                // Entering the Earth.
                R[0] = PREM_Radii[9];
                R[1] = PREM_Radii[8];
            } else if (layer == 0) {
            // Position directly at core border:
                R[0] = PREM_Radii[1];
                R[1] = PREM_Radii[0];
            } else {
                R[0] = PREM_Radii[layer + 1];
                R[1] = PREM_Radii[layer - 1];
            }
        } else {
        // Case B: Position is not at the layer border:
            // Core
            if (layer == 0) {
                return (-rv + std::sqrt(std::pow(rv, 2) - std::pow(v*r, 2)
                        +std::pow(v*PREM_Radii[layer], 2))) / std::pow(v, 2)
                       + 10 * unit::cm / v;
            } else {
                R[0] = PREM_Radii[layer];
                R[1] = PREM_Radii[layer - 1];
            } 
        }
        // With the two radii we can find tExit.
        Value texit = 0.0;
        // Outer radius:
        texit = (-rv + std::sqrt(std::pow(rv, 2) - std::pow(v*r, 2)
                 + std::pow(v*R[0], 2))) / std::pow(v, 2);
        // Inner Radius
        Value radic = std::pow(rv, 2) - std::pow(v*r, 2) + std::pow(v*R[1], 2);
        if (radic >= 0) {
            Value t1 = (-rv - std::sqrt(radic)) / std::pow(v, 2);
            Value t2 = (-rv + std::sqrt(radic)) / std::pow(v, 2);
            if (t1 > 0 && t1 < texit) texit = t1;
            if (t2 > 0 && t2 < texit) texit = t2;
        }
    return texit + 10 * unit::cm / v;
}

/**
 * Integrate lambda^-1 in one layer.
 */
template<typename Vector3, typename Value>
Value PREMEarth<Vector3, Value>::ProbabilityExponent(const Vector3& position,
                                                     const Vector3& vel,
                                                     Value L)
{
    Value r0 = position.norm();
    Value v = vel.norm();
    Value cosalpha = position.dot(vel)/r0/v;
    int i = PREMLayer(r0);
    Value Ltilde = std::sqrt(L*L + 2*L*r0*cosalpha + r0*r0);
    // Coefficients
    Value C1, C2, C3, C4;
    C1 = L * radius * radius;
    C3 = L * (r0*r0 + r0*L*cosalpha + L*L/3);
    Value epsilon = 1e-14;
    if (std::abs(cosalpha + 1) > epsilon) {
        C2 = radius / 2 * (Ltilde * (L + r0*cosalpha) - r0*r0*cosalpha
                           + (1 - cosalpha*cosalpha) * r0 * r0
                              * std::log((L + Ltilde + r0*cosalpha)
                                          / ((1 + cosalpha)*r0)));
        C4 = 1.0 / 8.0 / radius * ((5 - 3 * cosalpha * cosalpha)
                * (cosalpha*std::pow(r0,3) * Ltilde - std::pow(r0, 4)*cosalpha)
                + 2*L*L*Ltilde * (L + 3*r0*cosalpha)
                + L*Ltilde*r0*r0 * (5 + cosalpha*cosalpha)
                + 3*std::pow(r0, 4) * std::pow(1 - cosalpha*cosalpha, 2)
                  * std::log((L + Ltilde + r0*cosalpha)/((1 + cosalpha)*r0)));
    } else {
        C2 = (L * (L - 2 * r0) * std::sqrt(std::pow(L - r0, 2)) * radius)
              / (2. * (L - r0));
        C4 = (L * (L - 2 * r0) * std::sqrt(std::pow(L - r0, 2))
                * (std::pow(L, 2) - 2 * L * r0 + 2 * std::pow(r0, 2)))
              / (4. * (L - r0) * radius);
    }
    return g_Factor(i) * (a[i] * C1 + b[i] * C2 + c[i] * C3 + d[i] * C4)
           / radius / radius * unit::gram * std::pow(unit::cm, -3);
}

template<typename Vector3, typename Value>
void PREMEarth<Vector3, Value>::propagate(Particle<Vector3, Value>& p,
                                          RandomNumber<Value>& rn)
{
    if (cross_section_energy_dependent) setPREM(p);
    Value v = p.v.norm();
    Value xi = rn.uniform_xi();
    Value logxi = -log(1 - xi);
    Value t_tot = 0.0;
    // Add up terms in the exponent of P(L)
    Value Lambda = 0;
    Value t, L, Lambda_l, LambdaNew;
    Vector3 r = p.r;
    int layer = PREMLayer(r.norm());
    while (layer < 10) {
        t = tExit(r, p.v);
        t_tot += t;
        L = t * v;
        Lambda_l = ProbabilityExponent(r, p.v, L);
        LambdaNew = Lambda + Lambda_l;
        if (LambdaNew < logxi) {
        // No scattering in this layer
            Lambda = LambdaNew;
            r += t * p.v;
            layer = PREMLayer(r.norm());
        } else {
        // Scattering in this layer
            Value l = cross_section_energy_dependent ?
                bisection(r, p.v, xi, Lambda, L) :
                NewtonMethod(r, p.v, xi, Lambda, (logxi-Lambda)/Lambda_l * L);
            // Value l = bisection(r, p.v, xi, Lambda, L);
            t = l / v;
            r += t * p.v;
            t_tot += t;
            break;
        }
    }
    if (r.norm() > constant::rEarth) {
        Value costheta = p.ep.dot(-p.er);
        Value rnorm = p.r.norm();
        Value freep = rfinal + rnorm * costheta
                       + std::sqrt(rnorm*rnorm * (costheta*costheta - 1)
                       + radius*radius);
        p.r += freep * p.ep;
        p.t += freep / p.v.norm();
        p.in_earth = false;
    } else {
        p.r = r;
        p.t += t_tot;
    }
    p.updateEr();
}

} // namespace darkprop
