from collections.abc import Sequence
from pathlib import Path

import numpy as np

from volara.tmp import replace_values

from .utils import StrictBaseModel


class LUT(StrictBaseModel):
    """
    A class for defining look up tables
    """

    path: Path | str
    """
    The path at which we will read/write the look up table
    """

    @property
    def name(self) -> str:
        return self.file.stem

    @property
    def file(self) -> Path:
        if isinstance(self.path, str):
            return (
                Path(self.path)
                if self.path.endswith(".npz")
                else Path(f"{self.path}.npz")
            )
        elif isinstance(self.path, Path):
            return self.path
        else:
            raise TypeError(f"Invalid type for path ({self.path}): {type(self.path)}")

    def drop(self):
        if self.file.exists():
            self.file.unlink()

    def save(self, lut: np.ndarray, edges=None):
        if edges is not None:
            np.savez_compressed(
                self.file, fragment_segment_lut=lut.astype(int), edges=edges
            )
        else:
            np.savez_compressed(self.file, fragment_segment_lut=lut.astype(int))

    def load(self) -> np.ndarray | None:
        if not self.file.exists():
            return None
        return np.load(self.file)["fragment_segment_lut"]

    def __add__(self, other):
        """
        Add two disjoint LUTs together via simple concatenation.
        i.e. {0:1} + {1:2} = {0:2}
        """
        if isinstance(other, LUT):
            return LUTS(luts=[self, other])
        raise TypeError(f"Cannot add {type(other)} to LUT")


class LUTS:
    def __init__(self, luts: LUT | Sequence[LUT]):
        self.luts: Sequence[LUT] = luts if not isinstance(luts, LUT) else [luts]

    def __add__(self, other):
        if isinstance(other, LUTS):
            return LUTS(self.luts + other.luts)
        elif isinstance(other, LUT):
            return LUTS(list(self.luts) + [other])
        raise TypeError(f"Cannot add {type(other)} to LUTS")

    def load(self):
        return np.concatenate(
            [lut.load() for lut in self.luts if lut.load() is not None], axis=1
        )  # type: ignore

    def load_iterated(self):
        starting_map = self.luts[0].load()
        assert starting_map is not None, "No lookup tables to load"
        for lut in self.luts[1:]:
            next_map = lut.load()
            if next_map is not None:
                starting_map[1] = replace_values(
                    starting_map[1], next_map[0], next_map[1]
                )
        return starting_map
