#include "Frame_Analysis.hpp"
#include <algorithm>
#include <cmath>
#include <limits>
#include <numeric>
#include <stdexcept>
#include <unordered_map>

namespace {

std::uint64_t subtract_saturated(std::uint64_t value, std::uint64_t origin) {
    return value >= origin ? value - origin : 0;
}

double percentile(std::vector<std::uint64_t> values, double probability) {
    if (values.empty())
        return 0.0;
    std::sort(values.begin(), values.end());
    const double position = probability * static_cast<double>(values.size() - 1);
    const auto lower = static_cast<std::size_t>(std::floor(position));
    const auto upper = static_cast<std::size_t>(std::ceil(position));
    const double fraction = position - static_cast<double>(lower);
    return static_cast<double>(values[lower]) * (1.0 - fraction) +
           static_cast<double>(values[upper]) * fraction;
}

struct Node_Samples {
    std::vector<std::uint64_t> durations;
    std::vector<std::uint64_t> waits;
    std::size_t critical_count{};
};

Node_Statistics statistics_for(Render_Node_Id id, const Node_Samples& samples) {
    Node_Statistics result;
    result.node_id = id;
    result.execution_count = samples.durations.size();
    if (samples.durations.empty())
        return result;
    const auto duration_sum = std::accumulate(samples.durations.begin(), samples.durations.end(),
                                              std::uint64_t{});
    const auto wait_sum = std::accumulate(samples.waits.begin(), samples.waits.end(),
                                          std::uint64_t{});
    result.average_ns = static_cast<double>(duration_sum) / samples.durations.size();
    const std::size_t moving_count = std::min<std::size_t>(20, samples.durations.size());
    const auto moving_begin = samples.durations.end() - static_cast<std::ptrdiff_t>(moving_count);
    result.moving_average_ns = static_cast<double>(
        std::accumulate(moving_begin, samples.durations.end(), std::uint64_t{})) / moving_count;
    result.p50_ns = percentile(samples.durations, 0.50);
    result.p95_ns = percentile(samples.durations, 0.95);
    result.p99_ns = percentile(samples.durations, 0.99);
    result.minimum_ns = *std::min_element(samples.durations.begin(), samples.durations.end());
    result.maximum_ns = *std::max_element(samples.durations.begin(), samples.durations.end());
    result.critical_path_frequency = samples.critical_count;
    result.average_scheduler_wait_ns = samples.waits.empty()
        ? 0.0
        : static_cast<double>(wait_sum) / samples.waits.size();
    return result;
}

}

