#include "io.h"

#include <sstream>
#include <fstream>
#include <iostream>
#include <algorithm>
#include <complex>
#include <vector>
#include <iomanip>
#include <map>
#include <string>
#include <filesystem>


std::pair<std::string, double> get_optional_parameter(std::vector<std::string>::const_iterator &it) {
    std::string stripped_arg;
    if (it->size() > 2 && it->at(0) == '-' && it->at(1) == '-') {
        stripped_arg = it->substr(2, it->size());
    } else if (it->size() > 1 && it->at(0) == '-') {
        stripped_arg = it->substr(1, it->size());
    }
    std::transform(stripped_arg.begin(), stripped_arg.end(), stripped_arg.begin(), [](unsigned char c) { return std::tolower(c); });
    const auto stripped_value = std::stod(*++it, nullptr);

    return std::make_pair(stripped_arg, stripped_value);
}

/**
 * @brief Read and parse arguments coming from the command line
 * @param argc Number of command-line arguments
 * @param argv List of command-line arguments
 * @return Arguments
 */
Arguments parse_args(const int argc, char* argv[]) {
    if (argc < 3) {
        throw std::runtime_error("Not enough arguments: missing parameter file");
    }

    Arguments args;

    // convert to vector to use iterator for loop
    const std::vector<std::string> input_args(argv + 1, argv + argc);

    for (auto it = input_args.begin();  it != input_args.end(); ++it) {
        // check for optional parameter
        if (it->at(0) == '-') {

            // only --spectrum and --ste are parameter that are predefined
            if (*it == "--spectrum") {
                args.spectrum = true;
                continue;
            }

            if (*it == "--ste") {
                args.ste = true;
                continue;
            }

            // all other arguments are optional parameter
            auto [option_name, option_value] = get_optional_parameter(it);
            args.optional[option_name] = option_value;
            continue;
        }

        // Two positional parameters are defined: 1. Location of config file; 2. Name of motion model
        if (args.parameter_file.empty()) {
            args.parameter_file = *it;
            continue;
        }

        if (args.motion_type.empty()) {
            args.motion_type = *it;
            continue;
        }

        if (args.distribution_type.empty()) {
            args.distribution_type = *it;
            continue;
        }

        throw std::invalid_argument("too many positional arguments");
    }

    if (args.parameter_file.empty() || args.motion_type.empty() || args.distribution_type.empty()) {
        throw std::invalid_argument("Missing argument");
    }

    return args;
}


std::unordered_map<std::string, double> read_parameter(const std::filesystem::path& infile) {
    if (!exists(infile)) {
        std::cerr << "File " << infile << " does not exist" << std::endl;
        exit(1);
    }

    std::ifstream instream(infile);

    std::unordered_map<std::string, double> parameter;

    std::string line;
    std::string delim = "=";
    std::string key;
    std::string value;
    size_t delim_pos;

    while (std::getline(instream, line)) {
        // skip comment lines starting with #, and empty lines
        if (line[0] == '#' || line.length() == 1) continue;

        // strip spaces from line to have always key=value
        line.erase(std::remove(line.begin(), line.end(), ' '), line.end());

        // split at '=' character and add to map
        delim_pos = line.find('=');
        key = line.substr(0, delim_pos);
        value = line.substr(delim_pos+1);
        parameter[key] = std::stod(value);
    }

    return parameter;
}


std::string make_directory(const std::string& path_name) {
    if (!std::filesystem::create_directories(path_name)) {
        std::cout << "Created directory " << path_name << std::endl;
    }

    return path_name;
}


void save_parameter_to_file(
    const std::string& resulttype,
    const std::string& directory,
    const std::unordered_map<std::string, double>& parameter,
    const std::unordered_map<std::string, double>& optional
    ) {
    std::ostringstream parameter_filename;
    parameter_filename << directory << "/" << resulttype;

    parameter_filename << std::setprecision(6) << std::scientific;
    for (const auto& [key, value]: optional) {
        parameter_filename << "_" << key << "=" << value;
    }
    parameter_filename << "_parameter.txt";

    {
        // write data to file, columns are secondary axis (echo delay, evolution times)
        std::string datafile = parameter_filename.str();
        std::ofstream filestream(datafile, std::ios::out);

        for (const auto& [key, value]: parameter) {
            filestream << key << "=" << value << "\n";
        }
    }
}


void save_data_to_file(
    const std::string& resulttype,
    const std::string& directory,
    const std::vector<double>& x,
    const std::map<double, std::vector<double>>& y,
    const std::unordered_map<std::string, double>& optional
    ) {
    // make file name
    std::ostringstream datafile_name;
    datafile_name << directory << "/" << resulttype;

    datafile_name << std::setprecision(6) << std::scientific;
    for (const auto& [key, value]: optional) {
        datafile_name << "_" << key << "=" << value;
    }
    datafile_name << ".dat";

    {
        // write data to file, columns are secondary axis (echo delay, evolution times)
        std::string datafile = datafile_name.str();
        std::ofstream filestream(datafile, std::ios::out);

        // first line are values of secondary axis
        filestream << "#";
        for (const auto& [t_echo_j, _] : y) {
            filestream << t_echo_j << "\t";
        }
        filestream << std::endl;

        // write values to file
        auto size = x.size();
        for (unsigned int i = 0; i < size; i++) {
            filestream << x[i];
            for (const auto& [_, fid_j] : y) {
                filestream << "\t" << fid_j[i];
            }
            filestream << "\n";
        }
    }
}

void save_data_to_file(
    const std::string& resulttype,
    const std::string& directory,
    const std::vector<double>& x,
    const std::vector<double>& y,
    const std::unordered_map<std::string, double>& optional
    ) {
    // make file name
    std::ostringstream datafile_name;
    datafile_name << directory << "/" << resulttype;

    datafile_name << std::setprecision(6) << std::scientific;
    for (const auto& [key, value]: optional) {
        datafile_name << "_" << key << "=" << value;
    }
    datafile_name << ".dat";

    {
        // write data to file, columns are secondary axis (echo delay, evolution times)
        std::string datafile = datafile_name.str();
        std::ofstream filestream(datafile, std::ios::out);

        // write values to file
        auto size = x.size();
        for (unsigned int i = 0; i < size; i++) {
            filestream << x[i] << "\t" << y[i] << "\n";
        }
    }
}