#include <array>
#include <atomic>
#include <Eigen/Geometry>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <mutex>
#include <random>
#include <string>
#include <thread>
#include <time.h> // We just want to print the timestamp
#include <unistd.h>

#include "loadtxt.h"

using double3 = Eigen::Vector3d;
using Generator = std::default_random_engine;

// Our units are {kiloparsec, solar mass, gigayear}
constexpr double G = 4.498317481097514e-06;

void integrate(const std::array<double,6> y0, const double t_max, const double stride_size, double y[]);

std::string git_version()
{
    std::string result;
    std::ifstream git_log(".git/logs/HEAD");
    if (git_log) {
        std::string prev_line, line;
        getline(git_log, prev_line);
        while (getline(git_log, line)) {
            prev_line = line;
        }
        if (prev_line.length() > 51) result = prev_line.substr(41, 10);
        else result = "UNKNOWN";
    } else result = "UNKNOWN";
    return result;
}

std::string iso_time(std::chrono::system_clock::time_point time)
{
    auto time_c = std::chrono::system_clock::to_time_t(time);
    auto timeinfo = gmtime(&time_c);
    char time_string[256];
    strftime(time_string,256,"%FT%TZ", timeinfo);
    return std::string(time_string);
}

class Acceptor {
// This is a base class for functors that accept or decline the final state of the orbit
public:
    virtual bool operator()(const double w[6]) const = 0;
    virtual void output_information(std::ostream& stream) const = 0;
};

class In_box_cyl : public Acceptor {
public:
    In_box_cyl(const double pos_tol, const double vel_tol, const double3 target_pos, const double3 target_vel)
        : pos_tol(pos_tol), vel_tol(vel_tol)
    {
        target_d  = sqrt(target_pos[0]*target_pos[0] + target_pos[1]*target_pos[1]);
        target_z  = target_pos[2];
        target_vd = (target_pos[0]*target_vel[0] + target_pos[1]*target_vel[1])/target_d;
        target_vt = (target_pos[0]*target_vel[1] - target_pos[1]*target_vel[0])/target_d;
        target_vz = target_vel[2];
    }
    bool operator()(const double w[6]) const override
    {
        double d  = sqrt(w[0]*w[0] + w[1]*w[1]);
        double z  = w[2];
        double vd = (w[0]*w[3] + w[1]*w[4])/d;
        double vt = (w[0]*w[4] - w[1]*w[3])/d;
        double vz = w[5];
        if (((z<0)&&(target_z>0)) || ((z>0)&&(target_z<0))) { // Flip up-down
            z  = -z;
            vz = -vz;
        }
        if (((vt<0)&&(target_vt>0)) || ((vt>0)&&(target_vt<0))) { // Flip left-right
            vt = -vt;
        }
        return    (abs(d  - target_d)  < pos_tol)
               && (abs(z  - target_z)  < pos_tol)
               && (abs(vd - target_vd) < vel_tol)
               && (abs(vt - target_vt) < vel_tol)
               && (abs(vz - target_vz) < vel_tol);
    }
    void output_information(std::ostream& stream) const override
    {
        stream.setf(std::ios::scientific);
        stream << "# Acceptor type: box in cylindrical coordinates\n";
        stream << "# Acceptor parameters:\n";
        stream << "#   pos_tol   = " << std::setw(13) << pos_tol   << '\n';
        stream << "#   vel_tol   = " << std::setw(13) << vel_tol   << '\n';
        stream << "#   target_d  = " << std::setw(13) << target_d  << '\n';
        stream << "#   target_z  = " << std::setw(13) << target_z  << '\n';
        stream << "#   target_vd = " << std::setw(13) << target_vd << '\n';
        stream << "#   target_vt = " << std::setw(13) << target_vt << '\n';
        stream << "#   target_vz = " << std::setw(13) << target_vz << '\n';
    }
private:
    double pos_tol, vel_tol;
    double target_d, target_z, target_vd, target_vt, target_vz;
};

class Accept_all : public Acceptor {
public:
    Accept_all() = default;
    bool operator()(const double w[6]) const override
    {
        return true;
    }
    void output_information(std::ostream& stream) const override
    {
        stream << "# Acceptor type: accept all\n";
    }
};

