from collections import defaultdict

from accelforge.frontend import arch
from accelforge.frontend.arch import Leaf, Memory, TensorHolder, Component
from accelforge.frontend.mapping import Compute, Mapping
from accelforge.frontend.spec import Spec

from accelforge.model._looptree.accesses import isl_buffer_accesses_from_buffet_actions
from accelforge.model._looptree.mapping_utilities import get_leaves
from accelforge.model._looptree.reuse.isl import IslReuseAnalysisOutput
from accelforge.model._looptree.reuse import SymbolicAnalysisOutput
from accelforge.model._looptree.types import Buffet

from accelforge.model._looptree.reuse.symbolic import BuffetStats
from accelforge.util._eval_expressions import MATH_FUNCS, eval_expression
from accelforge.util._sympy.broadcast_max import Max, Min
import sympy as sp


def isl_to_summarized(
    looptree_results: IslReuseAnalysisOutput, mapping, workload
) -> SymbolicAnalysisOutput:
    accesses_stats = isl_buffer_accesses_from_buffet_actions(
        looptree_results, mapping, workload, is_path=False
    )
    buffet_stats = {
        Buffet(level=component, tensor=tensor, einsum=einsum): BuffetStats(
            max_per_unit_read_actions=accesses.max_per_unit_reads,
            max_per_unit_write_actions=accesses.max_per_unit_writes,
        )
        for (component, tensor, einsum), accesses in accesses_stats.items()
    }
    return SymbolicAnalysisOutput(buffet_stats=buffet_stats)


def component_latency(
    looptree_results: SymbolicAnalysisOutput,
    flattened_arch: list[Leaf],
    mapping: Mapping,
    spec: Spec,
):
    component_to_actions: dict[str, dict[str, float]] = defaultdict(
        lambda: defaultdict(lambda: 0)
    )
    name2component: dict[str, Component] = {node.name: node for node in flattened_arch}

    compute_obj = flattened_arch[-1]
    if not isinstance(compute_obj, arch.Compute):
        raise ValueError("Last node in flattened_arch must be a Compute")

    for buffet, buffet_stats in looptree_results.buffet_stats.items():
        component = buffet.level
        actions = component_to_actions[component]
        if component not in name2component:
            raise ValueError(f"Component {component} found in mapping but not arch")

        for action in name2component[component].actions:
            actions[f"{action.name}_actions"] += 0

        if isinstance(name2component[component], TensorHolder):
            actions["read_actions"] += (
                buffet_stats.max_per_unit_read_actions
                - buffet_stats.min_per_unit_skipped_first_read_actions
            )
            if not isinstance(name2component[component], arch.Toll):
                actions["write_actions"] += (
                    buffet_stats.max_per_unit_write_actions
                    - buffet_stats.min_per_unit_skipped_first_write_actions
                )
        elif isinstance(name2component[component], arch.Compute):
            pass
        else:
            raise NotImplementedError(
                f"Component {component} is not a TensorHolder or Compute"
            )

    longest_compute_latency = Max(
        0, *[s.max_latency for s in looptree_results.compute_stats.values()]
    )
    component_to_actions[compute_obj.name]["compute_actions"] = longest_compute_latency

    # TODO: Unhardcode "compute" name"
    component_to_action_latency = defaultdict(dict)
    for component, actions in component_to_actions.items():
        component_obj = name2component[component]
        for action, count in actions.items():
            action_name = action.rsplit("_", 1)[0]
            latency = component_obj.actions[action_name].latency
            component_to_action_latency[component][f"{action_name}_latency"] = (
                latency * count
            )

    component_latency = {}

    symbol_table_base = {
        **dict(spec.variables),
        "variables": spec.variables,
        "max": Max,
        "min": Min,
        "sum": sp.Add,
    }

    for component, actions in component_to_actions.items():
        component_obj = name2component[component]
        symbol_table = {
            "action2latency": component_to_action_latency[component],
            **symbol_table_base,
            **dict(name2component[component]),
            **actions,
            **component_to_action_latency[component],
        }
        if name2component[component].total_latency is not None:
            component_latency[component] = eval_expression(
                name2component[component].total_latency,
                symbol_table,
                attr_name="latency",
                location=component,
            )

    return component_latency
