"""Implement a class to describe a pure state photonic state."""

from __future__ import annotations

import logging
from typing import Self, Literal
from numbers import Number
from copy import deepcopy

import numpy as np
from numpy.typing import NDArray, ArrayLike
from scipy.sparse import sparray, coo_array, csr_array

from qoptcraft.algebra import photon_basis, basis_and_index
from qoptcraft.operators import creation_fock, annihilation_fock

from ._exceptions import (
    ProbabilityError,
    PureStateLengthError,
    NumberPhotonsError,
    NumberModesError,
    NotHermitianError,
)


class State:
    """Base type for our quantum states."""

    ...


class MixedState(State):
    """A mixed quantum state.

    Args:
        density_matrix (NDArray): provide directly the density matrix of the mixed state.
        modes (int): number of modes in the optical network.
        photons (int): number of photons.

    Attributes:
        density_matrix (NDArray): density matrix of the state.
        modes (int): number of modes in the optical network.
        photons (int): number of photons.
    """

    def __init__(
        self,
        density_matrix: NDArray | sparray,
        modes: int,
        photons: int,
    ) -> None:
        def is_hermitian(density_matrix, rtol: float = 1e-5, atol: float = 1e-8):
            """Reimplement numpy isclose for spare matrices"""
            mat_1 = density_matrix
            mat_2 = density_matrix.conj().T
            return abs(mat_1 - mat_2).max() <= atol + rtol * abs(mat_2).max()

        if not is_hermitian(density_matrix):
            raise NotHermitianError()
        if not np.isclose(density_matrix.trace(), 1):
            raise ValueError("Density matrix trace is not 1.")
        self._density_matrix = density_matrix
        self.photons = photons
        self.modes = modes

    @classmethod
    def from_mixture(cls, pure_states: list[PureState], probs: ArrayLike) -> Self:
        """Initialize a mixed state from a superposition of pure states instead of
        initializing directly the density matrix.

        Args:
            pure_states (list[PureState]): pure states in the statistical mixture.
            probs (ArrayLike): probability of measuring each pure state.

        Raises:
            PureStateLengthError: States and probabilities differ in length.
            ProbabilityError: Probabilities don't add up to 1.
            NumberPhotonsError: Not all states have the same number of photons.
            NumberModesError: Not all states have the same number of modes.

        Returns:
            MixedState
        """
        if not len(pure_states) == len(probs):
            raise PureStateLengthError()
        if not np.isclose(1, sum(probs)):
            raise ProbabilityError(sum(probs))

        photons_list = [state.photons for state in pure_states]
        if not all(photons == photons_list[0] for photons in photons_list):
            raise NumberPhotonsError()

        modes_list = [state.modes for state in pure_states]
        if not all(modes == modes_list[0] for modes in modes_list):
            raise NumberModesError()

        density_matrix = probs[0] * pure_states[0].density_matrix
        for i in range(1, len(probs)):
            density_matrix += probs[i] * pure_states[i].density_matrix

        return cls(density_matrix, modes_list[0], photons_list[0])

    @property
    def density_matrix(self):
        return self._density_matrix

    def evolution(self, unitary: NDArray) -> Self:
        new_density = unitary @ self.density_matrix @ unitary.T.conj()
        return MixedState(new_density, self.modes, self.photons)


