from bloqade.builder.typing import ScalarType
from bloqade.builder.start import ProgramStart
from bloqade.ir.scalar import Scalar, Literal, cast
from bloqade.ir.tree_print import Printer

from pydantic.v1.dataclasses import dataclass
from beartype.typing import List, Tuple, Generator, Union, Optional
from beartype.door import is_bearable
from beartype import beartype
from enum import Enum
from numpy.typing import NDArray
from bloqade.submission.ir.capabilities import QuEraCapabilities
from bloqade.visualization import get_atom_arrangement_figure
from bloqade.visualization import display_ir

from beartype.vale import Is
from typing import Annotated
import plotext as pltxt
import sys
import numpy as np


def check_position_array(array):
    return (
        array.ndim == 2
        and (
            np.issubdtype(array.dtype, np.floating)
            or np.issubdtype(array.dtype, np.integer)
        )
        and array.shape[1] == 2
    )


def check_bool_array(array):
    return array.ndim == 1 and np.issubdtype(array.dtype, np.bool_)


PositionArray = Annotated[np.ndarray, Is[check_position_array]]
BoolArray = Annotated[np.ndarray, Is[check_bool_array]]


class SiteFilling(int, Enum):
    filled = 1
    vacant = 0

    def print_node(self) -> str:
        return self.name

    def children(self) -> List:
        return []


@dataclass(frozen=True)
class LocationInfo:
    position: Tuple[Scalar, Scalar]
    filling: SiteFilling

    @beartype
    @staticmethod
    def create(position: Tuple[ScalarType, ScalarType], filled: bool):
        if filled:
            filling = SiteFilling.filled
        else:
            filling = SiteFilling.vacant

        position = tuple(cast(ele) for ele in position)

        return LocationInfo(position, filling)

    def print_node(self) -> str:
        return f"Location: {self.filling.name}"

    def children(self) -> List[Scalar]:
        return {"x": self.position[0], "y": self.position[1]}