Frame_Analysis analyze_frame(const Render_Plan& plan, const Frame_Snapshot& frame) {
    if (frame.render_plan_version != plan.version)
        throw std::invalid_argument("frame and render plan versions differ");
    if (frame.node_executions.size() != plan.graph.nodes.size())
        throw std::invalid_argument("frame execution slot count differs from render plan");

    Frame_Analysis result;
    result.frame_id = frame.frame_id;
    result.render_plan_version = frame.render_plan_version;
    result.total_render_duration_ns = frame.render_duration_ns();
    result.nodes.resize(plan.graph.nodes.size());
    std::unordered_map<Render_Node_Id, std::size_t> indices;
    indices.reserve(plan.graph.nodes.size());
    for (const auto& node : plan.graph.nodes) {
        indices.emplace(node.node_id, node.execution_index);
        const auto& execution = frame.node_executions.at(node.execution_index);
        if (execution.node_id != node.node_id)
            throw std::invalid_argument("frame execution slot has a different node id");
        auto& analysis = result.nodes[node.execution_index];
        analysis.node_id = node.node_id;
        analysis.duration_ns = execution.duration_ns();
        analysis.start_offset_ns = subtract_saturated(execution.start_time_ns,
                                                      frame.render_start_ns);
        analysis.end_offset_ns = subtract_saturated(execution.end_time_ns,
                                                    frame.render_start_ns);
        result.total_work_duration_ns += analysis.duration_ns;
    }

    std::vector<std::vector<std::size_t>> predecessors(plan.graph.nodes.size());
    std::vector<std::vector<std::size_t>> successors(plan.graph.nodes.size());
    std::vector<std::size_t> indegree(plan.graph.nodes.size());
    for (const auto& edge : plan.graph.edges) {
        const std::size_t from = indices.at(edge.from);
        const std::size_t to = indices.at(edge.to);
        predecessors[to].push_back(from);
        successors[from].push_back(to);
        ++indegree[to];
    }
    std::vector<std::size_t> topological;
    topological.reserve(plan.graph.nodes.size());
    for (std::size_t index = 0; index < indegree.size(); ++index) {
        if (indegree[index] == 0)
            topological.push_back(index);
    }
    for (std::size_t cursor = 0; cursor < topological.size(); ++cursor) {
        for (std::size_t successor : successors[topological[cursor]]) {
            if (--indegree[successor] == 0)
                topological.push_back(successor);
        }
    }
    if (topological.size() != plan.graph.nodes.size())
        throw std::invalid_argument("render plan contains a cycle");

    std::vector<std::uint64_t> path_duration(plan.graph.nodes.size());
    std::vector<std::size_t> path_parent(plan.graph.nodes.size(),
                                         std::numeric_limits<std::size_t>::max());
    for (std::size_t index : topological) {
        std::uint64_t dependency_ready = frame.render_start_ns;
        std::uint64_t parent_path{};
        for (std::size_t predecessor : predecessors[index]) {
            dependency_ready = std::max(dependency_ready,
                frame.node_executions[predecessor].end_time_ns);
            if (path_duration[predecessor] > parent_path) {
                parent_path = path_duration[predecessor];
                path_parent[index] = predecessor;
            }
        }
        auto& analysis = result.nodes[index];
        analysis.dependency_ready_time_ns = dependency_ready;
        analysis.scheduler_wait_ns = frame.node_executions[index].start_time_ns >= dependency_ready
            ? frame.node_executions[index].start_time_ns - dependency_ready
            : 0;
        path_duration[index] = parent_path + analysis.duration_ns;
    }
    if (!path_duration.empty()) {
        std::size_t index = static_cast<std::size_t>(std::distance(
            path_duration.begin(), std::max_element(path_duration.begin(), path_duration.end())));
        result.critical_path_duration_ns = path_duration[index];
        while (index != std::numeric_limits<std::size_t>::max()) {
            result.critical_path.push_back(plan.graph.nodes[index].node_id);
            result.nodes[index].on_critical_path = true;
            index = path_parent[index];
        }
        std::reverse(result.critical_path.begin(), result.critical_path.end());
    }
    for (auto& node : result.nodes) {
        node.work_contribution = result.total_work_duration_ns == 0
            ? 0.0
            : static_cast<double>(node.duration_ns) / result.total_work_duration_ns;
        node.critical_path_contribution = !node.on_critical_path ||
                                          result.critical_path_duration_ns == 0
            ? 0.0
            : static_cast<double>(node.duration_ns) / result.critical_path_duration_ns;
    }

    struct Event {
        std::uint64_t time{};
        int delta{};
    };
    std::vector<Event> events;
    std::unordered_map<std::uint32_t, std::uint64_t> worker_work;
    events.reserve(frame.node_executions.size() * 2);
    for (const auto& execution : frame.node_executions) {
        if (execution.duration_ns() == 0)
            continue;
        events.push_back({execution.start_time_ns, 1});
        events.push_back({execution.end_time_ns, -1});
        worker_work[execution.worker_id] += execution.duration_ns();
    }
    std::sort(events.begin(), events.end(), [](const Event& left, const Event& right) {
        return left.time < right.time || (left.time == right.time && left.delta < right.delta);
    });
    std::size_t active{};
    std::uint64_t busy_duration{};
    std::uint64_t previous_time{};
    for (const Event& event : events) {
        if (active != 0 && event.time >= previous_time)
            busy_duration += event.time - previous_time;
        if (event.delta < 0) {
            if (active != 0)
                --active;
        } else {
            ++active;
            result.peak_parallelism = std::max(result.peak_parallelism, active);
        }
        previous_time = event.time;
    }
    result.parallel_overlap_ns = result.total_work_duration_ns >= busy_duration
        ? result.total_work_duration_ns - busy_duration
        : 0;
    result.average_parallelism = result.total_render_duration_ns == 0
        ? 0.0
        : static_cast<double>(result.total_work_duration_ns) /
              result.total_render_duration_ns;
    result.workers.reserve(worker_work.size());
    for (const auto& [worker_id, work] : worker_work) {
        result.workers.push_back({worker_id, work,
            result.total_render_duration_ns == 0
                ? 0.0
                : static_cast<double>(work) / result.total_render_duration_ns});
    }
    std::sort(result.workers.begin(), result.workers.end(),
              [](const Worker_Utilization& left, const Worker_Utilization& right) {
                  return left.worker_id < right.worker_id;
              });
    std::vector<std::size_t> bottlenecks(result.nodes.size());
    std::iota(bottlenecks.begin(), bottlenecks.end(), 0);
    std::sort(bottlenecks.begin(), bottlenecks.end(), [&result](std::size_t left,
                                                               std::size_t right) {
        const auto left_cost = result.nodes[left].duration_ns + result.nodes[left].scheduler_wait_ns;
        const auto right_cost = result.nodes[right].duration_ns + result.nodes[right].scheduler_wait_ns;
        return left_cost > right_cost;
    });
    bottlenecks.resize(std::min<std::size_t>(5, bottlenecks.size()));
    for (std::size_t index : bottlenecks)
        result.bottleneck_nodes.push_back(result.nodes[index].node_id);
    return result;
}

