from accelforge.frontend.renames import TensorName
from collections import defaultdict
from accelforge.frontend import arch
from accelforge.frontend.spec import Spec
from accelforge.frontend.workload import EinsumName
from accelforge._accelerated_imports import pd
from accelforge.mapper.FFM._make_pmappings.make_pmappings import (
    get_num_computes,
    get_per_tensor_size,
)
from typing import Union
import numpy as np
from numbers import Integral


class Mappings:
    """
    A collection of mappings and their evaluation results, generated by
    :func:`~accelforge.mapper.FFM.join_pmappings`.

    Attributes
    ----------
    spec:
        The specification used to generate the mappings.
    einsum_names:
        The names of the Einsums in these mappings.
    data:
        A DataFrame containing the mappings and their evaluation results. Column names
        in the dataframe are are string separated by "<SEP>", such as
        "Total<SEP>energy".
    total_mappings:
        The total number of mappings that have been explored in order to get the
        mappings in the dataframe, equal to the full mapspace size.
    valid_mappings:
        The number of valid mappings that have been explored in order to get the
        mappings in the dataframe, equal to the valid mapspace size.
    flattened_arches:
        A dictionary of (EinsumName, Compute Name) to lists of architecture nodes. These
        contain the evaluated and flattened architecture node for that particular Einsum
        and compute combination.
    evaluated_specs:
        A dictionary of Einsum names to evaluated specifications. These contain the
        evaluated specification for that particular Einsum.
    """

    def __init__(
        self,
        spec: Spec,
        einsum_names: list[EinsumName],
        data: pd.DataFrame,
        total_mappings: int,
        valid_mappings: int,
        flattened_arches: dict[(EinsumName, str), list[arch.Leaf]],
        evaluated_specs: dict[EinsumName, Spec],
    ):
        self.spec: Spec = spec
        self.einsum_names: list[EinsumName] = einsum_names
        self.data: pd.DataFrame = data
        self.total_mappings: int = total_mappings
        self.valid_mappings: int = valid_mappings
        self.flattened_arches: dict[(EinsumName, str), list[arch.Leaf]] = (
            flattened_arches
        )
        self.evaluated_specs: dict[EinsumName, Spec] = evaluated_specs

    def num_computes(self, einsum_name: EinsumName | None = None) -> int:
        """
        Returns the number of computes for the given Einsum name, or total computes if
        ``einsum_name`` is ``None``.
        """
        # TODO: this is not correct if there are recomputations.
        if einsum_name is None:
            return sum(get_num_computes(self.spec, e) for e in self.einsum_names)
        return get_num_computes(self.spec, einsum_name)

    def per_tensor_size(self, return_n_elements: bool = False) -> dict[TensorName, int]:
        """
        Returns a dictionary of: {Tensor name: Number of bits} for each tensor in the
        spec. If return_n_elements is true, then the number of elements is returned
        instead of the number of bits.

        Parameters
        ----------
        return_n_elements:
            If True, then the number of elements is returned instead of the number of
            bits.

        Returns
        -------
        dict[TensorName, int]
            A dictionary of: {Tensor name: Number of bits} for each tensor in the spec.
        """
        assert (
            self.evaluated_specs is not None
        ), "Can't get per-tensor size if no Einsums have been mapped."
        return get_per_tensor_size(
            next(iter(self.evaluated_specs.values())).workload,
            return_n_elements=return_n_elements,
        )

    def _update(self, **kwargs):
        data = dict(
            spec=self.spec,
            einsum_names=self.einsum_names,
            data=self.data,
            total_mappings=self.total_mappings,
            valid_mappings=self.valid_mappings,
            flattened_arches=self.flattened_arches,
            evaluated_specs=self.evaluated_specs,
        )
        data.update(kwargs)
        return Mappings(**data)

    def __getitem__(self, key: str | Integral) -> Union[pd.Series, "Mappings"]:
        if isinstance(key, Integral):
            return self._update(data=pd.DataFrame(self.data.iloc[key]).T)
        return self.data[key]

    def __len__(self):
        return len(self.data)

    def __iter__(self):
        return iter(self.data)

    def _get_cols(self, key: str, col_idx: int | None = None) -> list[str]:
        found_index = col_idx
        found = []
        for col in self.data.columns:
            col = col.split("<SEP>")
            if key not in col:
                continue

            if col_idx is not None and col.index(key) != col_idx:
                continue

            if col_idx is None and sum(c == key for c in col) > 1:
                raise ValueError(
                    f"Key {key} found multiple times in the column names. "
                    f'Columns: "{col}"'
                )

            cur_idx = col.index(key) if col_idx is None else col_idx

            if found_index is not None and cur_idx != found_index:
                raise ValueError(
                    f"Key {key} found at varying indexes in the column names. "
                    f'Columns: "{col}" and "{found}"'
                )
            found_index = cur_idx
            found.append("<SEP>".join(col))
        return found, found_index

    def access(self, *keys: str, col_idx: int | None = None) -> "Mappings":
        """
        Returns a new Mappings object with only the columns that contain the given keys.
        Column names are strings separated by "<SEP>", and this method will return
        columns with one <SEP>-separated string matching the given key. Then, for all
        remaining columns, the key will be removed.

        For example, if the columns are "Compute<SEP>Energy", and "DRAM<SEP>Energy", and
        "DRAM<SEP>Latency", then access("Energy") will return a Mappings object with
        columns "Compute" and "DRAM", and access("DRAM") will return a Mappings object
        with columns "Compute" and "Latency".

        If multiple keys are given, then the procedure is repeated for each key.

        col_idx:
            The index of the key in the column name. This can be used if the given key
            is found at multiple indexes in different columns.

        Parameters
        ----------
        keys:
            The keys to access from the columns.

        Returns
        -------
        Mappings
            A new Mappings object with only the given keys.
        """
        assert len(set(self.data.columns)) == len(
            self.data.columns
        ), "Columns must be unique"

        if len(keys) != 1:
            for k in keys:
                self = self.access(k, col_idx=col_idx)
            return self

        key = keys[0]
        col_renames = {}
        cols, col_idx = self._get_cols(key, col_idx=col_idx)
        for col in cols:
            split = col.split("<SEP>")
            new = "<SEP>".join(c for i, c in enumerate(split) if i != col_idx)
            if new in col_renames:
                raise ValueError(
                    f"Removing {key} from {col} results in duplicate column name {new}"
                )
            col_renames[new] = col

        col_renames = {v: k for k, v in col_renames.items()}
        return self._update(
            data=self.data[list(col_renames.keys())].rename(columns=col_renames)
        )

    def _get_keys_of_length(self, length: int) -> list[str]:
        cols = [c for c in self.columns if len(c.split("<SEP>")) == length]
        return cols

    def drop(self, *keys: str) -> "Mappings":
        """
        Returns a new Mappings object with the given keys dropped from all columns.
        Column names are strings separated by "<SEP>", and this method will will drop
        columns with one <SEP>-separated string matching the given key. Then, for all
        remaining columns, the key will be removed.

        For example, if the columns are "Compute<SEP>Energy", and "DRAM<SEP>Energy", and
        "DRAM<SEP>Latency", then drop("Energy") will drop "Compute<SEP>Energy" and
        "DRAM<SEP>Energy", and drop("DRAM") will drop "DRAM<SEP>Energy" and
        "DRAM<SEP>Latency".

        If multiple keys are given, then the procedure is repeated for each key.

        Parameters
        ----------
        keys:
            The keys to drop from the columns.

        Returns
        -------
        Mappings
            A new Mappings object with the given keys dropped from all columns.
        """
        assert len(set(self.data.columns)) == len(
            self.data.columns
        ), "Columns must be unique"

        if len(keys) != 1:
            for k in keys:
                self = self.drop(k)
            return self

        cols, _ = self._get_cols(keys[0])
        return self._update(data=self.data.drop(columns=cols))

    def sum(self, keep_key_index: list[int] | int | None = None) -> "Mappings":
        if len(self.data.columns) == 1:
            return self

        if isinstance(keep_key_index, int):
            keep_key_index = [keep_key_index]
        elif keep_key_index is None:
            keep_key_index = []

        columns = list(self.data.columns)
        for col in self.data.columns:
            if len(col.split("<SEP>")) != len(columns[0].split("<SEP>")):
                raise ValueError(
                    "Can only sum columns with same-length keys. Try first calling "
                    'access("key") or drop("key") to make all columns '
                    "have the same number of keys."
                )

        max_n = len(columns[0].split("<SEP>"))
        if any(k < 0 or k >= max_n for k in keep_key_index):
            raise ValueError(f"Keep indices must be in the range [0, {max_n})")

        target2sources = {}
        for col in columns:
            target = col.split("<SEP>")
            target = "<SEP>".join(target[i] for i in keep_key_index)
            target2sources.setdefault(target, []).append(col)

        new_data = pd.DataFrame(index=self.data.index)
        for target, sources in target2sources.items():
            new_data[target] = self.data[sources].sum(axis=1)

        return self._update(data=new_data)

    @property
    def columns(self) -> list[str]:
        """The columns of the dataframe."""
        return list(self.data.columns)

    def to_dict(self, value_if_one_mapping: bool = True) -> dict[str, list[float]]:
        """
        Returns the data in this Mappings object as a dictionary. Each column in the
        this Mappings' data becomes a key in the dictionary. Values in the dictionary
        may be a single value if there is only one mapping, or a list of values if there
        are multiple mappings.

        Parameters
        ----------
        value_if_one_mapping:
            If True and there is only one mapping, then values in the returned
            dictionary will be a single value, rather than a list of values. Otherwise,
            they will always be a list of values.

        Returns
        -------
        dict[str, list[float]]
            A dictionary with the same keys as the columns of the dataframe, and values
            that are either a single value or a list of values.
        """
        new = self.data.to_dict(orient="list")
        if value_if_one_mapping and len(self) == 1:
            new = {k: v[0] for k, v in new.items()}
        return new

    def items(
        self, value_if_one_mapping: bool = True
    ) -> list[tuple[str, float | list[float]]]:
        """
        Shorthand for ``to_dict().items()``. Casts this Mappings object to a dictionary,
        then returns the items of the dictionary.

        Parameters
        ----------
        value_if_one_mapping:
            If True and there is only one mapping, then values in the returned list will
            be a single value, rather than a list of values. Otherwise, they will always
            be a list of values.

        Returns
        -------
        list[tuple[str, float | list[float]]]
            A list of tuples of (key, value).
        """
        return self.to_dict(value_if_one_mapping=value_if_one_mapping).items()

    def per_compute(self, per_einsum: bool = False) -> "Mappings":
        """
        Returns the per-compute evaluation results by dividing all statistics by the
        number of computes.

        Parameters
        ----------
        per_einsum:
            If True, then for each Einsum, statistics will be divided only by the number
            of computes for that Einsum. This makes it clear to see per-compute stats
            for each Einsum, but note that total energy/compute will not be a direct sum
            of the per-compute stats then (and the same for latency and other
            statistics).

        Returns
        -------
        Mappings
            A new Mappings object with the per-compute evaluation results.
        """
        new_df = self.data.copy()
        total_computes = self.num_computes()
        for col in new_df.columns:
            n_computes = total_computes
            if per_einsum:
                einsum_name = col.split("<SEP>")[0]
                if einsum_name not in self.einsum_names and einsum_name != "Total":
                    raise ValueError(
                        f"Einsum name {einsum_name} not found. Ensure that all "
                        f"columns are prefixed with the Einsum name if per_einsum "
                        f"is True."
                    )
                if einsum_name != "Total":
                    n_computes = self.num_computes(einsum_name)
            # Check if the column can be converted to numeric
            try:
                new_df[col] /= n_computes
            except (ValueError, TypeError):
                # Skip columns that can't be converted to numeric
                continue
        return self._update(data=new_df)

    def drop_zeros(self) -> "Mappings":
        """
        Returns a new Mappings object with all columns that have only zeros dropped.
        """
        new_df = self.data.copy()
        new_df = new_df[(c for c in new_df.columns if (new_df[c] != 0).any())]
        return self._update(data=new_df)

    def _repr_svg_(self) -> str:
        return self.render()

    def render(self, index: int | None = None) -> str:
        """
        Renders the mapping as a Pydot graph. Returns an SVG string. This is only
        supported if there is a single mapping; if there are multiple mappings, then
        either index into this Mappings object first, or pass in an index.

        Parameters
        ----------
        index:
            The index of the mapping to render. If None and there are multiple mappings,
            then an error is raised.

        Returns
        -------
        str
            An SVG string of the mapping.

        Raises
        ------
        ValueError
            If there are multiple mappings and no index is provided.
        """
        if index is not None:
            self = self[index]
        if len(self) != 1:
            raise ValueError(
                f"Can only render a single mapping, but got {len(self)}. Try calling "
                f"mappings[i].render() instead, for some integer 0 <= i < {len(self)}."
            )
        return self.data.iloc[0]["Total<SEP>mapping"].render()

    def energy(
        self: "Mappings",
        per_einsum: bool = False,
        per_component: bool = False,
        per_tensor: bool = False,
        per_action: bool = False,
        value_if_one_mapping: bool = True,
    ) -> dict[tuple[str, ...] | str, float | list[float]] | float | list[float]:
        """
        Returns the energy consumed. A dictionary is returned with keys that are tuples
        of (Einsum name, Component name, Tensor name, Action name), with any of these
        being omitted if the corresponding parameter is not set to True. If neither of
        the ``per_`` parameters are set to True, a float or a list of floats is
        returned.

        NOTE: Leak power is not per-tensor. If ``per_tensor`` is True, then the tensor
        name for leak will be None.

        Parameters
        ----------
        per_einsum:
            If True, then the energy will reported per-Einsum.
        per_component:
            If True, then the energy will reported per-component.
        per_tensor:
            If True, then the energy will reported per-tensor.
        per_action:
            If True, then the energy will reported per-action.
        value_if_one_mapping:
            If True and there is only one mapping, then values in the returned
            dictionary will be a single value, rather than a list of values. Otherwise,
            they will always be a list of values.

        Returns
        -------
        dict[tuple[str, ...], float | list[float]] | float | list[float]
            A dictionary with the energy consumed for each Einsum, Component, Tensor,
            and Action. Keys are tuples of (Einsum name, Component name, Tensor name,
            Action name), with any of these being omitted if the corresponding parameter
            is not set to True. If none of the ``per_`` parameters are set to True, a
            float or a list of floats is returned.
        """

        energy = self.access("energy")

        result = {}
        for einsum in self.einsum_names:
            einsum_accessed = energy.access(einsum, col_idx=0)
            # None for computes
            for tensor in list(self.spec.workload.einsums[einsum].tensor_names) + [
                "None"
            ]:
                tensor_accessed = einsum_accessed.access(tensor, col_idx=1)
                for col in tensor_accessed._get_keys_of_length(2):
                    component, action = col.split("<SEP>")
                    result[(einsum, component, tensor, action)] = tensor_accessed[col]
            for col in einsum_accessed._get_keys_of_length(2):
                component, action = col.split("<SEP>")
                if action == "leak":
                    result[(einsum, component, None, action)] = einsum_accessed[col]

        keep_indices = []
        for i, idx in enumerate([per_einsum, per_component, per_tensor, per_action]):
            if idx:
                keep_indices.append(i)

        if not keep_indices:
            v = sum(result.values())
            if value_if_one_mapping and len(self.data) == 1:
                return v.iloc[0] if isinstance(v, pd.Series) else v
            return v

        new_result = defaultdict(float)
        for key, value in result.items():
            newkey = tuple(key[i] for i in keep_indices)
            new_result[newkey] += value
        result = new_result

        if len(keep_indices) == 1:
            result = {k[0]: v for k, v in result.items()}

        if value_if_one_mapping and len(self.data) == 1:
            return {k: v.iloc[0] for k, v in result.items()}

        return {k: list(v) for k, v in result.items()}

    def latency(
        self: "Mappings",
        per_einsum: bool = False,
        per_component: bool = False,
        value_if_one_mapping: bool = True,
    ) -> dict[tuple[str, ...] | str, float | list[float]] | float | list[float]:
        """
        Returns the latency consumed. A dictionary is returned with keys that are tuples
        of (Einsum name, Component name), with either being omitted if the corresponding
        parameter is not set to True. If neither of the ``per_`` parameters are set to
        True, a float or a list of floats is returned.

        NOTE: If per-Einsum is False and per-component is True, then the latency of each
        component will be summed across all Einsums. THE TOTAL LATENCY MAY BE GREATER
        THAN THE MAX OF THE PER-COMPONENT LATENCIES. This is because different
        components can be the bottleneck for different Einsums.

        Parameters
        ----------
        per_einsum:
            If True, then the latency will reported per-Einsum.
        per_component:
            If True, then the latency will reported per-component.
        value_if_one_mapping:
            If True and there is only one mapping, then values in the returned
            dictionary will be a single value, rather than a list of values. Otherwise,
            they will always be a list of values.

        Returns
        -------
        dict[tuple[str, ...], float | list[float]] | float | list[float]
            A dictionary with the latency for each Einsum+Component pair. Keys are
            tuples of (Einsum name, Component name), with either being omitted if the
            corresponding parameter is not set to True. If neither of the ``per_``
            parameters are set to True, a float or a list of floats is returned.
        """

        energy = self.access("latency")

        result = {}
        for einsum in self.einsum_names:
            einsum_accessed = energy.access(einsum, col_idx=0)
            for component in einsum_accessed._get_keys_of_length(1):
                result[(einsum, component)] = einsum_accessed[component]

        # If not per-component, aggregate into the per-Einsum latency
        if not per_component:
            new_result = {}
            for (einsum, component), value in result.items():
                if einsum not in new_result:
                    new_result[einsum] = value
                else:
                    new_result[einsum] = np.maximum(new_result[einsum], value)
            result = new_result

        if not per_einsum:
            # Not per-Einsum and per-component: for each component, sum across Einsums
            if per_component:
                new_result = {}
                for (einsum, component), value in result.items():
                    if component not in new_result:
                        new_result[component] = value
                    else:
                        new_result[component] += value
                result = new_result

            # Not per-Einsum and not per-component: sum into a single value
            else:
                result = np.sum(list(result.values()))

        if value_if_one_mapping and len(self.data) == 1:
            if isinstance(result, dict):
                return {k: v.iloc[0] for k, v in result.items()}
            return result  # Numpy sum already pulls out the number

        return {k: list(v) for k, v in result.items()}

    def resource_usage(
        self,
        value_if_one_mapping: bool = True,
    ) -> dict[str, float | list[float]]:
        """
        Returns the usage of each resource in the mapping.

        Parameters
        ----------
        value_if_one_mapping:
            If True and there is only one mapping, then values in the returned
            dictionary will be a single value, rather than a list of values. Otherwise,
            they will always be a list of values.

        Returns
        -------
        dict[str, float | list[float]]
            A dictionary with the usage of each resource.
        """
        try:
            reservations = self.access("reservation")
        except KeyError:
            raise ValueError(f"No reservations found in columns {sorted(self.columns)}")

        usage = {}
        for col in reservations._get_keys_of_length(3):
            resource, _, _ = col.split("<SEP>")
            if resource not in usage:
                usage[resource] = 0
            usage[resource] = np.maximum(usage[resource], reservations[col])

        if value_if_one_mapping and len(self.data) == 1:
            return {k: v.iloc[0] for k, v in usage.items()}
        return {k: list(v) for k, v in usage.items()}

    def actions(
        self: "Mappings",
        per_einsum: bool = False,
        per_component: bool = True,
        per_tensor: bool = False,
        value_if_one_mapping: bool = True,
    ) -> dict[tuple[str, ...] | str, float | list[float]] | float | list[float]:
        """
        Returns the number of actions performed in the mapping. A dictionary is returned
        with keys that are tuples of (Einsum name, Component name, Tensor name, Action
        name), with any of these being omitted if the corresponding parameter is not set
        to True. If neither of the ``per_`` parameters are set to True, a float or a
        list of floats is returned.

        Parameters
        ----------
        per_einsum:
            If True, then the number of actions will reported per-Einsum.
        per_component:
            If True, then the number of actions will reported per-component.
        per_tensor:
            If True, then the number of actions will reported per-tensor.
        value_if_one_mapping:
            If True and there is only one mapping, then values in the returned
            dictionary will be a single value, rather than a list of values. Otherwise,
            they will always be a list of values.

        Returns
        -------
        dict[tuple[str, ...], float | list[float]] | float | list[float]
            A dictionary with the number of actions performed for each Einsum,
            Component, Tensor, and Action. Keys are tuples of (Einsum name, Component
            name, Tensor name, Action name), with any of these being omitted if the
            corresponding parameter is not set to True. If none of the ``per_``
            parameters are set to True, a float or a list of floats is returned.
        """

        actions = self.access("action")

        result = {}
        for einsum in self.einsum_names:
            einsum_accessed = actions.access(einsum, col_idx=0)
            # None for computes
            for tensor in list(self.spec.workload.einsums[einsum].tensor_names) + [
                "None"
            ]:
                tensor_accessed = einsum_accessed.access(tensor, col_idx=1)
                for col in tensor_accessed._get_keys_of_length(2):
                    component, action = col.split("<SEP>")
                    result[(einsum, component, tensor, action)] = tensor_accessed[col]

        keep_indices = []
        for i, idx in enumerate([per_einsum, per_component, per_tensor, True]):
            if idx:
                keep_indices.append(i)

        if not keep_indices:
            v = sum(result.values())
            if value_if_one_mapping and len(self.data) == 1:
                return v.iloc[0] if isinstance(v, pd.Series) else v
            return v

        new_result = defaultdict(float)
        for key, value in result.items():
            newkey = tuple(key[i] for i in keep_indices)
            new_result[newkey] += value
        result = new_result

        if len(keep_indices) == 1:
            result = {k[0]: v for k, v in result.items()}

        if value_if_one_mapping and len(self.data) == 1:
            return {k: v.iloc[0] for k, v in result.items()}

        return {k: list(v) for k, v in result.items()}