@dataclass(init=False)
class AtomArrangement(ProgramStart):
    def __str__(self) -> str:
        def is_literal(x):
            return isinstance(x, Literal)

        plot_results = all(
            is_literal(info.position[0]) and is_literal(info.position[1])
            for info in self.enumerate()
        )

        if plot_results:
            xs = [float(info.position[0].value) for info in self.enumerate()]
            ys = [float(info.position[1].value) for info in self.enumerate()]
            filling = [bool(info.filling.value) for info in self.enumerate()]
            xs_filled = [x for x, filled in zip(xs, filling) if filled]
            ys_filled = [y for y, filled in zip(ys, filling) if filled]
            xs_vacant = [x for x, filled in zip(xs, filling) if not filled]
            ys_vacant = [y for y, filled in zip(ys, filling) if not filled]

            pltxt.clear_figure()
            pltxt.limit_size(False, False)
            pltxt.plot_size(80, 24)
            pltxt.canvas_color("default")
            pltxt.axes_color("default")
            pltxt.ticks_color("white")
            pltxt.title("Atom Positions")
            pltxt.xlabel("x (um)")
            pltxt.ylabel("y (um)")

            marker = "."
            if sys.stdout.encoding.lower().startswith("utf"):
                marker = "dot"

            pltxt.scatter(xs_filled, ys_filled, color=(100, 55, 255), marker=marker)
            if len(xs_vacant) > 0:
                pltxt.scatter(
                    xs_vacant, ys_vacant, color="white", label="vacant", marker=marker
                )

            return pltxt.build()
        else:
            ph = Printer()
            ph.print(self)
            return ph.get_value()

    def print_node(self) -> str:
        return "AtomArrangement"

    def children(self) -> List[LocationInfo]:
        return list(self.enumerate())

    def enumerate(self) -> Generator[LocationInfo, None, None]:
        """enumerate all locations in the register."""
        raise NotImplementedError

    def figure(self, fig_kwargs=None, **assignments):
        """obtain a figure object from the atom arrangement."""
        return get_atom_arrangement_figure(self, fig_kwargs, **assignments)

    def show(self, **assignments) -> None:
        display_ir(self, assignments)

    def rydberg_interaction(self, **assignments) -> NDArray:
        """calculate the Rydberg interaction matrix.

        Args:
            **assignments: the values to assign to the variables in the register.

        Returns:
            NDArray: the Rydberg interaction matrix in the lower triangular form.

        """

        from bloqade.constants import RB_C6

        # calculate the Interaction matrix
        V_ij = np.zeros((self.n_sites, self.n_sites))
        for i, site_i in enumerate(self.enumerate()):
            pos_i = np.array([float(ele(**assignments)) for ele in site_i.position])

            for j, site_j in enumerate(self.enumerate()):
                if j >= i:
                    break  # enforce lower triangular form

                pos_j = np.array([float(ele(**assignments)) for ele in site_j.position])
                r_ij = np.linalg.norm(pos_i - pos_j)

                V_ij[i, j] = RB_C6 / r_ij**6

        return V_ij

    @property
    def n_atoms(self) -> int:
        """number of atoms (filled sites) in the register."""
        raise NotImplementedError

    @property
    def n_sites(self) -> int:
        """number of sites in the register."""
        raise NotImplementedError

    @property
    def n_vacant(self) -> int:
        """number of vacant sites in the register."""
        raise NotImplementedError

    @property
    def n_dims(self) -> int:
        """number of dimensions in the register."""
        raise NotImplementedError

    @beartype
    def scale(self, scale: ScalarType):
        """
        Scale the geometry of your atoms.

        ### Usage Example:
        ```
        >>> reg = start.add_position([(0,0), (1,1)])
        # atom positions are now (0,0), (2,2)
        >>> new_reg = reg.scale(2)
        # you may also use scale on pre-defined geometries
        >>> from bloqade.atom_arrangement import Chain
        # atoms in the chain will now be 2 um apart versus
        # the default 1 um
        >>> Chain(11).scale(2)
        ```

        - Next possible steps are:
        - Continuing to build your geometry via:
            - `...add_position(positions).add_position(positions)`:
                to add more positions
            - `...add_position(positions).apply_defect_count(n_defects)`:
            to randomly drop out n_atoms
            - `...add_position(positions).apply_defect_density(defect_probability)`:
            to drop out atoms with a certain probability
            - `...add_position(positions).scale(scale)`: to scale the geometry
        - Targeting a level coupling once you're done with the atom geometry:
            - `...add_position(positions).rydberg`:
            to specify Rydberg coupling
            - `...add_position(positions).hyperfine`:
            to specify Hyperfine coupling
        - Visualizing your atom geometry:
            - `...add_position(positions).show()`:
            shows your geometry in your web browser

        """

        scale = cast(scale)
        location_list = []
        for location_info in self.enumerate():
            x, y = location_info.position
            new_position = (scale * x, scale * y)
            location_list.append(
                LocationInfo.create(new_position, bool(location_info.filling.value))
            )

        return ListOfLocations(location_list)

    ## enforce input is single tuple and filling is single boolean
    def add_position_single_tupe(
        self, position: Tuple[ScalarType, ScalarType], filling: Optional[bool] = None
    ):
        if filling is None:
            filling = True

        location_list = list(self.enumerate())
        location_list.append(LocationInfo.create(position, filling))

        return ListOfLocations(location_list)

    ## enforce input is list of tuples and filling is list of booleans or none
    def add_position_list_tuples(  # noqa: F811
        self,
        position: List[Tuple[ScalarType, ScalarType]],
        filling: Optional[List[bool]] = None,
    ):
        location_list = list(self.enumerate())

        assert (
            len(position) == len(filling) if filling else True
        ), "Length of positions and filling must be the same"

        if filling:
            for position, filling in zip(position, filling):
                location_list.append(LocationInfo.create(position, filling))

        else:
            for position in position:
                location_list.append(LocationInfo.create(position, True))

        return ListOfLocations(location_list)

    ## enforce input is numpy array and filling is numpy array or none
    def add_position_ndarray(  # noqa: F811
        self, position: PositionArray, filling: Optional[BoolArray] = None
    ):
        return self.add_position(
            list(map(tuple, position.tolist())),
            filling.tolist() if filling is not None else None,
        )

    def add_position(
        self,
        position: Union[
            PositionArray,
            List[Tuple[ScalarType, ScalarType]],
            Tuple[ScalarType, ScalarType],
        ],
        filling: Optional[Union[BoolArray, List[bool], bool]] = None,
    ) -> "ListOfLocations":
        """
        Add a position or multiple positions to a pre-existing geometry.

        `add_position` is capable of accepting:
        - A single tuple for one atom coordinate: `(1.0, 2.5)`
        - A list of tuples: `[(0.0, 1.0), (2.0,1.5), etc.]
        - A numpy array of shape (N, 2) where N is the number of atoms

        You may also intersperse variables anywhere a value may be present.

        You can also pass in an optional argument which determines the atom "filling"
        (whether or not at a specified coordinate an atom should be present).

        ### Usage Example:
        ```
        # single coordinate
        >>> reg = start.add_position((0,0))
        # you may chain add_position calls
        >>> reg_plus_two = reg.add_position([(2,2),(5.0, 2.1)])
        # you can add variables anywhere a value may be present
        >>> reg_with_var = reg_plus_two.add_position(("x", "y"))
        # and specify your atom fillings
        >>> reg_with_filling = reg_with_var.add_position([(3.1, 0.0), (4.1, 2.2)],
        [True, False])
        # alternatively you could use one boolean to specify
        # all coordinates should be empty/filled
        >>> reg_with_more_filling = reg_with_filling.add_positions([(3.1, 2.9),
        (5.2, 2.2)], False)
        ```

        - Next possible steps are:
        - Continuing to build your geometry via:
            - `...add_position(positions).add_position(positions)`:
                to add more positions
            - `...add_position(positions).apply_defect_count(n_defects)`:
            to randomly drop out n_atoms
            - `...add_position(positions).apply_defect_density(defect_probability)`:
            to drop out atoms with a certain probability
            - `...add_position(positions).scale(scale)`: to scale the geometry
        - Targeting a level coupling once you're done with the atom geometry:
            - `...add_position(positions).rydberg`: to specify Rydberg coupling
            - `...add_position(positions).hyperfine`: to specify Hyperfine coupling
        - Visualizing your atom geometry:
            - `...add_position(positions).show()`:
            shows your geometry in your web browser

        """

        if is_bearable(position, PositionArray) and is_bearable(
            filling, Optional[BoolArray]
        ):
            return self.add_position_ndarray(position, filling)
        elif is_bearable(position, List[Tuple[ScalarType, ScalarType]]) and is_bearable(
            filling, Optional[List[bool]]
        ):
            return self.add_position_list_tuples(position, filling)
        elif is_bearable(position, Tuple[ScalarType, ScalarType]) and is_bearable(
            filling, Optional[bool]
        ):
            return self.add_position_single_tupe(position, filling)
        else:
            raise TypeError("Invalid input types for add_position provided!")

    @beartype
    def apply_defect_count(
        self, n_defects: int, rng: np.random.Generator = np.random.default_rng()
    ):
        """
        Drop `n_defects` atoms from the geometry randomly. Internally this occurs
        by setting certain sites to have a SiteFilling set to false indicating
        no atom is present at the coordinate.

        A default numpy-based Random Number Generator is used but you can
        explicitly override this by passing in your own.

        ### Usage Example:

        ```
        >>> from bloqade.atom_arrangement import Chain
        >>> import numpy as np
        # set a custom seed for a numpy-based RNG
        >>> custom_rng = np.random.default_rng(888)
        # randomly remove two atoms from the geometry
        >>> reg = Chain(11).apply_defect_count(2, custom_rng)
        # you may also chain apply_defect_count calls
        >>> reg.apply_defect_count(2, custom_rng)
        # you can also use apply_defect_count on custom geometries
        >>> from bloqade import start
        >>> start.add_position([(0,0), (1,1)]).apply_defect_count(1, custom_rng)
        ```

        - Next possible steps are:
        - Continuing to build your geometry via:
            - `...apply_defect_count(defect_counts).add_position(positions)`:
                to add more positions
            - `...apply_defect_count(defect_counts)
                .apply_defect_count(n_defects)`: to randomly drop out n_atoms
            - `...apply_defect_count(defect_counts)
                .apply_defect_density(defect_probability)`:
                to drop out atoms with a certain probability
            - `...apply_defect_count(defect_counts).scale(scale)`:
                to scale the geometry
        - Targeting a level coupling once you're done with the atom geometry:
            - `...apply_defect_count(defect_counts).rydberg`: to specify
                Rydberg coupling
            - `...apply_defect_count(defect_counts).hyperfine`:
                to specify Hyperfine coupling
        - Visualizing your atom geometry:
            - `...apply_defect_count(defect_counts).show()`:
                shows your geometry in your web browser
        """

        location_list = []
        for location_info in self.enumerate():
            location_list.append(location_info)

        filled_sites = []

        for index, location_info in enumerate(location_list):
            if location_info.filling is SiteFilling.filled:
                filled_sites.append(index)

        if n_defects >= len(filled_sites):
            raise ValueError(
                f"n_defects {n_defects} must be less than the number of filled sites "
                f"({len(filled_sites)})"
            )

        for _ in range(n_defects):
            index = rng.choice(filled_sites)
            location_list[index] = LocationInfo.create(
                location_list[index].position,
                (False if location_list[index].filling is SiteFilling.filled else True),
            )
            filled_sites.remove(index)

        return ListOfLocations(location_list)

    @beartype
    def apply_defect_density(
        self,
        defect_probability: float,
        rng: np.random.Generator = np.random.default_rng(),
    ):
        """
        Drop atoms randomly with `defect_probability` probability (range of 0 to 1).
        Internally this occurs by setting certain sites to have a SiteFilling
        set to false indicating no atom is present at the coordinate.

        A default numpy-based Random Number Generator is used but you can
        explicitly override this by passing in your own.

        ### Usage Example:

        ```
        >>> from bloqade.atom_arrangement import Chain
        >>> import numpy as np
        # set a custom seed for a numpy-based RNG
        >>> custom_rng = np.random.default_rng(888)
        # randomly remove two atoms from the geometry
        >>> reg = Chain(11).apply_defect_density(0.2, custom_rng)
        # you may also chain apply_defect_density calls
        >>> reg.apply_defect_count(0.1, custom_rng)
        # you can also use apply_defect_density on custom geometries
        >>> from bloqade import start
        >>> start.add_position([(0,0), (1,1)])
        .apply_defect_density(0.5, custom_rng)
        ```

        - Next possible steps are:
        - Continuing to build your geometry via:
            - `...apply_defect_count(defect_counts).add_position(positions)`:
            to add more positions
            - `...apply_defect_count(defect_counts).apply_defect_count(n_defects)`:
            to randomly drop out n_atoms
            - `...apply_defect_count(defect_counts)
            .apply_defect_density(defect_probability)`:
            to drop out atoms with a certain probability
            - `...apply_defect_count(defect_counts).scale(scale)`:
            to scale the geometry
        - Targeting a level coupling once you're done with the atom geometry:
            - `...apply_defect_count(defect_counts).rydberg`:
            to specify Rydberg coupling
            - `...apply_defect_count(defect_counts).hyperfine`:
            to specify Hyperfine coupling
        - Visualizing your atom geometry:
            - `...apply_defect_count(defect_counts).show()`:
            shows your geometry in your web browser
        """

        p = min(1, max(0, defect_probability))
        location_list = []

        for location_info in self.enumerate():
            if rng.random() < p:
                location_list.append(
                    LocationInfo.create(
                        location_info.position,
                        (
                            False
                            if location_info.filling is SiteFilling.filled
                            else True
                        ),
                    )
                )
            else:
                location_list.append(location_info)

        return ListOfLocations(location_list=location_list)

    def remove_vacant_sites(self):
        new_locations = []
        for location_info in self.enumerate():
            if location_info.filling is SiteFilling.filled:
                new_locations.append(location_info)

        return ListOfLocations(new_locations)