class PureState(State):
    """A pure quantum state.

    Args:
        fock_states (tuple[tuple[int, ...]]): Fock states that, in superposition,
            constitute our pure state.
        coefs (ArrayLike): amplitude of each Fock state in the superposition.

    Attributes:
        modes (int): number of modes in the optical network.
        photons (int): number of photons.
        basis (tuple[tuple[int, ...]]): basis of the Hilbert state. Defaults to None.
        density_matrix (NDArray): density matrix of the state.
    """

    # @classmethod
    # def _fast_init(cls, fock_states, coefs, modes, photons):
    #     obj = cls.__new__(cls)
    #     obj.fock_states = fock_states
    #     obj.coefs = coefs
    #     obj.modes = modes
    #     obj.photons = photons
    #     ...
    #     return obj

    def __init__(self, fock_states: list[tuple[int, ...]], coefs: ArrayLike) -> None:
        self._assert_inputs(fock_states, coefs)

        self.photons: int = sum(fock_states[0])
        self.modes: int = len(fock_states[0])

        # Convert fock_states to list of tuples
        self.fock_states = [tuple(f) if not isinstance(f, tuple) else f for f in fock_states]
        self.coefs = np.array(coefs)

        # sorted_inputs = sorted(
        #     zip(fock_states_list, coefs, strict=True), key=lambda pair: pair[0], reverse=True
        # )
        # self.fock_states, coefs = zip(*sorted_inputs, strict=True)

        sum_coefs = np.sum(np.abs(self.coefs) ** 2)
        if np.isclose(sum_coefs, 0):
            self.amplitudes = self.coefs
        else:
            self.amplitudes = np.array(coefs) / np.sqrt(sum_coefs)

        self.basis, self._basis_index = basis_and_index(self.modes, self.photons)

        # Get the amplitude of a certain fock state in the basis.
        self._amp_fock = None
        self._coef_fock = None

        self._state_vector = None
        self._density_matrix: NDArray | None = None
        self._density_matrix_sparse: sparray | None = None

    @staticmethod
    def _assert_inputs(fock_states: tuple[tuple[int, ...]], coefs: ArrayLike) -> None:
        """Assert the instance inputs are not contradictory.

        Args:
            fock_states (tuple[tuple[int, ...]]): Fock states that, in superposition,
                constitute our pure state.
            coefs (list[float]): amplitude of each Fock state in the superposition.

        Raises:
            NumberPhotonsError: Not all states have the same number of photons.
            NumberModesError: Not all states have the same number of modes.
            ProbabilityError: Probabilities don't add up to 1.
        """
        for fock in fock_states:
            for i in fock:
                if i < 0:
                    raise ValueError(f"Invalid fock state {fock}")

        photons_set = {sum(fock_state) for fock_state in fock_states}
        if len(photons_set) != 1:
            print(f"{photons_set = },\n {fock_states = }")
            raise NumberPhotonsError()

        modes_list = [len(fock_state) for fock_state in fock_states]
        if not all(modes == modes_list[0] for modes in modes_list):
            raise NumberModesError()

        if len(fock_states) != len(coefs):
            raise ValueError("Error: fock_states and amplitudes must have the same length.")

    def __str__(self) -> str:
        """String representation of the state."""
        sorted_inputs = sorted(
            zip(self.fock_states, self.coefs, strict=True), key=lambda pair: pair[0], reverse=True
            )
        str_ = ""
        for fock, amp in sorted_inputs:
            str_ = str_ + f"{amp:.2f} * {fock} + \n"
        str_ = str_[:-4]
        return str_

    def __repr__(self) -> str:
        """Representation of the state."""
        return str(self)

    def __eq__(self, other: PureState) -> bool:
        """Compare if two pure states are equal."""
        if isinstance(other, PureState):
            if self.photons != other.photons or self.modes != other.modes:
                return False
            if not np.allclose(self.state_vector, other.state_vector):
                print("Amplitudes differ")
                return False
            return True
        if isinstance(other, PureStateMultiSector):
            return self.to_multisector() == other
        raise NotImplementedError

    def __add__(self, other: PureState) -> Self:
        """Tensor product of self with other state."""
        if isinstance(other, Number):
            if other == 0:
                return self
            raise NotImplementedError

        if isinstance(other, PureState):
            if self.modes != other.modes:
                raise NumberModesError()

            if self.photons == other.photons:
                coefs = []
                fock_states = list(set(self.fock_states) | set(other.fock_states))
                for fock in fock_states:
                    try:
                        coefs.append(self.coef_fock[fock])
                        try:
                            coefs[-1] += other.coef_fock[fock]
                        except KeyError:
                            continue
                    except KeyError:
                        coefs.append(other.coef_fock[fock])
                return PureState(fock_states, coefs)
            # if self.photons != other.photons:
            return PureStateMultiSector(
                list(self.fock_states) + list(other.fock_states),
                list(self.coefs) + list(other.coefs),
            )
        raise NotImplementedError

    def __radd__(self, other: PureState) -> Self:
        """Tensor product of self with other state."""
        return self + other

    def __sub__(self, other: PureState) -> Self:
        """Substraction of states."""
        if isinstance(other, Vacuum):
            return self
        if isinstance(other, PureState):
            if self.photons != other.photons:
                raise NumberPhotonsError()
            if self.modes != other.modes:
                raise NumberModesError()
            return self + (-1) * other
        logging.error(f"Substraction not implemented for opperand type {type(other)}")
        raise NotImplementedError

    def __neg__(self) -> Self:
        return PureState(self.fock_states, -self.coefs)

    def __mul__(self, other: PureState | Number) -> Self:
        """Tensor product of self with other state."""
        if isinstance(other, Number):
            return PureState(self.fock_states, self.coefs * other)
        if isinstance(other, np.ndarray):
            if other.size == 1:
                logging.error("Operation not implemented for array of size bigger than 1.")
                return PureState(self.fock_states, self.coefs * other.item())
            else:
                raise NotImplementedError
        if isinstance(other, PureState):
            fock_states = []
            amplitudes = []
            for fock, amp in zip(self.fock_states, self.amplitudes, strict=True):
                for fock_other, amp_other in zip(other.fock_states, other.amplitudes, strict=True):
                    fock_states.append(fock + fock_other)
                    amplitudes.append(amp * amp_other)
            return PureState(fock_states, amplitudes)
        if isinstance(other, PureStateMultiSector):
            fock_states = []
            coefs = []
            for sector in other.sector_states.values():
                for f1, c1 in zip(self.fock_states, self.coefs, strict=True):
                    for f2, c2 in zip(sector.fock_states, sector.coefs, strict=True):
                        fock_states.append(f1 + f2)
                        coefs.append(c1 * c2)
            return PureStateMultiSector(fock_states, coefs)

        logging.error(f"Operation not implemented for opperand type {type(other)}")
        raise NotImplementedError

    def __rmul__(self, other: PureState | Number) -> Self:
        """Tensor product of self with other state."""
        return self * other

    def __truediv__(self, other: Number) -> Self:
        """Return the division of self and the other number.

        Args:
            other: Other scalar value (rhs).

        Raises:
            ValueError: Division by zero.
            TypeError: Not int/float passed in.

        Returns:
            The multiplication of self and the other vector/number.
        """
        if not isinstance(other, Number):
            raise TypeError("You must pass in a scalar value.")
        return self * (1 / other)

    def __pow__(self, exponent: int, modulo=None) -> Self:
        """Tensor product of self with itself a given number of times."""
        state = self
        for _ in range(exponent - 1):
            state = state * self
        return state

    def isclose(self, other: PureState, atol: float = 1e-3, rtol=1e-05) -> bool:
        """Compare if two pure states are equal."""
        if isinstance(other, PureState):
            if self.photons != other.photons or self.modes != other.modes:
                return False
            fock_states = list(set(self.fock_states) | set(other.fock_states))
            for fock in fock_states:
                try:
                    coef_self = np.abs(self.coef_fock[fock])
                    try:
                        coef_other = np.abs(other.coef_fock[fock])
                        if not np.isclose(coef_self, coef_other, atol=atol, rtol=rtol):
                            print(f"{fock = }, {coef_self = }, {coef_other = }")
                            return False
                    except KeyError:
                        if not np.isclose(coef_self, 0, atol=atol, rtol=rtol):
                            print(f"{fock = }, {coef_self = }, coef_other = 0")
                            return False
                except KeyError:
                    coef_other = np.abs(other.coef_fock[fock])
                    if not np.isclose(coef_other, 0, atol=atol, rtol=rtol):
                        print(f"{fock = }, coef_self = 0, {coef_other = }")
                        return False
            return True
        raise ValueError("other is not PureState")

    def conj(self) -> Self:
        """State with conjugated coefficients."""
        return PureState(self.fock_states, self.coefs.conj())

    def normalize(self):
        self.coefs = self.amplitudes

    @property
    def probabilities(self):
        return np.abs(self.amplitudes) ** 2

    @property
    def amp_fock(self):
        if self._amp_fock is None:
            self._amp_fock = dict(zip(self.fock_states, self.amplitudes))
        return self._amp_fock

    @property
    def coef_fock(self):
        if self._coef_fock is None:
            self._coef_fock = dict(zip(self.fock_states, self.coefs))
        return self._coef_fock

    def get_density_matrix(self, *, sparse):
        if sparse:
            return self.density_matrix_sparse
        return self.density_matrix

    @property
    def density_matrix_sparse(self):
        """Density matrix of the pure state in a certain basis."""
        if self._density_matrix_sparse is None:
            if self._density_matrix is not None:
                self._density_matrix_sparse = csr_array(self._density_matrix)
                return self._density_matrix_sparse
            else:
                if self.basis is None:
                    self.basis = photon_basis(self.modes, self.photons)
                rows = []
                cols = []
                data = []

                for i, fock_i in enumerate(self.fock_states):
                    idx_i = self.basis.index(fock_i)
                    amp_i = self.amplitudes[i]
                    for j, fock_j in enumerate(self.fock_states):
                        idx_j = self.basis.index(fock_j)
                        amp_j = self.amplitudes[j]
                        rows.append(idx_i)
                        cols.append(idx_j)
                        data.append(amp_i * amp_j.conjugate())

                self._density_matrix_sparse = coo_array(
                    (data, (rows, cols)),
                    shape=(len(self.basis), len(self.basis)),
                    dtype=np.complex128,
                )
                self._density_matrix_sparse = self._density_matrix_sparse.tocsr()
        return self._density_matrix_sparse

    @property
    def density_matrix(self):
        """Density matrix of the pure state in a certain basis."""
        if self._density_matrix is None:
            state_vector = self.state_vector
            self._density_matrix = np.outer(state_vector, state_vector.conj().T)
            return self._density_matrix
        return self._density_matrix

    @property
    def state_vector(self) -> NDArray:
        """Output the state vector in the Fock basis."""
        if self._state_vector is None:
            self._state_vector = np.zeros(len(self.basis), dtype=np.complex128)
            for amp, fock in zip(self.amplitudes, self.fock_states, strict=True):
                self._state_vector[self._basis_index[fock]] = amp
            return self._state_vector
        return self._state_vector

    def coefs_in_basis(self) -> NDArray:
        """Provide the coefs of the state in the fock basis."""
        state = np.zeros(len(self.basis), dtype=np.complex128)

        for i, fock in enumerate(self.fock_states):
            for j, basis_fock in enumerate(self.basis):
                if fock == basis_fock:
                    state[j] = self.coefs[i]
        return state

    def dot(self, other: PureState) -> complex:
        "Compute the overlap of two pure states. < self | other >"
        if isinstance(other, Vacuum):
            return 0

        if isinstance(other, PureState):
            prod = 0
            for i, fock_self in enumerate(self.fock_states):
                try:
                    prod += self.amplitudes[i].conjugate() * other.amp_fock[fock_self]
                except KeyError:
                    continue
            return prod

        if isinstance(other, PureStateMultiSector):
            N = self.photons
            if N not in other.sector_states:
                return 0
            return self.dot(other.sector_states[N]) * other.sector_amplitudes[N]

        return NotImplementedError

    def _dot_coefs(self, other: PureState) -> complex:
        if isinstance(other, Vacuum):
            return 0
        prod = 0
        for i, fock_self in enumerate(self.fock_states):
            try:
                prod += self.coefs[i] * other.coef_fock[fock_self].conjugate()
            except KeyError:
                continue
        return prod

    def exp_photons(self, mode_creat: int, mode_annih: int) -> float:
        r"""Compute the expecation value of $a^\dagger_i a_j$.

        Args:
            mode_creat (int): mode where we apply the creation operator.
            mode_annih (int): mode where we apply the annihilation operator.

        Returns:
            float: expectation value.
        """
        exp = 0
        if mode_creat == mode_annih:
            for i, fock in enumerate(self.fock_states):
                exp += self.probabilities[i] * fock[mode_creat]
        else:
            for i, fock in enumerate(self.fock_states):
                fock_, coef_ = annihilation_fock(mode_annih, fock)
                # if coef_ == 0:  # not really necessary to check this
                ##### continue  # since it will raise a ValueError below
                coef = self.amplitudes[i] * coef_
                fock_, coef_ = creation_fock(mode_creat, fock_)
                coef *= coef_
                try:
                    exp += coef * self.amp_fock[fock_].conjugate()
                except KeyError:
                    continue
        return exp

    def creation(self, mode: int) -> PureState:
        """Apply creation operator a†_mode.

        Args:
            mode (int): mode in which to create a photon.

        Returns:
            PureState: created state.
        """

        new_focks = []
        new_coefs = []

        if mode >= self.modes:
            raise ValueError(
                f"The given {mode = } is higher than the number of modes in this vacuum state: {self.modes}."+
                " Remember that modes start at 0, not 1."
            )

        for fock, coef in zip(self.fock_states, self.coefs, strict=True):
            # # extend mode dimension if needed
            # if mode >= len(fock):
            #     fock_new = tuple(list(fock) + [0] * (mode - len(fock)) + [1])
            #     coef_factor = 1.0
            photons = fock[mode]
            fock_new = list(fock)
            fock_new[mode] = photons + 1
            coef_factor = np.sqrt(photons + 1)

            new_focks.append(tuple(fock_new))
            new_coefs.append(coef * coef_factor)

        return PureState(new_focks, new_coefs)

    def annihilation(self, mode: int) -> PureState:
        """Apply annihilation operator a_mode.

        Args:
            mode (int): mode in which to annihilate.

        Returns:
            PureState: annihilated state.
        """

        if mode >= self.modes:
            raise ValueError(
                f"The given {mode = } is higher than the number of modes in this vacuum state: {self.modes}."+
                " Remember that modes start at 0, not 1."
            )

        new_focks = []
        new_coefs = []

        for fock, coef in zip(self.fock_states, self.coefs, strict=True):

            photons = fock[mode]

            if photons == 0:
                continue

            fock_list = list(fock)
            fock_list[mode] = photons - 1
            fock_new = tuple(fock_list)

            coef_factor = np.sqrt(photons)

            new_focks.append(fock_new)
            new_coefs.append(coef * coef_factor)

        if not new_focks:
            return 0

        return PureState(new_focks, new_coefs)

    def evolution(self, unitary: NDArray) -> Self:
        "Evolve the state with a quantum unitary."
        amps_vector = self.state_vector
        new_amps = unitary @ amps_vector
        return PureState(self.basis, new_amps)

    def evolution_vector(self, unitary: NDArray) -> Self:
        "Evolve the state with a quantum unitary."
        amps_vector = self.state_vector
        return unitary @ amps_vector

    def evolution_scattering(
        self,
        scattering_matrix: NDArray,
        method: Literal[
            "hamiltonian", "heisenberg", "permanent glynn", "permanent ryser"
        ] = "permanent glynn",
    ) -> Self:
        "Evolve the state with a quantum unitary."

        from qoptcraft.evolution import fock_evolution

        evolved_vector = np.zeros(len(self.basis), dtype=np.complex128)

        for amp, fock in zip(self.amplitudes, self.fock_states, strict=True):
            evolved_vector += amp * fock_evolution(scattering_matrix, fock, method)
        return evolved_vector

    def postselect(self, herald: Fock, reduce_dim=False):

        ps_amps = []
        ps_fock_list = []

        for amp, fock in zip(self.amplitudes, self.fock_states, strict=True):
            if fock[-herald.modes : :] == tuple(herald):
                ps_fock_list.append(fock[: -herald.modes] if reduce_dim else fock)
                ps_amps.append(amp)
        if len(ps_fock_list) == 0 or len(ps_amps) == 0:
            return Vacuum()
        else:
            return PureState(ps_fock_list, ps_amps)

    def to_multisector(self):
        """Convert a PureState to a PureStateMultisector"""
        return PureStateMultiSector.from_sector_states({self.photons: self})

    # def postselect_probability(self, mode_list, photon_list):

    #     if len(mode_list) != len(photon_list):
    #         raise ValueError("Mode and photon lists must have the same length.")

    #     ps_amps = []
    #     ps_fock_list = []

    #     for amp, fock in zip(self.amplitudes, self.fock_states):
    #         if all(fock[m] == n for m, n in zip(mode_list, photon_list)):
    #             if reduce_dim:
    #                 ps_fock = tuple(fock[m] for m in range(self.modes) if m not in mode_list)
    #             else:
    #                 ps_fock = fock
    #             ps_fock_list.append(ps_fock)
    #             ps_amps.append(amp)
    #     return PureState(ps_fock_list, ps_amps)


