//
// Created by dominik on 8/14/24.
//
#include <iostream>
#include <algorithm>
#include <unordered_map>
#include <map>
#include <string>
#include <vector>
#include <cmath>
#include <chrono>

#include "motions/base.h"
#include "times/base.h"
#include "functions.h"
#include "ranges.h"
#include "sims.h"
#include "io.h"


void run_spectrum(std::unordered_map<std::string, double>& parameter, Motion& motion, Distribution& dist) {
    const int num_walker = static_cast<int>(parameter["num_walker"]);

    // time axis for all time signals
    const int num_acq = static_cast<int>(parameter["num_acq"]);
    const std::vector<double> t_fid = arange(num_acq, parameter["dwell_time"]);
    const std::vector<double> echo_times = linspace(parameter["techo_start"], parameter["techo_stop"], static_cast<int>(parameter["techo_steps"]));

    // make timesignal vectors and set them to zero
    std::map<double, std::vector<double>> fid_dict;
    for (auto t_evo_i: echo_times) {
        fid_dict[t_evo_i] = std::vector<double>(num_acq);
        std::fill(fid_dict[t_evo_i].begin(), fid_dict[t_evo_i].end(), 0.);
    }

    // calculate min length of a trajectory
    const double tmax = *std::max_element(echo_times.begin(), echo_times.end()) * 2 + t_fid.back();

    // set parameter in distribution and motion model
    const double tau = parameter["tau"];
    dist.setTau(tau);
    motion.setDelta(parameter["delta"]);
    motion.setEta(parameter["eta"]);

    const auto start = std::chrono::system_clock::now();
    auto last_print_out = std::chrono::system_clock::now();
    const time_t start_time = std::chrono::system_clock::to_time_t(start);
    std::cout << "Start tau = " <<  tau << "s : " << ctime(&start_time);

    // let the walker walk
    for (int mol_i = 0; mol_i < num_walker; mol_i++){
        std::vector<double> traj_time{};
        std::vector<double> traj_phase{};

        make_trajectory(motion, dist, tmax, traj_time, traj_phase);

        for (auto& [t_echo_j, fid_j] : fid_dict) {
            // get phase at echo pulse
            int current_pos = nearest_index(traj_time, t_echo_j, 0);
            const double phase_techo = lerp(traj_time, traj_phase, t_echo_j, current_pos);


            for (int acq_idx = 0; acq_idx < num_acq; acq_idx++) {
                const double real_time = t_fid[acq_idx] + 2 * t_echo_j;

                current_pos = nearest_index(traj_time, real_time, current_pos);
                const double phase_acq = lerp(traj_time, traj_phase, real_time, current_pos);

                fid_j[acq_idx] += std::cos(phase_acq - 2 * phase_techo) / num_walker;
            }
            last_print_out = printSteps(last_print_out, start, num_walker, mol_i);
        }
    }

    // write fid to files
    fid_write_out("fid", t_fid, fid_dict, tau);

    const auto end = std::chrono::system_clock::now();

    std::chrono::duration<float> duration = end - start;
    const time_t end_time = std::chrono::system_clock::to_time_t(end);
    std::cout << "End   tau = " << tau << "s : " << ctime(&end_time);
    std::cout << "Duration: " << duration.count() << "s\n" << std::endl;

}


void run_ste(std::unordered_map<std::string, double>& parameter, Motion& motion, Distribution& dist) {
    const int num_walker = static_cast<int>(parameter[std::string("num_walker")]);

    const int num_mix_times = static_cast<int>(parameter[std::string("tmix_steps")]);
    const std::vector<double> evolution_times = linspace(parameter["tevo_start"], parameter["tevo_stop"], static_cast<int>(parameter["tevo_steps"]));
    const std::vector<double> mixing_times = logspace(parameter["tmix_start"], parameter["tmix_stop"], num_mix_times);


    // make ste decay vectors and set them to zero
    std::map<double, std::vector<double>> cc_dict;
    std::map<double, std::vector<double>> ss_dict;
    for (auto t_evo_i: evolution_times) {
        cc_dict[t_evo_i] = std::vector<double>(num_mix_times);
        ss_dict[t_evo_i] = std::vector<double>(num_mix_times);
        std::fill(ss_dict[t_evo_i].begin(), ss_dict[t_evo_i].end(), 0.);
    }

    // each trajectory must have a duration of at least tmax
    const double tmax = *std::max_element(evolution_times.begin(), evolution_times.end()) * 2 + *std::max_element(mixing_times.begin(), mixing_times.end());

    // set parameter in distribution and motion model
    const double tau = parameter["tau"];
    dist.setTau(tau);
    motion.setDelta(parameter["delta"]);
    motion.setEta(parameter["eta"]);

    const auto start = std::chrono::system_clock::now();
    const time_t start_time = std::chrono::system_clock::to_time_t(start);
    std::cout << "Start tau = " << tau << "s : " << ctime(&start_time);

    // let the walker walk
    for (int mol_i = 0; mol_i < num_walker; mol_i++){
        std::vector<double> traj_time{};
        std::vector<double> traj_phase{};

        make_trajectory(motion, dist, tmax, traj_time, traj_phase);

        for (auto& [t_evo_j, _] : cc_dict) {
            auto& cc_j = cc_dict[t_evo_j];
            auto& ss_j = ss_dict[t_evo_j];

            // get phase at beginning of mixing time
            int current_pos = nearest_index(traj_time, t_evo_j, 0);
            const double dephased = lerp(traj_time, traj_phase, t_evo_j, current_pos);
            const double cc_tevo = std::cos(dephased);
            const double ss_tevo = std::sin(dephased);

            for (int mix_idx = 0; mix_idx < num_mix_times; mix_idx++) {

                // get phase at end of mixing time
                const double time_end_mix = mixing_times[mix_idx] + t_evo_j;
                current_pos = nearest_index(traj_time, time_end_mix, current_pos);
                const double phase_mix_end = lerp(traj_time, traj_phase, time_end_mix, current_pos);

                // get phase at echo position
                const double time_echo = mixing_times[mix_idx] + 2 * t_evo_j;
                current_pos = nearest_index(traj_time, time_echo, current_pos);
                const double rephased = lerp(traj_time, traj_phase, time_echo, current_pos) - phase_mix_end;

                cc_j[mix_idx] += cc_tevo * std::cos(rephased) / num_walker;
                ss_j[mix_idx] += ss_tevo * std::sin(rephased) / num_walker;
            }
        }
    }

    // write to files
    fid_write_out("coscos", mixing_times, cc_dict, tau);
    fid_write_out("sinsin", mixing_times, ss_dict, tau);

    const auto end = std::chrono::system_clock::now();

    const std::chrono::duration<float> duration = end - start;
    const time_t end_time = std::chrono::system_clock::to_time_t(end);
    std::cout << "End   tau = " << tau << "s : " << ctime(&end_time);
    std::cout << "Duration: " << duration.count() << "s\n" << std::endl;

}


void make_trajectory(Motion& motion, const Distribution& dist, const double t_max, std::vector<double>& out_time, std::vector<double>& out_phase) {
    // Starting position
    double t_passed = 0;
    double phase = 0;

    motion.initialize();

    out_time.emplace_back(t_passed);
    out_phase.emplace_back(0);

    while (t_passed < t_max) {
        const double t = dist.tau_wait();
        t_passed += t;
        phase += motion.jump() * t;

        out_time.emplace_back(t_passed);
        out_phase.emplace_back(phase);
    }
}