@dataclass
class ParallelRegister(ProgramStart):
    atom_arrangement: AtomArrangement
    cluster_spacing: Scalar

    @property
    def n_atoms(self):
        return self.atom_arrangement.n_atoms

    @property
    def n_sites(self):
        return self.atom_arrangement.n_sites

    @property
    def n_vacant(self):
        return self.atom_arrangement.n_vacant

    @property
    def n_dims(self):
        return self.atom_arrangement.n_dims

    def __str__(self):
        return "ParallelRegister:\n" + self.atom_arrangement.__str__()

    def _repr_pretty_(self, p, _):
        p.text(str(self))

    def _compile_to_list(
        self, __capabilities: Optional[QuEraCapabilities] = None, **assignments
    ):
        from bloqade.compiler.rewrite.common import AssignBloqadeIR
        from bloqade.compiler.codegen.hardware import GenerateLattice
        from bloqade.submission.capabilities import get_capabilities

        lattice_data = GenerateLattice(__capabilities or get_capabilities()).emit(
            AssignBloqadeIR(assignments).emit(self)
        )

        list_of_locations = ListOfLocations()
        for site, filling in zip(lattice_data.sites, lattice_data.filling):
            list_of_locations = list_of_locations.add_position(site, filling == 1)

        return list_of_locations

    def figure(
        self,
        fig_kwargs=None,
        capabilities: Optional[QuEraCapabilities] = None,
        **assignments,
    ):
        return self._compile_to_list(capabilities).figure(fig_kwargs)

    def show(self, **assignments) -> None:
        display_ir(self, assignments)