class Fock(PureState):
    def __init__(self, *photons: int, coef: Number = 1) -> None:
        super().__init__(fock_states=(photons,), coefs=(coef,))

    def __repr__(self) -> str:
        """Representation of the state."""
        return str(self)

    def __str__(self) -> str:
        numbers = ", ".join(map(str, self.fock_states[0]))
        return f"|{numbers}⟩"

    def __iter__(self):
        """Iterate the photons of the Fock state."""
        yield from self.fock_states[0]

    def __len__(self):
        """Number of modes in the fock state."""
        return self.modes


class Vacuum(Fock):
    def __init__(self, modes: int = 1):
        super().__init__(*([0] * modes))

    def __repr__(self) -> str:
        """Representation of the state."""
        return str(self)

    def __str__(self) -> str:
        numbers = ", ".join(map(str, tuple([0] * self.modes)))
        return f"|{numbers}⟩"

    def creation(self, mode: int) -> PureState:
        """Creation of a photon in a certain mode.

        Args:
            mode (int): mode in which to create. Starts at 0.

        Returns:
            PureState: created state.
        """
        if self.modes is None:
            return Fock(*(mode * [0]), 1)
        elif mode >= self.modes:
            raise ValueError(
                f"The given {mode = } is higher than the number of modes in this vacuum state: {self.modes}."+
                " Remember that modes start at 0, not 1."
            )
        else:
            new_fock = self.modes * [0]
            new_fock[mode] = 1
            return Fock(*new_fock)

    def annihilation(self, mode: int) -> Self:
        """Annihilation of a photon in a certain mode.

        Args:
            mode (int): mode in which to annihilate.

        Returns:
            PureState: created state.
        """
        return self