////////////////////////////////////////////////////////////////////////////////////////////////////////////
class Match_data : public Acceptor {
public:
    Match_data() :
        data(Loadtxt("Peter_7D.txt.input", {2, 3, 4, 5, 6, 7}, 3).get_cols()),
        n(data[0].size()),
        tolerance(0.1)
    {
        constexpr double kms_to_kpcGyr = 1.0226911647958985;
        for (unsigned i = 0; i < n; i++) {
            data[3][i] *= kms_to_kpcGyr;
            data[4][i] *= kms_to_kpcGyr;
            data[5][i] *= kms_to_kpcGyr;
        }
    }
    bool operator()(const double w[6]) const override
    {
        for (unsigned i = 0; i < n; i++) {
            for (unsigned k = 0; k < 6; k++) {
                if (abs((w[k] - data[k][i])/w[k]) > 0.1) break;
                return true;
            }
        }
        return false;
    }
    void output_information(std::ostream& stream) const override
    {
        stream << "# Acceptor type: match data\n";
        stream << "#   tolerance = " << std::setw(13) << tolerance << '\n';
    }
private:
    std::vector<std::vector<double>> data;
    unsigned n;
    double tolerance;
};

class Initial_conditions {
public:
    virtual std::array<double,6> operator()(Generator& generator) const = 0;
    virtual void output_information(std::ostream& stream) const = 0;
};

class Uniform_sphere_escape_velocity_cutoff : public Initial_conditions {
public:
    Uniform_sphere_escape_velocity_cutoff(const double r_max, const double rho_0, const double b)
        : r_max(r_max),
          rho_0(rho_0),
          b(b),
          uniform_distribution(std::uniform_real_distribution<double>(0, 1))
    {}
    std::array<double,6> operator()(Generator& generator) const override
    {
        std::array<double,6> result;

        double r         = pow(uniform_distribution(generator), 1./3.)*r_max;
        double cos_theta = uniform_distribution(generator)*2-1;
        double sin_theta = sqrt(1 - cos_theta*cos_theta);
        double phi       = uniform_distribution(generator)*2*M_PI;
        result[0] = r*cos(phi)*sin_theta;
        result[1] = r*sin(phi)*sin_theta;
        result[2] = r*cos_theta;

        double potential = nfw_potential(r);
        if (potential > 0) throw std::runtime_error("Cannot deal with a positive potential");
        double v_max = sqrt(-2*potential);
        double v_magnitude;
        do {
            for (int i=3; i<6; i++) result[i] = uniform_distribution(generator);
            v_magnitude = sqrt(result[3]*result[3]+result[4]*result[4]+result[5]*result[5]);
        } while (v_magnitude >= v_max);

        return result;
    }
    void output_information(std::ostream& stream) const override
    {
        stream.setf(std::ios::scientific);
        stream << "# Initial conditions type: uniform sphere with escape velocity cutoff\n";
        stream << "# Initial conditions parameters:\n";
        stream << "#   r_max = " << std::setw(13) << r_max   << '\n';
    }
private:
    inline double nfw_potential(double r) const
    {
        return -4*M_PI*G*rho_0*b*b*b*log1p(r/b)/r;
    }
    double r_max, rho_0, b;
    mutable std::uniform_real_distribution<double> uniform_distribution;
};

class Box_initial_conditions : public Initial_conditions {
public:
    Box_initial_conditions(double x_max, double v_max)
        : x_max(x_max), v_max(v_max)
    {
        position_distribution = std::uniform_real_distribution<double>(-x_max, x_max);
        velocity_distribution = std::uniform_real_distribution<double>(-v_max, v_max);
    }
    std::array<double,6> operator()(Generator& generator) const override
    {
        std::array<double,6> result;
        for (unsigned i=0; i<3; i++) result[i] = position_distribution(generator);
        for (unsigned i=3; i<6; i++) result[i] = velocity_distribution(generator);
        return result;
    }
    void output_information(std::ostream& stream) const override
    {
        stream.setf(std::ios::scientific);
        stream << "# Initial conditions type: uniform box\n";
        stream << "# Initial conditions parameters:\n";
        stream << "#   x_max = " << std::setw(13) << x_max   << '\n';
        stream << "#   v_max = " << std::setw(13) << v_max   << '\n';
    }
private:
    double x_max, v_max;
    mutable std::uniform_real_distribution<double> position_distribution, velocity_distribution;
};

