/**
 * @file
 * Core classes defination.
 */

#pragma once

#include <cmath>
#include <ios>
#include <memory>
#include <random>
#include <string>
#include <iostream>
#include "Const.hpp"

namespace darkprop {

/**
 * Target nuclues or electron.
 */
struct Target
{
    std::string name;  /**< name string */
    int Z;             /**< charge number */
    int A;             /**< mass number */
    double m;          /**< mass */
    Target() {}
    Target(const std::string& t_name, int t_Z, int t_A, double t_m)
        : name {t_name}, Z {t_Z}, A {t_A}, m {t_m} {}
    ~Target() {}
};


/**
 * Random number generator.
 */
template<typename Value=double>
class RandomNumber
{
private:
    std::uniform_real_distribution<Value> distrphi;
    std::uniform_real_distribution<Value> distrcos;
    std::uniform_real_distribution<Value> distr01;
    std::mt19937_64 prng;
public:
    RandomNumber(int seed=1) :
        distrphi {std::uniform_real_distribution<Value>(0.0, 2*M_PI)},
        distrcos {std::uniform_real_distribution<Value>(-1., 1.)},
        distr01 {std::uniform_real_distribution<Value>(0., 1.)} {
        if (seed < 0) {
            std::random_device rd;
            prng = std::mt19937_64(rd());
        } else {
            prng = std::mt19937_64(seed);
        }
    }
    /**
     * Sample uniform distribution on (0, pi).
     */
    Value uniform_phi() {
        return distrphi(prng);
    }
    /**
     * Sample uniform distribution on (-1, 1).
     */
    Value uniform_costheta() {
        return distrcos(prng);
    }
    /**
     * Sample uniform distribution on (0, 1).
     */
    Value uniform_xi() {
        return distr01(prng);
    }
    /**
     * Sample uniform distribution on (a, b).
     */
    Value uniform_ab(Value a, Value b) {
        return std::uniform_real_distribution<Value>(a, b)(prng);
    }
};

/**
 * Rotate a vector v around another vector n by angle theta
 */
template<typename Vector3, typename Value>
void rotation_axis_angle(Vector3& v, const Vector3 n, Value theta)
{
	v = (v.dot(n)) * n + std::cos(theta) * (n.cross(v)).cross(n)
        + std::sin(theta) * n.cross(v);
}


/**
 * Store particle information for dump.
 */
template<typename Vector3, typename Value>
struct Event
{
    Value t;     ///> time
    Value T;     ///> kinetic energy
    Vector3 r;   ///> position
    Vector3 p3;  ///> 3 momentum
    Event() {}
    Event(Value t_t, Value t_T, Vector3 t_r, Vector3 t_p3)
        : t {t_t}, T{t_T}, r{t_r}, p3{t_p3} {}
    ~Event() {}
};


/**
 * Base abstract Particle class.
 */
template<typename Vector3, typename Value>
class Particle
{
public:
    // TODO make them protected members
    Value m;        ///< mass
    Value sigma0;   ///< cross section
    Value t;        ///< time
    Value T;        ///< kinetic energy
    Vector3 r;      ///< position
    Vector3 v;      ///< velocity
    Vector3 p3;     ///< 3 momentum
    Vector3 ep;     ///< velocity (momentum) direction
    Vector3 er;     ///< position direction
    bool in_earth;  ///< switch indicating whether to continue the simulation 
public:
    Particle(bool t_in_earth=false) : in_earth {t_in_earth} {};
    virtual ~Particle() {};

    Value pFromT() const { return std::sqrt(T * (T + 2.0 * m)); }
    Value TFromP3() const {
        Value pnorm = p3.norm();
        if (pnorm / m < 1e-4) {
            return 0.5 * pnorm * pnorm / m;
        }
        return std::sqrt(pnorm * pnorm + m * m) - m;
    }
    void updateEr() {
        er = r / r.norm();
    }
    void updateEp() {
        ep = p3 / p3.norm();
    }
    void updateVwithP3T() {
        v = p3 / (T + m);
    }
    void updateP3TwithV() {
        // TODO
        Value vnorm = v.norm();
        if (vnorm < 1e-4) {
            T = 0.5 * m * vnorm * vnorm;
            p3 = m * v;
        } else {
            Value gamma = 1.0 / std::sqrt((1 - vnorm) * (1 + vnorm));
            T = (gamma - 1) * m;
            p3 = gamma * m * v;
        }
        updateEp();
    }
    void setR(Vector3 t_r) {
        this->r = t_r;
        this->updateEr();
    }
    void setP(Vector3 t_p) {
        p3 = t_p;
        T = this->TFromP3();
        updateVwithP3T();
        updateEp();
    }
    void setV(Vector3 t_v) {
        v = t_v;
        updateP3TwithV();
        updateEp();
    }

