/**
 *
 * @file plugins/CriticalPath/DrawStats.cpp
 *
 * @copyright 2008-2024 Bordeaux INP, CNRS (LaBRI UMR 5800), Inria,
 *                      Univ. Bordeaux. All rights reserved.
 *
 * @author Camille Ordronneau
 * @author Johnny Jazeix
 * @author Mathieu Faverge
 * @author Mohamed Faycal Boullit
 *
 * @date 2024-07-17
 */
#include <string>
#include <iostream>

#include "common/Session.hpp"

#include "trace/EntityTypes.hpp"
#include "trace/Entitys.hpp"
#include "trace/Container.hpp"
#include "trace/EntityValue.hpp"
#include "trace/StateChange.hpp"
#include "trace/tree/BinaryTree.hpp"
#include "trace/tree/Node.hpp"
#include "trace/Trace.hpp"

#include "DrawStats.hpp"

int cmp = 0;

size_t is_in_critical_path(std::vector<int> criticalPath, int jobId) {
    for (size_t i = 0; i < criticalPath.size(); i++) {
        if (criticalPath[i] == jobId)
            return i;
    }
    return -1;
}

void get_critical_path_states_rec(Node<StateChange> *node, const std::vector<int> &criticalPath, std::vector<const State *> &critical_path_states) {
    int job_id, index;
    if (!node)
        return;
    const State *st = node->get_element()->get_left_state();
    if (st) {
        const std::map<std::string, Value *> *extra_fields = st->get_extra_fields();
        if (extra_fields != nullptr && !extra_fields->empty()) {
            for (const auto &extra_field: *extra_fields) {
                if (!extra_field.first.compare("JobId")) {
                    job_id = stoi(extra_field.second->to_string());
                    index = is_in_critical_path(criticalPath, job_id);
                    if (index != -1) {
                        critical_path_states[index] = st;
                    }
                }
            }
        }
    }
    get_critical_path_states_rec(node->get_left_child(), criticalPath, critical_path_states);

    get_critical_path_states_rec(node->get_right_child(), criticalPath, critical_path_states);
}

void get_critical_path_states(Trace *trace, const std::vector<int> &critical_path, std::vector<const State *> &critical_path_states) {
    std::list<Container *> container_list;
    trace->get_all_containers(container_list);

    for (const auto &it: container_list) {
        BinaryTree<StateChange> *states;
        states = (it)->get_states();
        Node<StateChange> *node = states->get_root();
        get_critical_path_states_rec(node, critical_path, critical_path_states);
    }
}

double get_cpu_count(Trace *trace) {
    double cpu_count = 0;
    std::list<Container *> container_list;
    trace->get_all_containers(container_list);

    for (const auto &it: container_list) {
        size_t found_cpu = it->get_name().find("CPU");
        size_t found_proc = it->get_name().find("Proc");
        size_t found_thread = it->get_name().find("Thread");
        size_t found_user_thread = it->get_name().find("UserThread");
        if (found_cpu != std::string::npos || found_proc != std::string::npos || (found_user_thread == std::string::npos && found_thread != std::string::npos)) {
            cpu_count++;
        }
    }
    return cpu_count;
}

int next_state(std::vector<const State *> &critical_path_states, int current) {
    while (current >= 0 && critical_path_states[current] == nullptr) {
        current--;
    }
    return current;
}

void link_critical_path(Trace *trace, std::vector<const State *> &critical_path_states) {
    // Preparing arguments for functions startlink/endlink
    std::map<std::string, Value *> empty;
    Container *ancestor = trace->search_container(String("program"));
    // Getting the ContainerType of src/dest Container
    ContainerType *cpu_type = const_cast<ContainerType *>(trace->search_container(String("CPU0"))->get_type());
    // A new link type for critical path
    if (trace->search_link_type(String("Critical Path Link")) == NULL) {
        Name cp_name("Critical Path Link");
        trace->define_link_type(cp_name, const_cast<ContainerType *>(ancestor->get_type()), cpu_type, cpu_type, empty);
    }
    LinkType *cp_type = trace->search_link_type(String("Critical Path Link"));
    EntityType *cp_entity_type = trace->search_entity_type(String("WSL"));
    EntityValue *cp_entity_value = trace->search_entity_value(String("CPL"), cp_entity_type, true);
    cp_entity_value->set_used_color(Color(0.96, 0.85, 0.43));

    if (ancestor == nullptr || cp_type == nullptr) {
        std::cout << "Error" << std::endl;
        return;
    }

    // Linking states process
    int current = next_state(critical_path_states, critical_path_states.size() - 1);
    int nb_link = 0;
    const State *current_state = critical_path_states[current];
    Date time = current_state->get_end_time();
    std::string key = "Critical link" + std::to_string(nb_link);
    ancestor->start_link(time, cp_type, const_cast<Container *>(current_state->get_container()), cp_entity_value, key, empty);
    current = next_state(critical_path_states, current - 1);
    while (current >= 0) {
        current_state = critical_path_states[current];
        time = current_state->get_start_time();
        ancestor->end_link(time, const_cast<Container *>(current_state->get_container()), key, empty);
        nb_link++;
        key = "Critical link" + std::to_string(nb_link);
        time = current_state->get_end_time();
        ancestor->start_link(time, cp_type, const_cast<Container *>(current_state->get_container()), cp_entity_value, key, empty);
        current = next_state(critical_path_states, current - 1);
    }
}
