/**
 * @file
 * A homogeneous Earth model.
 */

#pragma once

#include <ios>
#include <iostream>
#include <array>
#include "Const.hpp"
#include "Base.hpp"
#include "SIDM.hpp"
#include "Vector3d.hpp"

namespace darkprop {

/**
 * A homogeneous Earth model composed of 8 components.
 */
template<typename Vector3=Vector3d<double>, typename Value=double>
class HomoEarth : public Earth<Vector3, Value>
{
public:
    Value rfinal;
    static constexpr Value density {
        2.7 * unit::gram / (unit::cm * unit::cm * unit::cm)
    };
    /// radius squared
    static constexpr Value radius2 {constant::rEarth * constant::rEarth};   
    static constexpr std::array<Value, 8> target_density_array {
        46.6e-2 / 0.985 * density / (16*constant::mu),
        27.7e-2 / 0.985 * density / (28*constant::mu),
        8.1e-2  / 0.985 * density / (27*constant::mu),
        5.0e-2  / 0.985 * density / (57*constant::mu),
        3.6e-2  / 0.985 * density / (40*constant::mu),
        2.8e-2  / 0.985 * density / (39*constant::mu),
        2.6e-2  / 0.985 * density / (23*constant::mu),
        2.1e-2  / 0.985 * density / (24*constant::mu)
    };
    bool use_cache=false;
    Value cached_lambda = 0;
    std::array<Value, 8> cached_prob_array;
    inline Value calcLambda(const Particle<Vector3, Value>& p, size_t i) const;
    inline Value calcLambda(const Particle<Vector3, Value>& p) const;
public:
    HomoEarth(Value t_rfinal=unit::cm) : rfinal {t_rfinal} {
        this->targets = std::vector<Target> {
            Target({"O", 8, 16, 16*constant::mu}),
            Target({"Si", 14, 28, 28*constant::mu}),
            Target({"Al", 13, 27, 27*constant::mu}),
            Target({"Fe", 26, 56, 56*constant::mu}),
            Target({"Ca", 20, 40, 40*constant::mu}),
            Target({"K", 19, 39, 39*constant::mu}),
            Target({"Na", 11, 23, 23*constant::mu}),
            Target({"Mg", 12, 24, 24*constant::mu})};
    }
    ~HomoEarth() {}
    const std::vector<Target>& getTargets() override { return this->targets; }
    void setCache(const Particle<Vector3, Value>& p) override;
    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>
Value HomoEarth<Vector3, Value>::calcLambda(
        const Particle<Vector3, Value>& p, size_t i) const
{
    return 1.0/(target_density_array[i]*p.totalCrossSection(this->targets[i]));
}

template<typename Vector3, typename Value>
Value HomoEarth<Vector3, Value>::
calcLambda(const Particle<Vector3, Value>& p) const
{
    Value lambda {0.0};
    for (size_t i = 0; i != 8; ++i) {
        lambda += target_density_array[i]*p.totalCrossSection(this->targets[i]);
    }
    lambda = 1.0 / lambda;
    return lambda;
}

template<typename Vector3, typename Value>
void HomoEarth<Vector3, Value>::setCache(const Particle<Vector3, Value>& p)
{
    cached_lambda = calcLambda(p);
    for (size_t i = 0; i != 8; ++i) {
        cached_prob_array[i] = cached_lambda / calcLambda(p, i);
    }
    use_cache = true;
}

template<typename Vector3, typename Value>
void HomoEarth<Vector3, Value>::propagate(
        Particle<Vector3, Value>& p,
        RandomNumber<Value>& rn)
{
    Value xi = rn.uniform_xi();
    Value freep {0.0};
    if (use_cache) {
        freep = -std::log(1 - xi) * cached_lambda;
    } else {
        freep = -std::log(1 - xi) * calcLambda(p);
    }
    Vector3 rnow = p.r;
    p.r += p.ep * freep;
    if (p.r.norm() > constant::rEarth) {
        p.in_earth = false;
        Value rnorm = rnow.norm();
        Value costheta = p.ep.dot(-p.er);
        freep = rfinal + rnorm * costheta
            + std::sqrt(rnorm*rnorm * (costheta*costheta - 1) + radius2);
        p.r = rnow + p.ep * freep;
    }
    p.t += freep / p.v.norm();
    p.updateEr();
}

template<typename Vector3, typename Value>
Target HomoEarth<Vector3, Value>::sampleTarget(
        const Particle<Vector3, Value>& p,
        RandomNumber<Value>& rn) const
{
    Value xi = rn.uniform_xi();
    Value prob = 0.0;
    if (use_cache) {
        for (size_t i = 0; i != 8; ++i) {
            prob += cached_prob_array[i];
            if (xi < prob) {
                return this->targets[i];
            }
        }
    } else {
        for (size_t i = 0; i != 8; ++i) {
            prob += (calcLambda(p) / calcLambda(p, i));
            if (xi < prob) {
                return this->targets[i];
            }
        }
    }
    // this should not happen
    return {"", 0, 0, 0.};
}

} // namespace darkprop
