from collections import defaultdict
from enum import Enum
from dataclasses import dataclass, replace
import itertools
from numbers import Number
from typing import Literal, TypeVar

from accelforge._accelerated_imports import pandas as pd
from accelforge.frontend.mapping import (
    Mapping,
    Spatial,
    TensorHolder,
    Reservation as MappingReservation,
    Split as MappingSplit,
    TilePattern,
    Loop as MappingLoop,
)
from accelforge.frontend.renames import Rank, RankVariable, TensorName
from accelforge.frontend.workload import Einsum
from accelforge.mapper.FFM._pareto_df.df_convention import (
    is_fused_loop_col,
    make_fused_loop_col,
    stride2col,
    initial2col,
    iterations2col,
)

from accelforge.util import _expfmt, fzs, oset

# Abstractions:
# 1. Each tensor is stored above some loop index. 0 is the outermost loop, 1 the
#    next-innermost...
# 2. All loops above any shared tensor are co-tiled and must match between PmappingGroups.

T = TypeVar("T", bound="Updatable")


class Updatable:
    def update(self: T, **kwargs) -> T:
        return replace(self, **kwargs)


def _update_rename_dict(
    renames: dict[str, str],
    new_renames: dict[str, str],
):
    for mine, other in new_renames.items():
        if mine not in renames:
            renames[mine] = other
        elif renames[mine] != other:
            raise ValueError(
                f"Renaming {mine} to {other} conflicts with {renames[mine]}"
            )


@dataclass(frozen=True, order=True, eq=True)
class Loop(Updatable):
    rank_name: Rank
    tile_pattern: TilePattern | None
    is_spatial: bool

    def __post_init__(self):
        assert isinstance(self.rank_name, Rank)
        assert isinstance(self.tile_pattern, Number | TilePattern | str | None)
        assert isinstance(self.is_spatial, bool)
        assert isinstance(
            self.tile_pattern.initial_tile_shape,
            Number | str | None,
        )
        assert isinstance(
            self.tile_pattern.tile_shape,
            Number | str | None,
        )

    def __repr__(self):
        return (
            f"Loop({self.rank_name.__repr__()}, {self.tile_pattern}, {self.is_spatial})"
        )

    def __str__(self):
        return (
            "S-" if self.is_spatial else ""
        ) + f"{self.rank_name}-{self.tile_pattern}"

    def pydot_str(self):
        if self.is_spatial:
            return f"S-for R{self.rank_name} size {_expfmt(self.tile_pattern)}"
        return f"for {self.rank_name} size {_expfmt(self.tile_pattern)}"

    def to_yaml(self):
        return {"type": "loop", **self.__dict__}

    def merge_next(self, right: "Loop") -> "Loop":
        assert self.tile_pattern == right.tile_pattern
        return Loop(
            self.rank_name | right.rank_name,
            right.tile_pattern,
            self.is_spatial,
        )

    def clear_loop_bound(self, value=0):
        return self.update(tile_pattern=value)

    def populate(self, nloop: int, ranks_with_tile_pattern: set[Rank]) -> "Loop":
        initial = (
            initial2col(self.rank_name, nloop)
            if self.rank_name in ranks_with_tile_pattern
            else None
        )
        tile_shape = (
            stride2col(self.rank_name, nloop)
            if self.rank_name in ranks_with_tile_pattern
            else None
        )
        tile_pattern = TilePattern(
            tile_shape=tile_shape,
            initial_tile_shape=initial,
            calculated_n_iterations=iterations2col(nloop),
        )
        return self.update(tile_pattern=tile_pattern)

    def _prepend_symbols(self, prepend: str) -> "Loop":
        return self.update(tile_pattern=self.tile_pattern._prepend_symbols(prepend))

    def clear_symbolic_tile_patterns(self) -> "Loop":
        return self.update(tile_pattern=self.tile_pattern._clear_symbols())

    def make_fused_loop_symbols(self, prefix: str) -> tuple[dict[str, str], "Loop"]:
        r = {}
        new = self

        def replace(attr, new):
            g = getattr(self.tile_pattern, attr)
            if not isinstance(g, str):
                return new
            g2 = make_fused_loop_col(f"{prefix}<SEP>{g}")
            r[g] = g2
            return new.update(tile_pattern=new.tile_pattern.update(**{attr: g2}))

        for s in new.tile_pattern._symbol_attrs():
            new = replace(s, new)

        return r, new

    def add_n_iteration_symbols(self) -> "Loop":
        return self.update(tile_pattern=self.tile_pattern.add_n_iteration_symbols())

    def _rename_to_match(self, other: "Loop") -> tuple["Loop", dict[str, str]]:
        new_tp, renames = self.tile_pattern._rename_to_match(other.tile_pattern)
        return self.update(rank_name=other.rank_name, tile_pattern=new_tp), renames