class Scattering_experiment {
public:
    Scattering_experiment(const std::string& file_name, const double t_max, const unsigned long long n_experiments, const Initial_conditions& initial_conditions, const Acceptor& accept)
        : n_experiments(n_experiments),
          t_max(t_max),
          accept(accept),
          initial_conditions(initial_conditions),
          start_time(std::chrono::system_clock::now()),
          file(std::ofstream(file_name)),
          n_threads(std::thread::hardware_concurrency())
    {
        const unsigned long time_count = start_time.time_since_epoch().count();
        const unsigned long big_prime = 840580612873L;
        unique_id = time_count + big_prime*getpid(); // TODO hash it
        file.setf(std::ios::scientific | std::ios::right);
        file.precision(16);
    }
    ~Scattering_experiment() {
        file.close();
    }
    void set_concurrency(int n_threads)
    {
        this->n_threads = n_threads;
    }
    void write_header()
    {
        file << "#############################################\n";
        file << "# Cosmological replay scattering experiment #\n";
        file << "#############################################\n";
        file << "# Time: " << iso_time(start_time) << '\n';
        file << "# Last Git commit: " << git_version() << '\n';
        char hostname[256];
        gethostname(hostname, 256);
        file << "# Hostname: " << hostname << '\n';
        file << "# PID: " << getpid() << '\n';
        file << "# Unique identifier: " << unique_id << '\n';
        file << "# Concurrency: " << n_threads << '\n';
        initial_conditions.output_information(file);
        accept.output_information(file);
        file.flush();
        // TODO information about galaxy model
    }
    void thread_task(int tid)
    {
        auto thread_begin_time = std::chrono::system_clock::now();
        auto report_interval = std::chrono::seconds(report_interval_seconds);
        auto next_report_time = thread_begin_time + report_interval;
        std::default_random_engine generator(unique_id + tid);

        while (counter < n_experiments) {
            for (int i=0; i<bunch_size; i++) {
                auto ic = initial_conditions(generator);
                double y[12];

                integrate(ic, t_max, t_max, y);

                if (accept(y+6)) {
                    std::lock_guard<std::mutex> lock(mtx);
                    for (int j=0; j<12; j++) file << std::setw(24) << y[j];
                    file << std::endl;
                }
            }
            counter += bunch_size;
            if (tid == 0) {
                auto now = std::chrono::system_clock::now();
                if (now >= next_report_time) {
                    std::lock_guard<std::mutex> lock(mtx);
                    file << "# " << iso_time(now) << " counter = " << counter << std::endl;
                    next_report_time = now + report_interval;
                }
            }
        }
    }
    void launch()
    {
        counter = 0;
        for (int tid=0; tid<n_threads; tid++) {
            threads.push_back(std::thread(&Scattering_experiment::thread_task, this, tid));
        }
        for (auto& t : threads) {
            t.join();
        }
        std::cout << counter << '\n';
    }
    int n_threads = 1;
    int bunch_size = 256;
    int report_interval_seconds = 300;
    std::chrono::system_clock::time_point start_time;
    std::ofstream file;
    std::mutex mtx;
    std::vector<std::thread> threads;
    std::atomic<unsigned long long> counter;
    unsigned long long n_experiments;
    unsigned long unique_id;
    double t_max;
    const Acceptor& accept;
    const Initial_conditions& initial_conditions;
};

int main()
{
    std::cout << "Hi\n";

    const double t_max = 11.65551390; // Gyr

    Uniform_sphere_escape_velocity_cutoff uniform_sphere_escape_velocity_cutoff(75, 1.82777387E+07, 9.86242602E+00);

    auto zzzzz = Match_data();

    Scattering_experiment scattering_experiment("results.dat", t_max, 1024*1024, uniform_sphere_escape_velocity_cutoff, std::ref(zzzzz));
    scattering_experiment.set_concurrency(80);
    scattering_experiment.write_header();
    scattering_experiment.launch();

    std::cout << "Bye\n";

    return 0;
}