import copy
from typing import Callable
from uuid import UUID
from accelforge.frontend import arch
from accelforge.frontend.spec import Spec
from accelforge.mapper.FFM._join_pmappings.pmapping_group import PmappingGroup
from accelforge.frontend.workload import EinsumName
from accelforge.frontend.mapping import Mapping
from accelforge.mapper.FFM._make_pmappings.pmapper_job import Job


class MultiEinsumPmappings:
    """
    A collection of pmappings for each Einsum in a workload, generated by
    :func:`~accelforge.mapper.FFM.make_pmappings`.

    Parameters
    ----------
    einsum2pmappings:
        A dictionary of Einsum names to lists of PmappingGroups. PmappingGroups contain
        the Pareto-optimal pmappings for the Einsum.
    pmapping_objects:
        A dictionary of Einsum names to dictionaries of UUIDs to Mappings. The entries
        in the PmappingGroup objects reference these pmapping objects.
    einsum2jobs:
        A dictionary of Einsum names to lists of Jobs that generated the pmappings.
    can_combine_multiple_runs:
        Whether the pmappings can be combined from multiple runs. If this is True, then
        multiple make_pmappings calls can be summed together to get a single
        MultiEinsumPmappings object. If this is True, the mapper may run more slowly.
    einsums_with_pmappings_generated:
        Einsums for which pmappings were generated (or attempted to be generated).
    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,
        einsum2pmappings: dict[EinsumName, list[PmappingGroup]],
        pmapping_objects: dict[EinsumName, dict[UUID, Mapping]],
        einsum2jobs: dict[EinsumName, list[Job]],
        can_combine_multiple_runs: bool,
        einsums_with_pmappings_generated: set[EinsumName],
        flattened_arches: dict[EinsumName, list[arch.Leaf]],
        evaluated_specs: dict[EinsumName, Spec],
    ):
        self.spec: Spec = spec
        self.einsum2pmappings: dict[EinsumName, list[PmappingGroup]] = einsum2pmappings
        self.pmapping_objects: dict[EinsumName, dict[UUID, Mapping]] = pmapping_objects
        self.einsum2jobs: dict[EinsumName, list[Job]] = einsum2jobs
        self.can_combine_multiple_runs: bool = can_combine_multiple_runs
        self.einsums_with_pmappings_generated: set[EinsumName] = (
            einsums_with_pmappings_generated
        )
        self.flattened_arches: dict[(EinsumName, str), list[arch.Leaf]] = (
            flattened_arches
        )
        self.evaluated_specs: dict[EinsumName, Spec] = evaluated_specs

    def __or__(self, other: "MultiEinsumPmappings"):
        if not self.can_combine_multiple_runs or not other.can_combine_multiple_runs:
            raise ValueError(
                "Must call make_pmappings with can_combine_multiple_runs=True to combine pmappings "
                "from multiple runs."
            )
        self = copy.copy(self)
        for einsum_name, pmappings in other.einsum2pmappings.items():
            self.einsum2pmappings.setdefault(einsum_name, []).extend(pmappings)
        for einsum_name, jobs in other.einsum2jobs.items():
            self.einsum2jobs.setdefault(einsum_name, []).extend(jobs)
        self.pmapping_objects.update(other.pmapping_objects)
        self.einsums_with_pmappings_generated.update(
            other.einsums_with_pmappings_generated
        )
        self.evaluated_specs.update(other.evaluated_specs)
        self.flattened_arches.update(other.flattened_arches)
        return self

    def _filter(
        self,
        filter_lambda: Callable[[PmappingGroup], bool],
        einsums_with_pmappings_generated: list[EinsumName] | None = None,
    ):
        new_einsum2pmappings = {}
        if einsums_with_pmappings_generated is None:
            einsums_with_pmappings_generated = list(self.einsum2pmappings.keys())
        for einsum_name in einsums_with_pmappings_generated:
            new_einsum2pmappings[einsum_name] = [
                pm for pm in self.einsum2pmappings[einsum_name] if filter_lambda(pm)
            ]

        return MultiEinsumPmappings(
            spec=self.spec,
            einsum2pmappings=new_einsum2pmappings,
            pmapping_objects=self.pmapping_objects,
            einsum2jobs=self.einsum2jobs,
            can_combine_multiple_runs=self.can_combine_multiple_runs,
            einsums_with_pmappings_generated=self.einsums_with_pmappings_generated,
            flattened_arches=self.flattened_arches,
            evaluated_specs=self.evaluated_specs,
        )

    def drop_einsums(self, *einsums_with_pmappings_generated: EinsumName):
        """
        Removes all pmappings for the given Einsums.

        Parameters
        ----------
        einsums_with_pmappings_generated:
            The Einsums for which to remove pmappings.
        """
        for einsum_name in einsums_with_pmappings_generated:
            del self.einsum2pmappings[einsum_name]
            del self.pmapping_objects[einsum_name]
            del self.einsum2jobs[einsum_name]
            self.einsums_with_pmappings_generated.remove(einsum_name)

    def pmapping_keep_rates(
        self, per_einsum: bool = False
    ) -> dict[EinsumName, dict[str, float]] | dict[str, float]:
        """
        Returns the keep rates for each cause of pmapping removal. For example, if only
        25% of the pmappings have a valid spatial fanout, then the keep rate for the
        spatial fanout cause will be 0.25.

        Parameters
        ----------
        per_einsum:
            If True, returns a dictionary of keep rates for each Einsum.

        Returns
        -------
        dict[EinsumName, dict[str, float]] | dict[str, float]
            A dictionary of keep rates for each cause of pmapping removal.
        """
        result = {}
        einsum2npmappings = self.n_total_pmappings(per_einsum=True)

        for einsum_name, jobs in self.einsum2jobs.items():
            cur_result = result.setdefault(einsum_name, {})
            for job in jobs:
                for cause, keep_rate in job.pmapping_keep_rates.items():
                    cur_result.setdefault(cause, 0)
                    cur_result[cause] += job.n_total_pmappings * keep_rate

        if per_einsum:
            for einsum_name, npmappings in einsum2npmappings.items():
                for cause, keep_rate in result[einsum_name].items():
                    result[einsum_name][cause] = keep_rate / npmappings
        else:
            new_result = {}
            n_total_pmappings = sum(einsum2npmappings.values())
            for einsum_name, keep_rates in result.items():
                for cause, keep_rate in keep_rates.items():
                    new_result.setdefault(cause, 0)
                    new_result[cause] += keep_rate / n_total_pmappings
            result = new_result

        return result

    def n_total_pmappings(
        self, per_einsum: bool = False
    ) -> int | dict[EinsumName, int]:
        """
        Returns the number of total pmappings in the mapspace.

        Parameters
        ----------
        per_einsum:
            If True, returns a dictionary of total pmappings for each Einsum.

        Returns
        -------
        int | dict[EinsumName, int]
            The number of total pmappings in the mapspace.
        """
        result = {
            einsum_name: sum(job.n_total_pmappings for job in jobs)
            for einsum_name, jobs in self.einsum2jobs.items()
        }
        if per_einsum:
            return result
        return sum(result.values())

    def n_valid_pmappings(
        self, per_einsum: bool = False
    ) -> int | dict[EinsumName, int]:
        """
        Returns the number of valid pmappings for each Einsum. A valid pmapping is one
        that satisfies all constraints and resource usage limits.

        Parameters
        ----------
        per_einsum:
            If True, returns a dictionary of valid pmappings for each Einsum.

        Returns
        -------
        int | dict[EinsumName, int]
            The number of valid pmappings in the mapspace.
        """
        result = {
            einsum_name: sum(job.n_valid_pmappings for job in jobs)
            for einsum_name, jobs in self.einsum2jobs.items()
        }
        if per_einsum:
            return result
        return sum(result.values())

    def n_pareto_optimal_pmappings(
        self, per_einsum: bool = False
    ) -> int | dict[EinsumName, int]:
        """
        Returns the number of Pareto-optimal pmappings for each Einsum. This is the
        number of mappings that will be returned by the make_pmappings function.

        Parameters
        ----------
        per_einsum:
            If True, returns a dictionary of Pareto-optimal pmappings for each Einsum.

        Returns
        -------
        int | dict[EinsumName, int]
            The number of Pareto-optimal pmappings in the mapspace.
        """
        result = {
            einsum_name: sum(len(p) for p in pmappings)
            for einsum_name, pmappings in self.einsum2pmappings.items()
        }
        if per_einsum:
            return result
        return sum(result.values())

    def n_evaluated_pmappings(
        self, per_einsum: bool = False
    ) -> int | dict[EinsumName, int]:
        """
        Returns the number of pmappings that were evaluated for each Einsum. This is
        greater than the number of Pareto-optimal pmappings because some mappings are
        found to be suboptimal after they have been evaluated.

        Parameters
        ----------
        per_einsum:
            If True, returns a dictionary of evaluated pmappings for each Einsum.

        Returns
        -------
        int | dict[EinsumName, int]
            The number of evaluated pmappings in the mapspace.
        """

        result = {
            einsum_name: sum(job.n_evaluated_pmappings for job in jobs)
            for einsum_name, jobs in self.einsum2jobs.items()
        }
        if per_einsum:
            return result
        return sum(result.values())

    def n_pmapping_string(self) -> str:
        """
        Returns a string representation of the number of pmappings in the mapspace.
        Printing this can help diagnose if the mapper is not finding any pmappings or
        mappings.

        Returns
        -------
        str
            A string representation of the number of pmappings in the mapspace.
        """
        if "Total" in self.einsum2pmappings:
            raise ValueError(
                f"Cannot print stats for a MultiEinsumPmappings object that has "
                f"an Einsum named 'Total'. Use a different name for the Einsum."
            )

        n_total_pmappings = self.n_total_pmappings(per_einsum=True)
        n_valid_pmappings = self.n_valid_pmappings(per_einsum=True)
        n_evaluated_pmappings = self.n_evaluated_pmappings(per_einsum=True)
        n_pareto_optimal_pmappings = self.n_pareto_optimal_pmappings(per_einsum=True)

        for x in (
            n_total_pmappings,
            n_valid_pmappings,
            n_evaluated_pmappings,
            n_pareto_optimal_pmappings,
        ):
            x["Total"] = sum(x.values())

        s = []
        for e in n_total_pmappings:
            t = n_total_pmappings[e]
            v = n_valid_pmappings[e]
            ev = n_evaluated_pmappings[e]
            p = n_pareto_optimal_pmappings[e]

            def fmt(x, total: bool = True):
                x = round(x)

                def _f(y):
                    y = round(y)
                    return str(y) if y < 1000 else f"{y:.2e}".replace("e+", "e")

                divved = _f(round(t) / x) if x != 0 else "inf"
                return f"{_f(x)} (1/{divved})" if total else _f(x)

            s.append(
                f"{e}: {fmt(t, False)} total, {fmt(v)} valid, {fmt(ev)} evaluated, "
                f"{fmt(p)} Pareto-Optimal"
            )
        return "\n".join(s)
