#include <ios>
#include <map>
#include <iomanip>
#include <limits>
#include <stdexcept>
#include <algorithm>
#include <mpi.h>
#include "darkprop/core.hpp"
#include "darkprop/IO.hpp"
#include "Interpolator.hpp"
#include "Integrator.hpp"

using namespace darkprop;

enum class EarthModel { homo, prem };
enum class DMModel { si, sihelm };

inline std::string str_tolower(std::string s) {
    std::transform(s.begin(), s.end(), s.begin(),
                   [](unsigned char c) { return std::tolower(c); });
    return s;
}

inline DMModel parse_dm_model(const std::string& dm_model) {
    std::string name = str_tolower(dm_model);
    if (name == "si" || name =="simple") { return DMModel::si; }
    if (name == "sihelm" || name =="elastic") { return DMModel::sihelm; }
    throw std::invalid_argument(std::string("dm model ") + dm_model
                    + " unknown.\nsupported are \"si\" and \"sihelm\"");
}

inline EarthModel parse_earth_model(const std::string& earth_model) {
    std::string name = str_tolower(earth_model);
    if (name == "homo" || name == "simple") { return EarthModel::homo; }
    if (name == "prem") { return EarthModel::prem; }
    throw std::invalid_argument(std::string("earth model ")
        + earth_model
        + " unknown.\nsupported are "
        + "\"homo\" and \"prem\"");
}

template<typename Vector3, typename Value>
class Injector {
public:
    Value m_rEarth;
    Value m_t0;
    RandomNumber<Value> m_rn;
    virtual ~Injector() {}
    virtual Value next(Particle<Vector3, Value>&) = 0;
    virtual Value operator()(Particle<Vector3, Value>& p) {
        return next(p);
    }
    void setSeed(int t_seed) { m_rn = RandomNumber<Value>(t_seed); }
    void sett0(Value t_0) { m_t0 = t_0; }
};

template<typename Vector3, typename Value>
class HaloInjector : public Injector<Vector3, Value>
{
public:
    Value m_vmin;
    Value m_vmax;
    Vector3 m_vearth;
    Value m_vesc;
    Value m_v0;
    HaloInjector(Value t_vmin=0,
                 Value t_vmax=-1,
                 Vector3 t_vearth={0, 0, 240*unit::km/unit::sec},
                 Value t_vesc=constant::vesc,
                 Value t_v0=constant::v0,
                 Value t_rEarth=constant::rEarth,
                 Value t_t0=0,
                 int t_seed=-1)
        : m_vmin {t_vmin}, m_vmax {t_vmax}, m_vearth {t_vearth},
          m_vesc {t_vesc}, m_v0 {t_v0}
        {
            this->m_rEarth = t_rEarth;
            this->m_t0 = t_t0;
            this->m_rn = RandomNumber<Value>(t_seed);
        }
    Value next(Particle<Vector3, Value>& p) override {
        init_halo(p, this->m_t0, m_vmin, m_vearth, this->m_rn,
                  1.1*this->m_rEarth, m_vmax, m_vesc, m_v0, this->m_rEarth);
        inject(p, this->m_rEarth);
        return p.v.norm();  // TODO actual weight
    }
};

template<typename Vector3, typename Value>
class FluxInjector : public Injector<Vector3, Value>
{
public:
    Value m_Tmin;
    Value m_Tmax;
    Value m_norm;  ///> [cm^-2 sr^-1]
    Interpolator<Value> spectrum;
    FluxInjector(int seed=-1, Value t_rEarth=constant::rEarth) {
        this->setSeed(seed);
        this->m_rEarth = t_rEarth;
    }
    FluxInjector(const std::vector<Value>& t_T,
                 const std::vector<Value>& t_flux,
                 Value t_rEarth=constant::rEarth,
                 int seed=-1) {
        this->m_rEarth = t_rEarth;
        this->setSeed(seed);
        init(t_T, t_flux);
    }
    void init(const std::vector<Value>& t_T,
              const std::vector<Value>& t_flux) {
        m_Tmin = *std::min_element(t_T.begin(), t_T.end());
        m_Tmax = *std::max_element(t_T.begin(), t_T.end());
        spectrum.interpolate(t_T, t_flux);
        Integrator<Interpolator<Value>> integrator(true);

        m_norm = integrator(spectrum, m_Tmin, m_Tmax);
    }
    ~FluxInjector() {}
    Value next(Particle<Vector3, Value>& p) override {
        Value T = std::exp(this->m_rn.uniform_ab(std::log(m_Tmin),
                                                 std::log(m_Tmax)));
        init_T_isotropic(p, T, this->m_rn, this->m_t0, 1.1*this->m_rEarth);
        inject(p, this->m_rEarth);
        return spectrum(T) * T * std::log(m_Tmax / m_Tmin) / m_norm;
    }
};