@dataclass(init=False)
class ParallelRegisterInfo:
    """ParallelRegisterInfo"""

    register_locations: List[List[Scalar]]
    register_filling: List[int]
    shift_vectors: List[List[Scalar]]

    def __init__(self, parallel_register: ParallelRegister):
        atom_arrangement = parallel_register.atom_arrangement
        cluster_spacing = parallel_register.cluster_spacing

        if atom_arrangement.n_atoms > 0:
            # calculate bounding box
            # of this register
            location_iter = atom_arrangement.enumerate()
            (x, y) = next(location_iter).position
            x_min = x
            x_max = x
            y_min = y
            y_max = y

            for location_info in location_iter:
                (x, y) = location_info.position
                x_min = x.min(x_min)
                x_max = x.max(x_max)
                y_min = y.min(y_min)
                y_max = y.max(y_max)

            shift_x = (x_max - x_min) + cluster_spacing
            shift_y = (y_max - y_min) + cluster_spacing

            register_locations = [
                list(location_info.position)
                for location_info in atom_arrangement.enumerate()
            ]
            register_filling = [
                location_info.filling.value
                for location_info in atom_arrangement.enumerate()
            ]
            shift_vectors = [[shift_x, cast(0)], [cast(0), shift_y]]
        else:
            raise ValueError("No locations to parallelize.")

        self.register_locations = register_locations
        self.register_filling = register_filling
        self.shift_vectors = shift_vectors


@dataclass(init=False)
class ListOfLocations(AtomArrangement):
    location_list: List[LocationInfo]

    @beartype
    def __init__(
        self,
        location_list: List[Union[LocationInfo, Tuple[ScalarType, ScalarType]]] = [],
    ):
        self.location_list = []
        for ele in location_list:
            if isinstance(ele, LocationInfo):
                self.location_list.append(ele)
            else:
                self.location_list.append(LocationInfo.create(ele, True))

        if self.location_list:
            self.__n_atoms = sum(
                1 for loc in self.location_list if loc.filling == SiteFilling.filled
            )
            self.__n_sites = len(self.location_list)
            self.__n_vacant = self.__n_sites - self.__n_atoms
            self.__n_dims = len(self.location_list[0].position)
        else:
            self.__n_sites = 0
            self.__n_atoms = 0
            self.__n_dims = None

        super().__init__()

    def _repr_pretty_(self, p, _):
        p.text(str(self))

    @property
    def n_atoms(self):
        return self.__n_atoms

    @property
    def n_sites(self):
        return self.__n_sites

    @property
    def n_vacant(self):
        return self.__n_vacant

    @property
    def n_dims(self):
        return self.__n_dims

    def enumerate(self):
        return iter(self.location_list)

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