#include <algorithm>
#include <chrono>
#include <filesystem>
#include <iostream>
#include <cmath>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <vector>
#include <string>

#include "main.hpp"
#include "yuc_config.h"

typedef double Value;
typedef float TStore;
typedef Vector3d<Value> Vector3;

int main(int argc, char *argv[])
{
    if (argc != 2) {
        std::cout << "Usage: mpiexec -n N " << argv[0]
            << " <config.toml>" << std::endl;
        exit(1);
    }

    int comm_sz;
    int my_rank;

    MPI_Init(NULL, NULL);
    MPI_Comm_size(MPI_COMM_WORLD, &comm_sz);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
    {

    yuc::config config;
    config.parse_toml(argv[1]);

    const std::string id = config["id"];

    Value Tmin {0}, Tmax {0};
    // [input]
    const auto input = config["input"];
    const std::string input_type = input["type"];
    std::vector<Value> mchis = input["mchis"];
    std::vector<Value> sigmas = input["sigmas"];
    for (auto& s : sigmas) { s *= unit::cm*unit::cm; }
    const std::string mchi_file = input["mchi_file"].or_get(std::string(""));
    const std::string sigma_file = input["sigma_file"].or_get(std::string(""));
    std::ifstream ifs;
    if (std::filesystem::exists(mchi_file)) {
        ifs.open(mchi_file);
        // overwirte mchis
        Value tmp;
        mchis.clear();
        while (ifs >> tmp) { mchis.push_back(tmp); }
        ifs.close();
    }
    if (std::filesystem::exists(sigma_file)) {
        ifs.open(sigma_file);
        // overwirte mchis
        Value tmp;
        sigmas.clear();
        while (ifs >> tmp) { sigmas.push_back(tmp * unit::cm * unit::cm); }
        ifs.close();
    }
    Value Tcut = input["Tcut"].or_get(0.0);

    // [input.halo]
    const auto input_halo = input["halo"];
    std::vector<Value> vec_earth = input_halo["v_earth"].or_get({0., 0., 0.});
    const Value kmsec = unit::km / unit::sec;
    const Vector3 v_earth(vec_earth[0] * kmsec,
                          vec_earth[1] * kmsec,
                          vec_earth[2] * kmsec);
    const Value halo_v_esc = input_halo["v_esc"].num(544) * kmsec;
    const Value halo_v_0 = input_halo["v_0"].num(220) * kmsec;
    const Value halo_v_cut = input_halo["v_cut"].num(1e-5) * kmsec;
    const Value halo_v_min = input_halo["v_min"].num(1e-5) * kmsec;
    Value halo_v_max = input_halo["v_max"].num(-1) * kmsec;
    halo_v_max = halo_v_max > 0 ? halo_v_max : halo_v_esc + v_earth.norm();
    
    // [input.flux]
    auto input_flux = input["flux"];
    const std::string flux_file = input_flux["flux_file"];
    std::vector<Value> energy;
    std::vector<std::vector<Value>> fluxes(mchis.size()*sigmas.size());
    if (input_type == "flux") {
        Value e_tmp;
        Value flux_tmp;
        // TODO check the file
        ifs.open(flux_file);
        while (ifs >> e_tmp) {
            energy.push_back(e_tmp);
            for (std::size_t i = 0; i != fluxes.size(); ++i) {
                ifs >> flux_tmp;
                fluxes[i].push_back(flux_tmp);
            }
        }
        ifs.close();

        Tmin = input_flux["Tmin"].or_get(0.0);
        Tmax = input_flux["Tmax"].or_get(0.0);
        if (Tmin <= 0.0) {
            Tmin = *min_element(energy.begin(), energy.end());
        }
        if (Tmax <= 0.0) {
            Tmax = *max_element(energy.begin(), energy.end());
        }
        if (Tcut <= 0.0) {
            Tcut = Tmin;
        }
    }

    // [output]
    const auto output = config["output"];
    const std::size_t sample_size = output["sample_size"];
    const std::string outdir = output["outdir"];
    const std::string file_name_style
        = output["file_name_style"].or_get(std::string("long"));
    std::size_t rings = output["rings"].or_get(1L);
    rings = rings < 1 ? 1 : rings;
    const std::string ring_space = output["ring_space"].or_get(std::string("linear"));
    const std::size_t chunk_size = output["chunk_size"];
    const std::vector<Value> r_axis = output["axis"].or_get({1.0, 0.0, 0.0});
    const bool store_time = output["store_time"].or_get(false);

    Vector3 ring_axis;
    if (input_type == "halo") {
        ring_axis = {v_earth[0], v_earth[1], v_earth[2]};
    } else {
        ring_axis = {r_axis[0], r_axis[1], r_axis[2]};
    }
    const std::size_t maximal_simulation = output["maximal_simulation"];

    const std::size_t full_tracks = output["full_tracks"].or_get(0L);

    // [output.log]
    auto output_log = output["log"];
    std::size_t screen_log_lag = output_log["screen_log_lag"].or_get(10000L);

    // [numerical]
    auto numerical = config["numerical"];
    int rn_seed = numerical["random_seed"].or_get(-1L);
    rn_seed = rn_seed >= 0 ? rn_seed + my_rank : rn_seed;
    const Value interp_Tmin = numerical["interp_Tmin"].or_get(1e-100);
    const std::size_t interp_num = numerical["interp_num"].or_get(10000L);

    // initialization
    // depth
    std::vector<Value> depth = output["depth"];

    // depth group name
    std::stringstream ss;
    std::vector<std::string> depth_groups;
    for (std::size_t di = 0; di != depth.size(); ++di) {
        ss << "depth_" << di;
        depth_groups.push_back(ss.str());
        depth[di] *= unit::km;
        ss.str("");
    }

    // ring group name
    std::vector<std::string> ring_groups;
    for (std::size_t r = 0; r != rings; ++r) {
        ss << "ring_" << r;
        ring_groups.push_back(ss.str());
        ss.str("");
    }

    std::vector<Value> ring_bins;
    if (ring_space == "linear") {
        ring_bins = linspace<Value>(0.0, M_PI, rings + 1, true);
    } else if (ring_space == "cos_linear") {
        ring_bins = linspace<Value>(1.0, -1.0, rings + 1, true);
        for (auto& r : ring_bins) { r = std::acos(r); }
    } else {
        throw std::invalid_argument(
            std::string("ring_space value ")
            + ring_space
            + " unknown.\nsupported are \"linear\" and \"cos_linear\"");
    }

    // DM
    std::shared_ptr<Particle<Vector3, Value>> dmp;
    const std::string dm_model = input["dm_model"];
    DMModel dmmodel = parse_dm_model(dm_model);

    // Earth
    std::shared_ptr<Earth<Vector3, Value>> earthp;
    auto premp = std::make_shared<PREMEarth<Vector3, Value>>();
    const std::string earth_model = input["earth_model"];
    EarthModel earthmodel = parse_earth_model(earth_model);
    switch (earthmodel) {
        case EarthModel::homo:
            earthp = std::make_shared<HomoEarth<Vector3, Value>>();
            break;
        case EarthModel::prem:
            earthp = premp;
            break;
    }

    // injector
    std::shared_ptr<Injector<Vector3, Value>> injp;
    auto flux_injp = std::make_shared<FluxInjector<Vector3, Value>>(rn_seed);
    if (input_type == "halo") {
        auto halo_injp = std::make_shared<HaloInjector<Vector3, Value>>(
            halo_v_min, halo_v_max, v_earth, halo_v_esc, halo_v_0,
            constant::rEarth, 0.0, rn_seed);
        injp = halo_injp;
    } else if (input_type == "flux") {
        injp = flux_injp;
    } else {
        throw std::invalid_argument(
            std::string("input.type value ")
            + input_type
            + " unknown.\nsupported are \"halo\" and \"flux\""
        );
    }

    auto counts = make_triple_vector<std::size_t>(depth.size(), rings, 2);
    
    std::string dir = outdir + "/" + id;
    std::filesystem::create_directories(dir);
    std::string track_dir = dir + "/fulltracks";
    if (full_tracks > 0) { std::filesystem::create_directories(track_dir); }

    auto start_time = std::chrono::high_resolution_clock::now();
    for (std::size_t mi = 0; mi != mchis.size(); mi++) {
        Value mchi = mchis[mi];
        for (std::size_t si = 0; si != sigmas.size(); ++si) {
            Value sigma = sigmas[si];
            if (input_type == "halo") {
                Tmax = 0.5 * mchi * halo_v_max * halo_v_max;
                Tmin = 0.5 * mchi * halo_v_min * halo_v_min;
                Tcut = 0.5 * mchi * halo_v_cut * halo_v_cut;
            } else if (input_type == "flux") {
                flux_injp->init(energy, fluxes[si + mi * sigmas.size()]);
                flux_injp->m_Tmin = Tmin;
                flux_injp->m_Tmax = Tmax;
            }

            switch (dmmodel) {
                case DMModel::si:
                    dmp = std::make_shared<SIDM<Vector3, Value>>(sigma, mchi);
                    break;
                case DMModel::sihelm:
                    auto elasticdmp =
                        std::make_shared<SIHelmDM<Vector3, Value>>(sigma, mchi);
                    for (auto& t : earthp->getTargets()) {
                        elasticdmp->initCrossSectionCDF(interp_Tmin, Tmax,
                                                        t, interp_num);
                    }
                    if (earthmodel == EarthModel::prem) {
                        premp->cross_section_energy_dependent = true;
                    }
                    dmp = elasticdmp;
                    break;
            }

            auto filename = make_filename(dir, id, mchi, sigma);
            if (file_name_style == "short") {
                filename = make_filename(dir, id, mi, si);
            }
            if (my_rank == 0) {
                std::cout << "Collectively generating file: " << filename
                    << " ...." << std::endl;
            }
            hid_t fapl = H5Pcreate(H5P_FILE_ACCESS);
            H5Pset_fapl_mpio(fapl, MPI_COMM_WORLD, MPI_INFO_NULL);
            auto fd = H5Fcreate(filename.c_str(), H5F_ACC_TRUNC,
                                H5P_DEFAULT, fapl);
            H5Pclose(fapl);

            hid_t ds = make_dset<TStore>(fd, "sigma");
            write_dset(ds, sigma / unit::cm / unit::cm);
            write_attr(ds, "unit", "cm^2");
            H5Dclose(ds);

            ds = make_dset<TStore>(fd, "mchi");
            write_dset(ds, mchi);
            write_attr(ds, "unit", "cm^2");
            H5Dclose(ds);

            ds = make_dset<TStore>(fd, "Tcut");
            write_dset(ds, mchi);
            write_attr(ds, "unit", "GeV");
            H5Dclose(ds);
            ds = make_dset<TStore>(fd, "R0");
            write_dset(ds, constant::rEarth / unit::km);
            write_attr(ds, "unit", "km");
            H5Dclose(ds);
            if (input_type == "halo") {
                ds = make_dset<TStore>(fd, "Vesc");
                write_dset(ds, halo_v_esc / kmsec);
                write_attr(ds, "unit", "km/s");
                H5Dclose(ds);
                ds = make_dset<TStore>(fd, "V0");
                write_dset(ds, halo_v_0 / kmsec);
                write_attr(ds, "unit", "km/s");
                H5Dclose(ds);
                ds = make_dset<TStore>(fd, "Vearth", {vec_earth.size()});
                write_dset(ds, vec_earth, {0}, {vec_earth.size()});
                write_attr(ds, "unit", "km/s");
                H5Dclose(ds);
                ds = make_dset<TStore>(fd, "Vmin");
                write_dset(ds, halo_v_min / kmsec);
                write_attr(ds, "unit", "km/s");
                H5Dclose(ds);
                ds = make_dset<TStore>(fd, "Vmax");
                write_dset(ds, halo_v_max / kmsec);
                write_attr(ds, "unit", "km/s");
                H5Dclose(ds);
                ds = make_dset<TStore>(fd, "Vcut");
                write_dset(ds, halo_v_cut / kmsec);
                write_attr(ds, "unit", "km/s");
                H5Dclose(ds);
            } else if (input_type == "flux") {
                ds = make_dset<TStore>(fd, "norm");
                write_dset(ds, flux_injp->m_norm);
                write_attr(ds, "unit", "cm^-2 sr^-1");
                H5Dclose(ds);
            }
            auto dsets = make_dsets(
                fd, depth_groups, ring_groups, depth, ring_bins,
                static_cast<Value>(constant::rEarth), sample_size,
                chunk_size, store_time);
            if (my_rank == 0) {
                std::cout << "Collectively generating file: " << filename
                    << " done\n" << "Start simulation" << std::endl;
            }

            MPI_Barrier(MPI_COMM_WORLD);
            auto my_counts = run(dsets, dmp, earthp, injp, depth, depth_groups,
                                 ring_groups, Tcut, my_rank, comm_sz,
                                 sample_size, maximal_simulation, full_tracks,
                                 track_dir, id, ring_axis, ring_bins,
                                 store_time, rn_seed, dmmodel, chunk_size,
                                 screen_log_lag);
            for (std::size_t di = 0; di != depth.size(); ++di) {
                for (std::size_t ri = 0; ri != rings; ++ri) {
                    MPI_Barrier(MPI_COMM_WORLD);
                    MPI_Reduce(my_counts[di][ri].data(), counts[di][ri].data(),
                               2, MPI_UNSIGNED_LONG_LONG, MPI_SUM,
                               0, MPI_COMM_WORLD);
                    auto groupname = rings > 1 ?
                                     depth_groups[di] + "/" + ring_groups[ri]
                                     : depth_groups[di];
                    if (my_rank == 0) {
                        write_dset(dsets[groupname + "/Nsample"],
                                   counts[di][ri][0]);
                        write_dset(dsets[groupname + "/Nsim"],
                                   counts[di][ri][1]);
                    }
                }
            }
            close_dsets(dsets);
            H5Fclose(fd);
        }
    }
    auto end_time = std::chrono::high_resolution_clock::now();

    if (my_rank == 0) {
        std::cout << "costs "
            << std::chrono::duration_cast<std::chrono::seconds>(
                end_time - start_time).count() << " seconds" << std::endl;
    }

    }  // release resources before MPI_Finalize();
    MPI_Finalize();
    return 0;
}