@dataclass(frozen=True, eq=True, order=True)
class TensorReservation(Updatable):
    # This order is important. Above loop index should be before resource name
    # so when we sort reservations for tensors then the backing tensor holder comes
    # first.
    # Size is not included in hash or equality functions. This is because there
    # may be floating point rounding errors in reservation sizes. The other
    # attributes are sufficient to determine equality.
    loops: tuple[Loop]
    name: TensorName
    resource_name: str
    persistent: bool = False

    def __post_init__(self):
        if self.persistent:
            assert len(self.loops) == 0, "Persistent tensors be above all loops"

    @property
    def above_loop_index(self) -> int:
        return -1 if self.persistent else len(self.loops)

    def __str__(self):
        return f"[{self.resource_name}] {self.name} below {self.loops}"

    def __repr__(self):
        return f"Reservation({repr(self.name)}, {repr(self.loops)}, {repr(self.resource_name)})"

    def pydot_str(self):
        return f"[{self.resource_name}] {self.name}"

    def permute(self, permutation) -> "Reservation":
        new_loops = [self.loops[permutation[i]] for i in range(len(self.loops))]
        return self.update(loops=tuple(new_loops))

    def clear_loop_bounds(self) -> "Reservation":
        return self.update(loops=tuple(loop.clear_loop_bound() for loop in self.loops))

    def populate_loops(self, ranks_with_tile_pattern: set[Rank]) -> "TensorReservation":
        return self.update(
            loops=tuple(
                loop.populate(nloop, ranks_with_tile_pattern)
                for nloop, loop in enumerate(self.loops)
            )
        )

    @staticmethod
    def get_backing_tensors(
        all_tensors: set["TensorReservation"],
    ) -> list["TensorReservation"]:
        id2tensor = defaultdict(lambda: [])
        for t in all_tensors:
            id2tensor[t.name].append(t)
        return sorted(sorted(v)[0] for v in id2tensor.values())

    def drop_loop_indices(self, loop_indices: set[int]) -> "TensorReservation":
        loops = tuple(l for i, l in enumerate(self.loops) if i not in loop_indices)
        return self.update(loops=loops)

    def _prepend_symbols(self, prepend: str) -> "TensorReservation":
        return self.update(
            loops=tuple(l._prepend_symbols(prepend) for l in self.loops),
        )

    def clear_symbolic_tile_patterns(self) -> "TensorReservation":
        return self.update(
            loops=tuple(l.clear_symbolic_tile_patterns() for l in self.loops),
        )

    def make_fused_loop_symbols(
        self, prefix: str
    ) -> tuple[dict[str, str], "TensorReservation"]:
        result = {}
        loops = []
        for l in self.loops:
            r, l = l.make_fused_loop_symbols(prefix)
            result.update(r)
            loops.append(l)
        return result, self.update(loops=tuple(loops))

    def add_n_iteration_symbols(self) -> "TensorReservation":
        return self.update(
            loops=tuple(l.add_n_iteration_symbols() for l in self.loops),
        )

    def _rename_to_match(
        self, other: "TensorReservation"
    ) -> tuple["TensorReservation", dict[str, str]]:
        renames = {}
        new_loops = []
        for l_mine, l_other in zip(self.loops, other.loops):
            l_mine, new_renames = l_mine._rename_to_match(l_other)
            _update_rename_dict(renames, new_renames)
            new_loops.append(l_mine)
        return self.update(loops=tuple(new_loops)), renames


