/**
 * @file
 * Initial condition of dark matter.
 */
#pragma once

#include <stdexcept>
#include <cmath>
#include <algorithm>
#include "Const.hpp"
#include "Base.hpp"

namespace darkprop {

/**
 * The normalization of Maxwell velocity distribution (vesc = +inf).
 */
template<typename Value=decltype(constant::v0)>
inline Value norm_maxwell(Value v0=constant::v0) {
    return M_PI * std::sqrt(M_PI) * v0 * v0 * v0;
}

/**
 * The normalization of halo DM velocity distribution with a finite escape
 * velocity.
 */
template<typename Value=decltype(constant::vesc)>
inline Value norm_esc(Value vesc=constant::vesc, Value v0=constant::v0) {
    return norm_maxwell(v0) * (std::erf(vesc / v0)
                -M_2_SQRTPI * vesc / v0 * std::exp(-vesc * vesc / v0 / v0));
}

/**
 * Halo DM 1D speed distribution.
 */
template<typename Value>
Value fv_halo_1d(Value v, Value vesc=constant::vesc, Value v0=constant::v0) {
    return 4.0 * M_PI / norm_esc(vesc, v0) * v*v * std::exp(-v * v / v0 / v0);
}

/**
 * Sample 1D halo speed from the isotropic SHM.
 */
template<typename Value>
Value sample_halo_speed(RandomNumber<Value>& rn, /**< random number gen */
                        Value vmin=0.0, /**< lower bound */
                        Value vmax=constant::vesc, /**< upper bound */
                        Value vesc=constant::vesc, /**< escape speed */
                        Value v0=constant::v0 /**< most probable speed */)
{
    // TODO use inverse CDF sampling
	Value v, fv;
    do {
        v = rn.uniform_ab(vmin, vmax);
        fv = rn.uniform_xi() * fv_halo_1d(v0, vesc, v0);
    } while (fv > fv_halo_1d(v, vesc, v0));
	return v;
}

/**
 * Sample a random vector isotropicly with length \p len (default 1).
 */
template<typename Vector3, typename Value>
Vector3 random_isotropic_vector(RandomNumber<Value>& rn, Value len=1)
{
    Value theta = std::acos(rn.uniform_costheta());
    Value phi = rn.uniform_phi();
    return Vector3(len * std::sin(theta) * std::cos(phi),
                   len * std::sin(theta) * std::sin(phi),
                   len * std::cos(theta));
}

/**
 * Sample 3D halo velocity in the Earth frame from the SHM.
 *
 * The sampling range can be set to (\p vmin, \p vmax). If \p vmax <= 0, it
 * will be set to \p v_earth.norm() + \p vesc.
 */
template<typename Vector3, typename Value>
Vector3 sample_halo_velocity_earth_frame(
        Vector3 v_earth,           /**< earth velocity */
        RandomNumber<Value>& rn,   /**< random number generator */
        Value vmin=0.0,            /**< speed lower bound */
        Value vmax=-1.0,           /**< speed upper bound */
        Value vesc=constant::vesc, /**< halo escape speed */
        Value v0=constant::v0      /**< halo most probable speed */)
{
    Vector3 velocity;
    if (vmax <= 0) {
        vmax = v_earth.norm() + vesc;
    }
    Value halo_vmin = std::max(vmin - v_earth.norm(), static_cast<Value>(0.0));
    Value halo_vmax = std::min(vmax + v_earth.norm(), vesc);
    do {
        Value vi = sample_halo_speed(rn, halo_vmin, halo_vmax, vesc, v0);
        velocity = random_isotropic_vector<Vector3>(rn, vi) - v_earth;
    } while (velocity.norm() < vmin || velocity.norm() > vmax);
    return velocity;
}

/**
 * Sample an incident position according to the given direction.
 */
template<typename Vector3, typename Value>
Vector3 random_incident_position(RandomNumber<Value>& rn,
                                 const Vector3& vi,
                                 Value r_earth=constant::rEarth,
                                 Value R=1.1*constant::rEarth)
{
    Vector3 ez = -vi.normalized();
    Value n = std::sqrt(ez[1] * ez[1] + ez[2] * ez[2]);
    Vector3 ex(0, ez[2] / n, -ez[1] / n);
    Vector3 ey = ez.cross(ex);
	Value phi = rn.uniform_phi();
    Value xi = rn.uniform_xi();
    return R * ez + std::sqrt(xi) * r_earth
                    * (std::cos(phi) * ex + std::sin(phi) * ey);
}

/**
 * Set halo DM random initial condition.
 *
 * \p vmin and \p vmax are in the Earth frame.
 */
template<typename Vector3, typename Value>
void init_halo(Particle<Vector3, Value>& p,
               Value t,                    /**< initial time */
               Value vmin,                 /**< speed lower bound */
               const Vector3& v_earth,     /**< velocity of the Earth */
               RandomNumber<Value>& rn,    /**< random number generator */
               Value R,                    /**< initail disk distance */
               Value vmax=-1.0,            /**< speed upper bound */
               Value vesc=constant::vesc,  /**< halo escape velocity */
               Value v0=constant::v0,       /**< halo most probable speed */
               Value radius=constant::rEarth)
{
    Vector3 vinit = sample_halo_velocity_earth_frame(v_earth, rn, vmin,
                                                     vmax, vesc, v0);
    Vector3 xinit = random_incident_position(rn, vinit, radius, R);
    p.t = t;
    p.setV(vinit);
    p.setR(xinit);
}


template<typename Value>
inline void check_non_positive(Value T)
{
    if (T <= 0) { throw std::invalid_argument {"parameter <= 0"}; }
}

/**
 * Set initial condition of a particle with energy T and direction in galactic
 * coordinate (b, l).
 */
template<typename Vector3, typename Value>
void init_Tbl(Particle<Vector3, Value>& p,
              Value T, Value b, Value l,
              RandomNumber<Value>& rn,
              Value t=0.0,
              Value R=1.1*constant::rEarth,
              Value radius=constant::rEarth)
{
    check_non_positive(T);
    Value pi = std::sqrt(T * (T + 2. * p.m));
    Vector3 ez = Vector3(std::cos(b) * std::cos(l),
                         std::cos(b) * std::sin(l),
                         std::sin(b));
    Vector3 pinit = -pi * ez;
    Vector3 xinit = random_incident_position(rn, pinit, radius, R);
    p.t = t;
    p.setP(pinit);
    p.setR(xinit);
}

template<typename Vector3, typename Value>
void init_Tbl_vertical(Particle<Vector3, Value>& p,
                       Value T, Value b, Value l,
                       Value t=0.0,
                       Value R=constant::rEarth)
{
    check_non_positive(T);
    Value pi = std::sqrt(T * (T + 2. * p.m));
    Vector3 er = Vector3(std::cos(b) * std::cos(l),
                         std::cos(b) * std::sin(l),
                         std::sin(b));
    p.setP(-pi * er);
    p.setR(R * er);
    p.t = t;
}

/**
 * Set initial condition of a particle isotropically with energy T.
 */
template<typename Vector3, typename Value>
void init_T_isotropic(Particle<Vector3, Value>& p,
                      Value T,
                      RandomNumber<Value>& rn,
                      Value t=0.0,
                      Value R=1.1*constant::rEarth,
                      Value radius=constant::rEarth)
{
    check_non_positive(T);
    Value pi = std::sqrt(T * (T + 2. * p.m));
    Vector3 pinit = random_isotropic_vector<Vector3>(rn, pi);
    Vector3 xinit = random_incident_position(rn, pinit, radius, R);
    p.t = t;
    p.setP(pinit);
    p.updateVwithP3T();
    p.setR(xinit);
}

/**
 * Place the particle on the the Earth's surface.
 */
template<typename Vector3, typename Value>
void inject(Particle<Vector3, Value>& p, Value radius=constant::rEarth)
{
    // already in the Earth
    if (p.r.norm() <= radius) {
        p.in_earth = true;
        return;
    }

    Value rep = p.r.dot(p.ep);
    if (rep >= 0) {
        throw std::runtime_error("Particle did not point to the Earth.");
    }

    Value delta = rep * rep - p.r.dot(p.r) + radius * radius;

    if (delta <= 0) { throw std::runtime_error("Particle missed the Earth."); }

    p.r += (-rep - std::sqrt(delta)) * p.ep;

    if (p.r.norm() > radius) {
        Value x = std::sqrt(delta);
        p.r += (x > 1.0 * unit::mm ? 1.0 * unit::mm * p.ep : x * p.ep);
    }
    if (p.r.norm() > radius) {
        std::cerr << "Warning!! particle injection failed." << std::endl;
        p.updateEr();
        p.in_earth = false;
        return;
    }

    p.updateEr();
    p.in_earth = true;
}

} // namespace darkprop
