//
// 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 "functions.h"
#include "motions/base.h"
#include "times/base.h"
#include "ranges.h"
#include "sims.h"

#include <bits/fs_fwd.h>

#include "io.h"


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

    const std::vector<double> correlation_times = logspace(parameter["tau_start"], parameter["tau_stop"], static_cast<int>(parameter["tau_steps"]));
    motion.setDelta(parameter["delta"]);
    motion.setEta(parameter["eta"]);

    // time axis for all time signals
    const auto 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"]));

    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.);
    }

    const double tmax = *std::max_element(echo_times.begin(), echo_times.end()) * 2 + t_fid.back();

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

        dist.setTau(tau_i);

        for (auto& [_, fid_j]: fid_dict) {
            std::fill(fid_j.begin(), fid_j.end(), 0.);
        }

        // reset array for each correlation time
        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_tevo = lerp(traj_time, traj_phase, t_echo_j, current_pos);

                // time axis by echo delay to get time in trajectory

                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_tevo) / 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_i);

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

        std::chrono::duration<float> duration = end - start;
        time_t end_time = std::chrono::system_clock::to_time_t(end);
        std::cout << "End   tau = " << tau_i << "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_acq = static_cast<int>(parameter[std::string("num_acq")]);
    const int num_walker = static_cast<int>(parameter[std::string("num_walker")]);

    const double tau = parameter["tau"];
    dist.setTau(tau);

    motion.setDelta(parameter["delta"]);
    motion.setEta(parameter["eta"]);

    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 = linspace(parameter["tevo_start"], parameter["tevo_stop"], static_cast<int>(parameter["tevo_steps"]));

    std::map<double, std::vector<double>> fid_dict;
    for (auto t_evo_i: evolution_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.);
    }

    // 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());

    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);

    for (auto& [_, fid_j]: fid_dict) {
        std::fill(fid_j.begin(), fid_j.end(), 0.);
    }

    // reset array for each correlation time
    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, fid_j] : fid_dict) {
            int current_pos = nearest_index(traj_time, t_evo_j, 0);
            const double phase_tevo = lerp(traj_time, traj_phase, t_evo_j, current_pos);

            // time axis by echo delay to get time in trajectory

            for (int acq_idx = 0; acq_idx < num_acq; acq_idx++) {
                const double real_time = mixing_times[acq_idx] + t_evo_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] += cos(phase_acq - 2 * phase_tevo);
            }
        }
    }

    // write fid to files
    fid_write_out("ste", mixing_times, fid_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);
    }
}