class SplitKind(Enum):
    SEQUENTIAL = 0
    PIPELINE = 1


@dataclass(frozen=True, order=True, eq=True)
class Split:
    split: MappingSplit
    above_loop_index: int


@dataclass(frozen=True)
class Compatibility(Updatable):
    tensors: fzs[TensorReservation]
    splits: fzs[Split] = fzs()
    reservation_indices: fzs[int] = fzs()
    check_reservation_indices: bool = True

    @property
    def n_loops(self) -> int:
        try:
            return object.__getattribute__(self, "_n_loops_cached")
        except AttributeError:
            val = max(
                (len(s.loops) for s in frozenset.__iter__(self.tensors)), default=0
            )
            object.__setattr__(self, "_n_loops_cached", val)
            return val

    @property
    def loops(self) -> tuple[Loop, ...]:
        return max([t.loops for t in self.tensors], key=len) if self.tensors else ()

    def _get_hash_tuple(self):
        return self.n_loops, self.tensors, self.reservation_indices

    def __hash__(self):
        return hash(self._get_hash_tuple())

    def __eq__(self, other):
        if self is other:
            return True
        return self._get_hash_tuple() == other._get_hash_tuple()

    def __post_init__(self):
        assert isinstance(self.tensors, fzs)
        assert isinstance(self.splits, fzs)
        assert isinstance(self.reservation_indices, fzs)
        if self.check_reservation_indices:
            assert isinstance(self.n_loops, int)
            assert (
                max(self.reservation_indices, default=-1) <= self.n_loops
            ), f"Extra reservation indices {self.reservation_indices} are greater than n_loops {self.n_loops}"
            p = f"are not in reservation indices {self.reservation_indices}"
            assert all(
                i >= 0 for i in self.reservation_indices
            ), f"Reservation indices {self.reservation_indices} are not all >= 0"
            assert all(
                s.above_loop_index in self.reservation_indices for s in self.splits
            ), f"Split above loop indices {self.splits} {p}"
            assert all(
                len(s.loops) in self.reservation_indices for s in self.tensors
            ), f"Tensor loops {self.tensors} {p}"

    def get_backing_levels(self) -> dict[str, int]:
        backings = {}
        for t in self.tensors:
            prev = backings.get(t.name, t.above_loop_index)
            backings[t.name] = min(prev, t.above_loop_index)
        return backings

    @property
    def tensor_names(self) -> set[str]:
        return oset(t.name for t in self.tensors)

    @property
    def max_above_loop_index(self) -> int:
        if len(self.tensors) == 0:
            return 0
        return max(s.above_loop_index for s in self.tensors)

    def shared_loop_index(self, live_tensors: set[str]) -> int:
        n = [l for t, l in self.get_backing_levels().items() if t in live_tensors]
        return max(n) - 1 if n else -1

    def __len__(self) -> int:
        return self.max_above_loop_index

    def _rename_to_match(
        self, other: "Compatibility"
    ) -> tuple["Compatibility", dict[str, str]]:
        renames = {}
        assert (
            self.clear_symbolic_tile_patterns() == other.clear_symbolic_tile_patterns()
        )
        tensors = []
        for t in self.tensors:
            other_t = other.get_reservation_of_tensor(t.name)
            t, new_renames = t._rename_to_match(other_t)
            tensors.append(t)
            _update_rename_dict(renames, new_renames)

        return (
            Compatibility(
                tensors=fzs(tensors),
                splits=self.splits,
                reservation_indices=self.reservation_indices,
            ),
            renames,
        )

    def clear_dead_tensors(
        self,
        live_tensors: set[str] | Literal["All"],
        keep_reservation_indices_and_splits: bool = False,
    ) -> "Compatibility":
        """
        Return a new compatibility with "dead" tensors removed by:
        1. keeping only loops relevant to `live_tensors` and
        2. keeping only `live_tensors`.

        If `keep_loops` is `True`, then all loops are kept.
        If `keep_tensors` is a set, tensors in the set are kept.
        """
        if live_tensors == "All":
            live_tensors = self.tensor_names

        remaining_tensors = fzs(s for s in self.tensors if s.name in live_tensors)
        new_n_loops = max((len(s.loops) for s in remaining_tensors), default=0)
        splits, reservation_indices = self.splits, self.reservation_indices
        if not keep_reservation_indices_and_splits:
            splits = fzs(s for s in splits if s.above_loop_index < new_n_loops)
            reservation_indices = fzs(
                oset(min(i, new_n_loops) for i in self.reservation_indices)
            )
            reservation_indices = fzs(x for x in reservation_indices if x >= 0)

        new_check_reservation_indices = (
            self.check_reservation_indices and not keep_reservation_indices_and_splits
        )
        return self.update(
            tensors=remaining_tensors,
            splits=splits,
            reservation_indices=reservation_indices,
            check_reservation_indices=new_check_reservation_indices,
        )

    def __lt__(self, other):
        return self._get_hash_tuple() < other._get_hash_tuple()

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return f"Compatibility(n_loops={self.n_loops}, tensors={repr(self.tensors)}, splits={repr(self.splits)}, reservation_indices={repr(self.reservation_indices)})"

    def _and_tensors_with_names(self, names: set[str]) -> "Compatibility":
        return fzs(s for s in self.tensors if s.name in names)

    def merge_next(
        self,
        right: "Compatibility",
        live_tensors: set[str],
        _force_allow_invalid_only_for_runtime_test: bool = False,
    ) -> "Compatibility":
        self_freed = self.clear_dead_tensors(
            live_tensors, keep_reservation_indices_and_splits=True
        )
        right_freed = right.clear_dead_tensors(
            live_tensors, keep_reservation_indices_and_splits=True
        )
        if (
            self_freed.n_loops > right_freed.n_loops
            and not _force_allow_invalid_only_for_runtime_test
        ):
            # This can be relaxed if we have a way to do order-independent joining
            # and/or non-looptree mappings.
            raise ValueError(
                f"Can't merge. I have more loops than the next, so my dataflow can't "
                f"be carried through a LoopTree to where it's needed."
            )

        live_minus_mine = live_tensors - oset(s.name for s in self.tensors)
        tensors_a = self._and_tensors_with_names(live_tensors)
        tensors_b = right._and_tensors_with_names(live_minus_mine)

        # TODO: split handling?
        joined = Compatibility(
            tensors=tensors_a | tensors_b,
            reservation_indices=self.reservation_indices | right.reservation_indices,
            check_reservation_indices=False,
        ).clear_dead_tensors(live_tensors)

        return joined

    def has_tensor(self, *tensors: TensorReservation) -> bool:
        return all(any(s == t for s in self.tensors) for t in tensors)

    def _permute_stops(self):
        stops = oset(len(s.loops) for s in self.tensors)
        stops |= self.reservation_indices
        stops |= oset(s.above_loop_index for s in self.splits)
        return stops

    def permute(
        self,
        loop_changes: list[int],
    ) -> "Compatibility":
        assert len(loop_changes) <= self.n_loops
        assert oset(loop_changes) == oset(
            range(len(loop_changes))
        ), f"Loop changes {loop_changes} are not a permutation of {range(len(loop_changes))}"
        if len(loop_changes) < len(self.loops):
            loop_changes = loop_changes + list(
                range(len(loop_changes), len(self.loops))
            )

        permute_stops = self._permute_stops()
        for i, c in enumerate(loop_changes):
            for r in permute_stops:
                assert (i < r) == (
                    c < r
                ), f"Loop changes {loop_changes} cross reservation {r}"
        new_tensors = fzs(s.permute(loop_changes) for s in self.tensors)
        return self.update(tensors=new_tensors)

    def make_equivalent_permutations(self) -> list[tuple["Compatibility", list[int]]]:
        # Get contiguous blocks of loops with no tensor reservation between them
        blocks = []
        current_block = []
        permute_stops = self._permute_stops()
        for i in range(self.n_loops):
            # Can't permute loops if there's a reservation between them
            if i in permute_stops:
                blocks.append(current_block)
                current_block = []
            current_block.append(i)
        if current_block:
            blocks.append(current_block)

        per_block_permutations = [
            list(itertools.permutations(block)) for block in blocks
        ]
        all_permutations = list(itertools.product(*per_block_permutations))
        all_permutations = [
            list(itertools.chain(*loop_changes)) for loop_changes in all_permutations
        ]
        return [(self.permute(p), p) for p in all_permutations]

    def get_reservation_of_tensor(self, tensor: str) -> TensorReservation:
        for s in self.tensors:
            if s.name == tensor:
                return s
        raise ValueError(f"No reservation found for {tensor}")

    def per_tensor_compatibility(self) -> dict[str, "Compatibility"]:
        result = {}
        for s in self.tensors:
            result[s.name] = self.clear_dead_tensors(oset([s.name]))
        return result

    def clear_loop_bounds(self) -> "Compatibility":
        return self.update(tensors=fzs(t.clear_loop_bounds() for t in self.tensors))

    def compatible_with(self, other: "Compatibility") -> bool:
        return True
        # for a in self.tensors:
        #     a = a.loops
        #     for b in other.tensors:
        #         b = b.loops
        #         if a[:len(b)] != b[:len(a)]:
        #             return False
        # return True

    def populate_loops(self, ranks_with_tile_pattern: set[Rank]):
        return self.update(
            tensors=fzs(
                t.populate_loops(ranks_with_tile_pattern) for t in self.tensors
            ),
        )

    @classmethod
    def from_mapping(
        cls,
        mapping: Mapping,
        tensors: set[TensorName],
        einsum: Einsum,
    ) -> "Compatibility":
        """
        Create Compatibility from a mapping, a set of fusable tensors, and the
        workload.
        """
        if not isinstance(einsum, Einsum):
            raise TypeError(f"einsum should be an Einsum, but {type(einsum)} instead")
        rank_variable_to_ranks = {
            t.name: t.rank_variable2ranks for t in einsum.tensor_accesses
        }

        # TODO: update compatibility to handle spatial-for loop per-tensor update
        tensor_indices = []
        split_above_loop_indices = []
        reservation_indices = []
        backing_remaining = oset(tensors)
        n_seen_loops = 0
        n_fused_loops = 0
        for i, n in enumerate(mapping.nodes):
            if isinstance(n, MappingReservation):
                reservation_indices.append(n_seen_loops)
                if not (backing := oset(n.purposes) & backing_remaining):
                    continue
                backing_remaining -= backing
                assert (
                    len(n.purposes) == 1
                ), "Backing reservations should have only one purpose"
                tensor_indices.append(i)
            elif isinstance(n, MappingSplit):
                split_above_loop_indices.append(n_seen_loops)
            elif isinstance(n, MappingLoop):
                n_seen_loops += 1
                n_fused_loops += bool(backing_remaining)
            elif isinstance(n, TensorHolder):
                reservation_indices.append(n_seen_loops)

        reservation_indices = fzs(min(n, n_fused_loops) for n in reservation_indices)
        reservation_indices = fzs(x for x in reservation_indices if x >= 0)

        assert (
            not backing_remaining
        ), f"Tensors {backing_remaining} not found in mapping"

        def get_rank(rank_variable, tensor):
            rv = rank_variable_to_ranks[tensor].get(rank_variable, oset())
            assert (
                len(rv) <= 1
            ), f"Rank variable {rank_variable} indexes into multiple ranks {rv} for tensor {tensor} "
            return next(iter(rv), Rank("NO RANK. RECOMPUTED."))

        def make_loops(above_index: int, tensor_name: TensorName) -> list[MappingLoop]:
            loops = [
                n for n in mapping.nodes[:above_index] if isinstance(n, MappingLoop)
            ]
            loops = [
                Loop(
                    rank_name=get_rank(n.rank_variable, tensor_name),
                    tile_pattern=n.tile_pattern._symbol2str(),
                    is_spatial=isinstance(n, Spatial),
                )
                for n in loops
            ]
            return tuple(loops)

        return cls(
            tensors=fzs(
                TensorReservation(
                    name=mapping.nodes[i].purpose,
                    loops=make_loops(i, mapping.nodes[i].purpose),
                    resource_name=mapping.nodes[i].resource,
                    persistent=mapping.nodes[i].persistent,
                )
                for i in tensor_indices
            ),
            splits=fzs(
                Split(split=n, above_loop_index=i) for i in split_above_loop_indices
            ),
            reservation_indices=fzs(reservation_indices),
        )

    def symbols(self) -> list[str]:
        symbols = []

        def add(x: str):
            if isinstance(x, str) and x not in symbols:
                symbols.append(x)

        for t in self.tensors:
            for l in t.loops:
                add(l.tile_pattern.initial_tile_shape)
                add(l.tile_pattern.tile_shape)
                add(l.tile_pattern.calculated_n_iterations)
        return symbols

    def drop_loop_indices(self, loop_indices: set[int]) -> "Compatibility":
        loop_indices = oset(loop_indices)
        tensors = fzs(t.drop_loop_indices(loop_indices) for t in self.tensors)
        splits = fzs(s for s in self.splits if s.above_loop_index not in loop_indices)

        def adjust(i: int) -> int:
            return i - sum(x < i for x in loop_indices)

        reservation_indices = fzs(adjust(i) for i in self.reservation_indices)
        reservation_indices = fzs(x for x in reservation_indices if x >= 0)

        splits = fzs(
            s.update(above_loop_index=adjust(s.above_loop_index)) for s in self.splits
        )
        return Compatibility(
            tensors=tensors,
            splits=splits,
            reservation_indices=reservation_indices,
        )

    def _prepend_symbols(self, prepend: str) -> "Compatibility":
        return self.update(
            tensors=fzs(t._prepend_symbols(prepend) for t in self.tensors)
        )

    def clear_tile_patterns_and_reservation_indices(self) -> "Compatibility":
        return self.update(
            tensors=fzs(t.clear_symbolic_tile_patterns() for t in self.tensors),
            reservation_indices=fzs(),
            check_reservation_indices=False,
        )

    def clear_symbolic_tile_patterns(self) -> "Compatibility":
        return self.update(
            tensors=fzs(t.clear_symbolic_tile_patterns() for t in self.tensors)
        )

    def make_fused_loop_symbols(
        self, prefix: str
    ) -> tuple[dict[str, str], "Compatibility"]:
        result = {}
        tensors = []
        for t in self.tensors:
            r, t = t.make_fused_loop_symbols(prefix)
            tensors.append(t)
            result.update(r)

        return result, self.update(tensors=fzs(tensors))

    def add_n_iteration_symbols(self) -> "Compatibility":
        return self.update(
            tensors=fzs(t.add_n_iteration_symbols() for t in self.tensors)
        )

    def clear_unrelated_columns(self, mappings: pd.DataFrame) -> "Compatibility":
        my_symbols = oset(self.symbols())
        for c in my_symbols:
            assert c in mappings.columns, f"Column {c} not found in mappings"
        keep = [
            c
            for c in mappings.columns
            if not (is_fused_loop_col(c) and c not in my_symbols)
        ]
        return mappings if len(keep) == len(mappings.columns) else mappings[keep]
