from collections import deque
from typing import Callable

import networkx as nx

from torch._functorch._activation_checkpointing.graph_info_provider import (
    GraphInfoProvider,
)


class KnapsackEvaluator:
    """
    This class evaluates the theoretical runtime and peak memory usage of a given checkpointing strategy.
    It takes in a graph and a list of nodes that are saved and recomputed, and then simulates the
    backward pass to calculate the peak memory usage.
    """

    def __init__(
        self,
        graph_info_provider: GraphInfoProvider,
    ) -> None:
        self._graph_info_provider = graph_info_provider

    def _get_backward_memory_from_topologically_sorted_graph(
        self,
        node_graph: nx.DiGraph,
        node_memories: dict[str, float],
        saved_nodes_set: set[str],
        peak_memory_after_forward_pass: float,
    ) -> list[tuple[float, str]]:
        """
        Simulates the backward pass and keeps track of the peak memory usage.

        High Level Steps:
            1. Set Initial Peak/Current Memory
                Allows you to set the peak memory after the forward pass, but typically this is
                the sum of the estimated memory of the saved nodes.
            2. Perform a reverse topological sort of the node_graph.
                If full graph is defined then will sort the full graph and only process the subset
                of nodes in the node_graph.
            3. Iterate through the sorted graph nodes.
                If the node is saved then just drop it's memory from current memory.
                If the node is not saved then add it's memory to current memory and then traverse it's
                predecessors to simulate recomuptation chain. Will check if new peak memory after all
                predecessors are processed.

        Args:
            node_graph (nx.DiGraph): A directed graph representing the recomputable forward nodes.
            saved_nodes_set (Set[str]): A set of node names that are saved.
            peak_memory_after_forward_pass (float): The peak memory usage after the forward pass.
        """
        current_memory = [
            (peak_memory_after_forward_pass, "Initial Peak/Current Memory")
        ]
        already_computed = set()
        sorted_nodes = list(reversed(list(nx.topological_sort(node_graph))))
        dependencies_computed = set()

        for node in sorted_nodes:
            if node in saved_nodes_set or node in already_computed:
                current_memory.append(
                    (
                        current_memory[-1][0] - node_memories[node],
                        f"Dropping Node(already saved): {node}",
                    )
                )
                continue

            already_computed.add(node)
            current_memory.append(
                (
                    current_memory[-1][0] + node_memories[node],
                    f"Recomputing Node: {node}",
                )
            )
            # Create a queue of dependencies required for recomputation
            predecessor_queue = deque(
                [
                    dependency
                    for dependency, v in node_graph.in_edges(node)
                    if dependency not in already_computed
                ]
            )
            while predecessor_queue:
                dep = predecessor_queue.popleft()
                already_computed.add(dep)
                dependencies_computed.add(dep)
                current_memory.append(
                    (
                        current_memory[-1][0] + node_memories[dep],
                        f"Recomputing Predecessor of {node}: {dep}",
                    )
                )
                # Add predecessors of the predecessor to the queue if they haven't been recomputed yet
                for dependency_of_dependency, _ in node_graph.in_edges(dep):
                    if (
                        dependency_of_dependency in already_computed
                        or dependency_of_dependency in saved_nodes_set
                        or dependency_of_dependency in predecessor_queue
                    ):
                        continue
                    predecessor_queue.append(dependency_of_dependency)
            dependencies_computed.clear()
            current_memory.append(
                (current_memory[-1][0] - node_memories[node], f"Dropping Node: {node}")
            )
        return current_memory

    def _validate_all_indexes_accounted_for_in_provided_output(
        self, saved_nodes_idxs: list[int], recomputable_node_idxs: list[int]
    ) -> None:
        """
        Validate that all indexes are accounted for in the provided output.
        This function checks that the union of saved nodes and recomputable nodes
        covers all candidate nodes without any overlaps.
        """
        recomputable_node_idxs_set = set(recomputable_node_idxs)
        saved_nodes_idxs_set = set(saved_nodes_idxs)
        all_candidate_nodes_idxs = set(
            range(len(self._graph_info_provider.all_recomputable_banned_nodes))
        )
        # Check that there are no overlaps between saved nodes and recomputable nodes
        assert (
            len(recomputable_node_idxs_set.intersection(saved_nodes_idxs_set)) == 0
        ), "Saved nodes and recomputable nodes cannot have any overlaps"
        # Check that all candidate nodes are accounted for
        assert (
            recomputable_node_idxs_set.union(saved_nodes_idxs_set)
            == all_candidate_nodes_idxs
        ), "All candidate nodes must be accounted for in the provided output"

    def evaluate_knapsack_output(
        self,
        saved_nodes_idxs: list[int],
        recomputable_node_idxs: list[int],
        account_for_backward_pass: bool = False,
    ) -> dict[str, float]:
        """
        Evaluate the theoretical runtime and peak memory usage of a given checkpointing strategy.
        Args:
        - saved_nodes_idxs (List[int]): The indices of nodes that are saved.
        - recomputable_node_idxs (List[int]): The indices of nodes that need to be recomputed.
        """
        self._validate_all_indexes_accounted_for_in_provided_output(
            saved_nodes_idxs, recomputable_node_idxs
        )
        recomputation_runtime = sum(
            self._graph_info_provider.all_node_runtimes[
                self._graph_info_provider.all_recomputable_banned_nodes[node]
            ]
            for node in recomputable_node_idxs
        )
        if account_for_backward_pass:
            memory_list = self._get_backward_memory_from_topologically_sorted_graph(
                node_graph=self._graph_info_provider.recomputable_node_only_graph_with_larger_graph_context,
                saved_nodes_set={
                    self._graph_info_provider.all_recomputable_banned_nodes[i]
                    for i in saved_nodes_idxs
                },
                node_memories=self._graph_info_provider.all_node_memories,
                peak_memory_after_forward_pass=sum(
                    self._graph_info_provider.all_node_memories[
                        self._graph_info_provider.all_recomputable_banned_nodes[i]
                    ]
                    for i in saved_nodes_idxs
                ),
            )
            peak_memory = max(memory_list, key=lambda x: x[0])[0]
        else:
            peak_memory = sum(
                self._graph_info_provider.all_node_memories[
                    self._graph_info_provider.all_recomputable_banned_nodes[node]
                ]
                for node in saved_nodes_idxs
            )
        return {
            "peak_memory": peak_memory,
            "recomputation_runtime": recomputation_runtime,
            "non_ac_peak_memory": self._graph_info_provider.get_non_ac_peak_memory(),
            "theoretical_max_runtime": self._graph_info_provider.get_theoretical_max_runtime(),
            "percentage_of_theoretical_peak_memory": peak_memory
            / self._graph_info_provider.get_non_ac_peak_memory(),
            "percentage_of_theoretical_peak_runtime": recomputation_runtime
            / self._graph_info_provider.get_theoretical_max_runtime(),
        }

    def evaluate_distribution_of_results_for_knapsack_algo(
        self,
        knapsack_algo: Callable[
            [list[float], list[float], float], tuple[float, list[int], list[int]]
        ],
        memory_budget_values: list[float],
    ) -> list[dict[str, float]]:
        """
        Evaluates the distribution of results for a given knapsack algorithm.
        Args:
            knapsack_algo (Callable): The knapsack algorithm to use for evaluation.
            memory_budget_values (List[float]): A list of memory budgets to evaluate.
        """
        results = list()
        for memory_budget in memory_budget_values:
            _, saved_nodes, recomputed_nodes = knapsack_algo(
                self._graph_info_provider.get_knapsack_memory_input(),
                self._graph_info_provider.get_knapsack_runtime_input(),
                memory_budget,
            )
            result = self.evaluate_knapsack_output(
                saved_nodes_idxs=saved_nodes,
                recomputable_node_idxs=recomputed_nodes,
            )
            result["memory_budget"] = memory_budget
            results.append(result)
        return results

    def get_knee_point_memory_budget(
        self,
        knapsack_algo: Callable[
            [list[float], list[float], float], tuple[float, list[int], list[int]]
        ],
        max_mem_budget: float = 0.1,
        min_mem_budget: float = 0.001,
        iterations: int = 100,
    ) -> float:
        """
        Finds the memory budget at the knee point in the Pareto frontier.

        The knee point is defined as the point where the trade-off between
        runtime and memory usage is optimal.

        Args:
            knapsack_algo (callable): Knapsack algorithm to use for evaluation.
            max_mem_budget (float, optional): Maximum memory budget. Defaults to 0.1.
            min_mem_budget (float, optional): Minimum memory budget. Defaults to 0.001.
            iterations (int, optional): Number of memory budgets to evaluate. Defaults to 100.

        Returns:
            float: Memory budget at the knee point.
        """
        import numpy as np

        results = self.evaluate_distribution_of_results_for_knapsack_algo(
            knapsack_algo=knapsack_algo,
            memory_budget_values=np.linspace(  # type: ignore[arg-type]
                min_mem_budget, max_mem_budget, iterations
            ).tolist(),
        )
        runtime_values = np.array(
            [result["percentage_of_theoretical_peak_runtime"] for result in results]
        )
        memory_values = np.array(
            [result["percentage_of_theoretical_peak_memory"] for result in results]
        )
        runtime_range = np.ptp(runtime_values)
        memory_range = np.ptp(memory_values)
        if runtime_range == 0 or memory_range == 0:
            return max_mem_budget
        runtime_norm = (runtime_values - runtime_values.min()) / runtime_range
        memory_norm = (memory_values - memory_values.min()) / memory_range
        distances = np.sqrt(runtime_norm**2 + memory_norm**2)
        knee_index = np.argmin(distances)
        return results[knee_index]["memory_budget"]