class PureStateMultiSector(State):
    """
    Multi-photon-number pure state.

    Internally stored as a direct sum of fixed-photon-number sectors.
    """

    @classmethod
    def from_sector_states(cls, sector_states: dict[int, PureState]):
        """
        Build a multi-sector state directly from
        { photon_number: PureState }.
        """
        focks = []
        coefs = []
        sector_states = dict(sector_states)
        for sector in sector_states.values():
            focks.extend(sector.fock_states)
            coefs.extend(sector.coefs)
        return cls(focks, coefs)

    def __init__(self, fock_states, coefs):

        if len(fock_states) == 0:
            raise ValueError("You must input some states.")
        if len(fock_states) != len(coefs):
            raise ValueError("fock_states and coefs must have same length.")

        # check modes consistency
        modes_set = {len(fock) for fock in fock_states}
        if len(modes_set) != 1:
            raise NumberModesError()
        self.modes = len(fock_states[0])

        # Convert fock_states to list of tuples
        self.fock_states = [tuple(f) if not isinstance(f, tuple) else f for f in fock_states]
        self.coefs = np.array(coefs)

        # group by photon number
        sectors = {}  # will be a dictionary of dictionaries
        for fock, coef in zip(self.fock_states, self.coefs, strict=True):
            N = int(sum(fock))
            sectors.setdefault(N, {"fock_states": [], "coefs": []})
            sectors[N]["fock_states"].append(fock)
            sectors[N]["coefs"].append(coef)

        # create fixed sectors
        self.sector_states = {
            N: PureState(data["fock_states"], data["coefs"]) for N, data in sectors.items()
        }

        self.sector_amplitudes = self._sector_amplitudes(sectors)

    def _sector_amplitudes(self, sectors) -> NDArray:
        """
        Compute and store the normalized amplitudes of each photon-number sector.

        Updates each sector's `amplitudes` attribute based on its `coefs`.
        """
        sector_amplitudes = {}
        total_amp_2 = 0
        for N, sector in sectors.items():
            sector_amp_2 = np.sum(np.abs(sector["coefs"]) ** 2)
            sector_amplitudes[N] = np.sqrt(sector_amp_2)
            total_amp_2 += sector_amp_2

        for N in sector_amplitudes:
            sector_amplitudes[N] /= np.sqrt(total_amp_2)

        return sector_amplitudes

    def __str__(self):
        lines = []
        for sector in self.sector_states.values():
            lines.append(f"{sector}")
        return "\n".join(lines)

    def __repr__(self) -> str:
        return self.__str__()

    def __add__(self, other):

        if isinstance(other, PureState):

            if self.modes != other.modes:
                raise NumberModesError()

            new_sectors = deepcopy(self.sector_states)
            N = other.photons
            if N in new_sectors:
                new_sectors[N] = new_sectors[N] + other
            else:
                new_sectors[N] = other

            return PureStateMultiSector.from_sector_states(new_sectors)

        if isinstance(other, PureStateMultiSector):

            if self.modes != other.modes:
                raise NumberModesError()

            new_sectors = deepcopy(self.sector_states)
            for N, sector_other in other.sector_states.items():
                if N in new_sectors:
                    new_sectors[N] = new_sectors[N] + sector_other
                else:
                    new_sectors[N] = sector_other

            return PureStateMultiSector.from_sector_states(new_sectors)

        raise NotImplementedError

    def __radd__(self, other) -> Self:
        return self + other

    def __sub__(self, other: PureState) -> Self:
        """Substraction of states."""
        if isinstance(other, PureState | PureStateMultiSector):
            return self + (-1) * other
        raise NotImplementedError

    def __rsub__(self, other: PureState) -> Self:
        """Substraction of states."""
        return self + (-1) * other

    def __mul__(self, other):

        if isinstance(other, Number):
            new_sectors = {}
            for N in self.sector_states:
                new_sectors[N] = other * self.sector_states[N]
            return PureStateMultiSector.from_sector_states(new_sectors)

        if isinstance(other, PureState):
            fock_states = []
            coefs = []
            for sector in self.sector_states.values():
                for f1, c1 in zip(sector.fock_states, sector.coefs, strict=True):
                    for f2, c2 in zip(other.fock_states, other.coefs, strict=True):
                        fock_states.append(f1 + f2)
                        coefs.append(c1 * c2)
            return PureStateMultiSector(fock_states, coefs)

        if isinstance(other, PureStateMultiSector):

            fock_states = []
            coefs = []
            for sec1 in self.sector_states.values():
                for sec2 in other.sector_states.values():

                    for f1, c1 in zip(sec1.fock_states, sec1.coefs):
                        for f2, c2 in zip(sec2.fock_states, sec2.coefs):
                            fock_states.append(f1 + f2)
                            coefs.append(c1 * c2)
            return PureStateMultiSector(fock_states, coefs)

        raise NotImplementedError

    def __rmul__(self, other) -> Self:
        return self * other

    def __eq__(self, other) -> bool:
        """Compare if two pure states are equal."""
        if isinstance(other, PureState):
            other_multisector = PureStateMultiSector.from_sector_states({other.photons: other})
            return self == other_multisector
        if isinstance(other, PureStateMultiSector):
            if self.modes != other.modes:
                return False
            for N in self.sector_states:
                try:
                    other_amp = other.sector_amplitudes[N]
                except KeyError:
                    return False
                if not np.isclose(self.sector_amplitudes[N], other_amp):
                    print("Amplitudes differ.")
                    return False
                if not self.sector_states[N] == other.sector_states[N]:
                    return False
            return True
        return False

    @property
    def photons(self):
        return sorted(self.sector_states.keys())

    @property
    def amplitudes(self):
        return np.array(self.coefs) / np.sqrt(np.sum(np.abs(self.coefs) ** 2))

    def evolution_scattering(self, S: NDArray) -> Self:
        """
        unitaries:
            Either:
                - single matrix (if only one photon sector)
                - dict {N: unitary_for_N}
        """
        new_sectors = {}
        for N, sector in self.sector_states.items():
            evolved_vector = self.sector_amplitudes[N] * sector.evolution_scattering(S)
            new_sectors[N] = PureState(photon_basis(self.modes, N), evolved_vector)
        return PureStateMultiSector.from_sector_states(new_sectors)

    def evolution_scattering_vector(self, S: NDArray) -> dict[NDArray]:
        """
        unitaries:
            Either:
                - single matrix (if only one photon sector)
                - dict {N: unitary_for_N}
        """
        evolved_vectors = {}
        for N, sector in self.sector_states.items():
            evolved_vectors[N] = self.sector_amplitudes[N] * sector.evolution_scattering(S)
        return evolved_vectors

    def normalize(self):
        for state in self.sector_states.values():
            state.normalize()

    def density_matrix_blocks(self):
        return {N: sector.density_matrix for N, sector in self.sector_states.items()}

    def dot(self, other):
        """Compute the overlap between two states: < self | other >"""

        if isinstance(other, PureState):
            N = other.photons
            if N not in self.sector_states:
                return 0
            return self.sector_amplitudes[N].conjugate() * self.sector_states[N].dot(other)

        if isinstance(other, PureStateMultiSector):

            total = 0
            common_sectors = set(self.sector_states) & set(other.sector_states)

            for N in common_sectors:
                overlap = self.sector_states[N].dot(other.sector_states[N])
                total += (
                    self.sector_amplitudes[N].conjugate() * other.sector_amplitudes[N] * overlap
                )
            return total
        return NotImplementedError

    # def dot_sector(self, other):
    #     """Return a dictionary with the overlap in each sector."""

    #     if isinstance(other, PureState):
    #         N = other.photons
    #         if N not in self.sector_states:
    #             return {}

    #         return {N: self.sector_amplitudes[N].conjugate() * self.sector_states[N].dot(other)}

    #     if isinstance(other, PureStateMultiSector):

    #         overlaps = {}
    #         common_sectors = set(self.sector_states) & set(other.sector_states)

    #         for N in common_sectors:
    #             overlap = self.sector_states[N].dot(other.sector_states[N])
    #             overlaps[N] = (
    #                 self.sector_amplitudes[N].conjugate() * other.sector_amplitudes[N] * overlap
    #             )
    #         return overlaps
    #     raise NotImplementedError

    def norm(self):
        return np.sqrt(self.dot(self))


class SingleModeSuperposition(PureStateMultiSector):
    """
    Superposition of photon numbers in a single mode up to a cutoff.

    Represents:

        sum_{n=0}^{cutoff} c_n |n>_mode ⊗ |0>_others
    """

    def __init__(
        self,
        coef_vector: list[complex] | np.ndarray,
    ):
        """
        Args:
            coefficients: [c0, c1, ..., cN]
            mode: which mode carries the photons
            modes: total number of modes
        """

        self.cutoff = len(coef_vector) - 1
        coef_vector = np.asarray(coef_vector, dtype=np.complex128)
        self.amp_vector = coef_vector / np.linalg.norm(coef_vector)

        focks = []
        amps = []

        for n, amp in enumerate(self.amp_vector):
            focks.append((n,))
            amps.append(amp)

        super().__init__(focks, amps)

    def evolution_unitary(self, U):
        return SingleModeSuperposition(U @ self.amp_vector)