    /**
     * Total cross section scattering with Target \p t.
     * Use current \p this->T as initial kinetic energy.
     */
    virtual Value totalCrossSection(const Target& t) const = 0;

    /**
     * Sample recoil kinetic energy of the target particle in laboratory frame.
     */
    virtual Value scatteringSampleTr(const Target& t,
                                     RandomNumber<Value>& rn) = 0;

    /**
     * Return cosine of scattering angle of DM particle in laboratory frame.
     */
    virtual Value scatteringCosTheta(const Target& t, Value Tr) = 0;

    /**
     * Generate a Event from the current state.
     * @param kmsec if true the quantities \p t and \p r will be in unit of
     * second and kilomiter respecively, otherwise remain natural unit with
     * GeV = 1. Default \p true.
     */
    Event<Vector3, Value> toEvent(bool kmsec=true) {
        if (kmsec) {
            return Event<Vector3, Value> {t/unit::sec, T, r/unit::km, p3};
        }
        return Event<Vector3, Value> {t, T, r, p3};
    }
};


/**
 * Base abstract Earth class.
 */
template<typename Vector3, typename Value>
class Earth
{
public:
    std::vector<Target> targets;
    Earth() {};
    Earth(const Earth&) = delete;
    Earth& operator=(const Earth&) = delete;
    Earth& operator=(const Earth&&) = delete;
    virtual ~Earth() {};

    virtual void setCache(const Particle<Vector3, Value>&) {}
    virtual const std::vector<Target>& getTargets() = 0;
    /**
     * Propagate a particle.
     * Sample a free path and propagate the particle and determine whether
     * the particle in Earth or not.
     * Set \p particle.in_earth to false if the simulation should stop.
     */
    virtual void propagate(Particle<Vector3, Value>& particle,
                           RandomNumber<Value>& rn) = 0;
    /**
     * Sample a target before scattering.
     */
    virtual Target sampleTarget(const Particle<Vector3, Value>& particle,
                                RandomNumber<Value>& rn) const = 0;
};

/**
 * Basic interface for Particle scattering in the Earth.
 *
 * Update \p T, \p p3, \p v, \p ep of the \p particle after a scattering.
 * Use \p Earth.sampleTarget, \p Particle.scatteringSampleTr and
 * \p Particle.scatteringCosTheta.
 */
template<typename Vector3, typename Value>
bool scatter(Particle<Vector3, Value>& p,
             Earth<Vector3, Value>& earth,
             RandomNumber<Value>& rn)
{
    if (p.in_earth) {
        Target target = earth.sampleTarget(p, rn);
        // select a direction perpendicular to incident momentum
        // TODO handle ep[2] = 0 && ep[0] + ep[1] == 0
        Vector3 en = {p.ep[2], p.ep[2], -(p.ep[0] + p.ep[1])};
        en.normalize();

        // TODO sample a vector point to theta dierction
        Value Tr = p.scatteringSampleTr(target, rn);
        Value costheta = p.scatteringCosTheta(target, Tr);
        if (costheta * costheta > 1 || isnanf(costheta)) {
            std::cerr << std::scientific;
            std::cerr.precision(20);
            std::cerr << "costheta = " << costheta << std::endl;
            std::cerr << " p.T = " << p.T << std::endl;
            std::cerr << " p.p3 = " << p.p3 << std::endl;
            std::cerr << " Tr = " << Tr << std::endl;
            exit(1);
        }
        Vector3 n = costheta * p.ep + std::sqrt(1 - costheta*costheta) * en;

        // update T; set after costheta sampled
        p.T = p.T - Tr;
        // update 3 momentum
        p.p3 = n * p.pFromT();

        // give a random azimuthal angle 0 ~ 2pi
        Value phi = rn.uniform_phi();
        rotation_axis_angle(p.p3, p.ep, phi);

        p.updateEp();
        p.updateVwithP3T();
        return true;
    }
    return false;
}

/**
 * Basic interface for Particle propagating the in Earth.
 *
 * Simply wrap Earth.propagate.
 */
template<typename Vector3, typename Value>
inline bool propagate(Particle<Vector3, Value>& p,
                      Earth<Vector3, Value>& earth,
                      RandomNumber<Value>& rn)
{
    if (p.in_earth) {
        earth.propagate(p, rn);
        return true;
    }
    return false;
}

} // namespace darkprop