template<typename T>
std::vector<std::vector<std::vector<T>>> make_triple_vector(
        std::size_t dim1, std::size_t dim2, std::size_t dim3=0)
{
    std::vector<std::vector<std::vector<T>>> result(dim1);
    for (auto& vd : result) {
        vd = std::vector<std::vector<T>>(dim2);
        if (dim3 > 0) {
            for (auto& v : vd) {
                v = std::vector<T>(dim3);
            }
        }
    }
    return result;
}

template<typename Vector3, typename Value, typename TStore>
std::size_t ring_classify(
        const std::vector<Event<Vector3, Value>>& events,
        TStore weight,
        const std::vector<Value>& ring_bins,
        const Vector3& axis,
        std::vector<std::vector<Event<Vector3, Value>>>& ring_events,
        std::vector<std::vector<TStore>>& ring_weights,
        const std::vector<std::size_t>& sample_nums,
        std::size_t sample_sz)
{
    std::size_t sample_num = 0;
    for (std::size_t i = 0; i != events.size(); ++i) {
        std::size_t ring = std::distance(
            ring_bins.begin() + 1,
            std::lower_bound(ring_bins.begin() + 1, ring_bins.end(),
                std::acos(events[i].r.dot(axis) / events[i].r.norm())));
        if (sample_nums[ring] + ring_events[ring].size() < sample_sz) {
            ring_events[ring].push_back(events[i]);
            ring_weights[ring].push_back(weight);
            ++sample_num;
        }
    }
    return sample_num;
}

inline
bool continue_run(const std::vector<std::vector<std::size_t>>& sample_nums,
                  std::size_t sample_sz) {
    for (auto& d1 : sample_nums) {
        for (auto& size : d1) { if (size < sample_sz) return true; }
    }
    return false;
}

template<typename Value>
std::string make_filename(const std::string& dir, const std::string& id,
                          Value mchi, Value sigma,
                          const std::string& what="event",
                          std::size_t track=0)
{
    std::stringstream ss;
    ss.setf(std::ios_base::scientific);
    ss.precision(3);
    if (what == "event") {
        ss << dir << "/" << id << "_mchi" << mchi
            << "GeV_sigma" << sigma / unit::cm / unit::cm << "cm2.hdf5";
    } else if (what == "track") {
        ss << dir << "/" << id << "_mchi" << mchi
            << "GeV_sigma" << sigma / unit::cm / unit::cm
            << "cm2_fulltrack_" << track << ".csv";
    }
    return ss.str();
}

template<> inline
std::string make_filename<std::size_t>(const std::string& dir,
                                       const std::string& id,
                                       std::size_t mi, std::size_t si,
                                       const std::string& what,
                                       std::size_t track)
{
    std::stringstream ss;
    if (what == "event") {
        ss << dir << "/" << id << "_m" << mi << "_s" << si << ".hdf5";
    } else if (what == "track") {
        ss << dir << "/" << id << "_m" << mi << "_s" << si
            << "_fulltrack_" << track << ".csv";
    }
    return ss.str();
}


inline std::size_t my_nums(int my_rank, int comm_sz, std::size_t total)
{
    std::size_t base_num {total / static_cast<std::size_t>(comm_sz)};
    std::size_t remaind {total % static_cast<std::size_t>(comm_sz)};
    return base_num + (static_cast<std::size_t>(my_rank) < remaind ? 1 : 0);
}

inline std::size_t my_offset(int my_rank, int comm_sz, std::size_t total)
{
    std::size_t base_num {total / static_cast<std::size_t>(comm_sz)};
    std::size_t remaind {total % static_cast<std::size_t>(comm_sz)};
    return static_cast<std::size_t>(my_rank) < remaind ?
        my_rank * (base_num + 1) : remaind + my_rank * base_num;
}

inline std::size_t get_width(std::size_t size)
{
    std::size_t w {2};
    while (static_cast<std::size_t>(std::pow(10, w - 1)) / (size + 1) == 0) {
        ++w;
    }
    return w;
}