std::vector<Node_Statistics> analyze_node_statistics(
    std::span<const Frame_Analysis> frames) {
    std::unordered_map<Render_Node_Id, Node_Samples> samples;
    for (const auto& frame : frames) {
        for (const auto& node : frame.nodes) {
            auto& values = samples[node.node_id];
            values.durations.push_back(node.duration_ns);
            values.waits.push_back(node.scheduler_wait_ns);
            values.critical_count += node.on_critical_path ? 1 : 0;
        }
    }
    std::vector<Node_Statistics> result;
    result.reserve(samples.size());
    for (const auto& [id, values] : samples)
        result.push_back(statistics_for(id, values));
    std::sort(result.begin(), result.end(), [](const Node_Statistics& left,
                                              const Node_Statistics& right) {
        return left.node_id < right.node_id;
    });
    return result;
}

std::vector<Plan_Version_Statistics> analyze_plan_versions(
    std::span<const Frame_Analysis> frames) {
    std::unordered_map<Render_Plan_Version, std::vector<Frame_Analysis>> grouped;
    for (const auto& frame : frames)
        grouped[frame.render_plan_version].push_back(frame);
    std::vector<Plan_Version_Statistics> result;
    result.reserve(grouped.size());
    for (auto& [version, group] : grouped) {
        Plan_Version_Statistics statistics;
        statistics.render_plan_version = version;
        statistics.frame_count = group.size();
        std::vector<std::uint64_t> render_durations;
        std::vector<std::uint64_t> waits;
        render_durations.reserve(group.size());
        double parallelism_sum{};
        for (const auto& frame : group) {
            render_durations.push_back(frame.total_render_duration_ns);
            parallelism_sum += frame.average_parallelism;
            statistics.peak_parallelism = std::max(statistics.peak_parallelism,
                                                   frame.peak_parallelism);
            for (const auto& node : frame.nodes)
                waits.push_back(node.scheduler_wait_ns);
        }
        statistics.render_average_ns = static_cast<double>(std::accumulate(
            render_durations.begin(), render_durations.end(), std::uint64_t{})) / group.size();
        statistics.render_p50_ns = percentile(render_durations, 0.50);
        statistics.render_p95_ns = percentile(render_durations, 0.95);
        statistics.render_maximum_ns = *std::max_element(render_durations.begin(),
                                                        render_durations.end());
        statistics.average_parallelism = parallelism_sum / group.size();
        statistics.scheduler_wait_average_ns = waits.empty()
            ? 0.0
            : static_cast<double>(std::accumulate(waits.begin(), waits.end(),
                                                  std::uint64_t{})) / waits.size();
        statistics.scheduler_wait_p95_ns = percentile(waits, 0.95);
        statistics.nodes = analyze_node_statistics(group);
        result.push_back(std::move(statistics));
    }
    std::sort(result.begin(), result.end(), [](const Plan_Version_Statistics& left,
                                              const Plan_Version_Statistics& right) {
        return left.render_plan_version < right.render_plan_version;
    });
    return result;
}