template<typename TStore=float, typename Vector3, typename Value>
std::vector<std::vector<std::vector<std::size_t>>> run(
    std::map<std::string, hid_t> dsets,
    std::shared_ptr<Particle<Vector3, Value>> dmp,
    std::shared_ptr<Earth<Vector3, Value>> earthp,
    std::shared_ptr<Injector<Vector3, Value>> injp,
    std::vector<Value> depth,
    std::vector<std::string> depth_groups,
    std::vector<std::string> ring_groups,
    Value Tcut,
    int t_rank,
    int t_comm_sz,
    std::size_t total_samples,
    std::size_t maximal_simulation,
    std::size_t full_tracks,
    const std::string& track_dir,
    const std::string& id,
    Vector3 axis,
    const std::vector<Value>& ring_bins,
    bool store_time,
    int seed,
    DMModel dmmodel,
    std::size_t chunk_size,
    std::size_t loglag)
{
    auto start_time = std::chrono::high_resolution_clock::now();

    RandomNumber<Value> rn(seed);

    const std::size_t tracks = my_nums(t_rank, t_comm_sz, full_tracks);
    const std::size_t track_offset = my_offset(t_rank, t_comm_sz, full_tracks);
    std::size_t track = 0;
    while (track < tracks) {
        std::string track_f = make_filename(
            track_dir, id, dmp->m, dmp->sigma0, "track", track + track_offset);
        injp->next(*dmp);
        dump_track_csv(simulate_track(*dmp, *earthp, Tcut, rn), track_f);
        ++track;
    }

    std::size_t rings = ring_bins.size() - 1;
    auto cross_events = make_triple_vector<Event<Vector3, Value>>(depth.size(),
                                                                  rings);
    auto weights = make_triple_vector<TStore>(depth.size(), rings);
    
    std::vector<std::vector<std::size_t>> sample_nums(depth.size());
    for (auto& v : sample_nums) { v = std::vector<std::size_t>(rings, 0); }
    // {Nsample, Nsim}
    auto counts = make_triple_vector<std::size_t>(depth.size(), rings, 2);

    const std::size_t sample_sz = my_nums(t_rank, t_comm_sz, total_samples);
    const std::size_t baseoffset = my_offset(t_rank, t_comm_sz, total_samples);
    const std::size_t max_runnum = my_nums(t_rank, t_comm_sz,
                                           maximal_simulation);

    std::size_t runnum = 0;
    std::size_t total_sample_num = 0;
    bool keep_going = true;

    do {
        auto weight = injp->next(*dmp);

        if (dmmodel == DMModel::si) { earthp->setCache(*dmp); }

        auto events = simulate_cross_depth(*dmp, *earthp, Tcut, depth, rn);
        ++runnum;

        for (std::size_t di = 0; di < depth.size(); ++di) {
            total_sample_num += ring_classify(events[di],
                                              static_cast<TStore>(weight),
                                              ring_bins,
                                              axis.normalized(),
                                              cross_events[di],
                                              weights[di],
                                              sample_nums[di],
                                              sample_sz);
            for (std::size_t ri = 0; ri != rings; ++ri) {
                // set chunk size for balance between memery usage and disk IO
                std::size_t events_on_ring = cross_events[di][ri].size();
                if (sample_nums[di][ri] == sample_sz) { continue; }
                if (sample_nums[di][ri] > sample_sz) {
                    throw std::runtime_error("sample_nums > sample_sz");
                }
                if (sample_nums[di][ri] + events_on_ring == sample_sz
                    || events_on_ring > chunk_size || runnum == max_runnum) {
                    auto groupname = rings > 1 ?
                                     depth_groups[di] + "/" + ring_groups[ri]
                                     : depth_groups[di];
                    std::size_t offset = baseoffset + sample_nums[di][ri];
                    h5dump_events_galactic<TStore>(dsets, groupname,
                        cross_events[di][ri], offset, store_time);
                    write_dset(dsets[groupname + "/weight"], weights[di][ri],
                               {offset}, {weights[di][ri].size()});

                    sample_nums[di][ri] += events_on_ring;

                    if (sample_nums[di][ri] == sample_sz
                        || runnum == max_runnum) {
                        counts[di][ri][0] = sample_nums[di][ri];
                        counts[di][ri][1] = runnum;
                    }

                    cross_events[di][ri].clear();
                    weights[di][ri].clear();
                }
            }
        }

        keep_going = continue_run(sample_nums, sample_sz)
                     && runnum < max_runnum;
        if (runnum % loglag == 0 || !keep_going) {
            std::cout << "\rprocess:"
                << std::setw(get_width(t_comm_sz)) << t_rank
                << " run # " << std::setw(get_width(max_runnum)) << runnum
                << " sample # "
                << std::setw(get_width(sample_sz)) << total_sample_num
                << std::scientific << std::setprecision(3)
                << " mass = " << dmp->m
                << " sigma = " << dmp->sigma0 / unit::cm / unit::cm
                << std::fixed << std::setprecision(1)
                << " " << total_sample_num * 100.0
                          / (sample_sz * rings * depth.size())
                << "%" << std::flush;
        }
    } while (keep_going);
    auto end_time = std::chrono::high_resolution_clock::now();
    std::cout << " " << std::chrono::duration_cast<std::chrono::seconds>(
                end_time - start_time).count() << " seconds" << std::endl;
    return counts;
}
