from __future__ import annotations

from functools import partial
from typing import TYPE_CHECKING, Any, TypedDict, TypeVar, cast

import awkward as ak
import awkward.operations.str as ak_str
import numba as nb
import numpy as np
from attrs import define
from awkward.contents import NumpyArray
from numpy.typing import NDArray
from phantom import Phantom
from seqpro.rag import RDTYPE, Ragged, is_rag_dtype

from ._torch import TORCH_AVAILABLE
from ._types import AnnotatedHaps

if TORCH_AVAILABLE or TYPE_CHECKING:
    import torch
    from torch.nested import nested_tensor_from_jagged as nt_jag

__all__ = ["Ragged", "RaggedIntervals", "RaggedTracks"]
INTERVAL_DTYPE = np.dtype(
    [("start", np.int32), ("end", np.int32), ("value", np.float32)], align=True
)


@define
class RaggedIntervals:
    starts: Ragged[np.int32]
    ends: Ragged[np.int32]
    values: Ragged[np.float32]

    def __getitem__(self, idx) -> RaggedIntervals:
        out = RaggedIntervals(self.starts[idx], self.ends[idx], self.values[idx])  # type: ignore
        return out

    @property
    def shape(self):
        """Shape of the haplotypes and all annotations."""
        return self.values.shape

    def to_padded(
        self, start: int, end: int, value: float
    ) -> tuple[NDArray[np.int32], NDArray[np.int32], NDArray[np.float32]]:
        """Convert this RaggedIntervals to a tuple of rectilinear arrays by right-padding each entry with appropriate values.
        The final axis will have the maximum length across all entries."""
        starts = to_padded(self.starts, start)
        ends = to_padded(self.ends, end)
        values = to_padded(self.values, value)
        return starts, ends, values

    def reshape(self, shape: int | tuple[int, ...]) -> RaggedIntervals:
        """Reshape the haplotypes and all annotations.

        Parameters
        ----------
        shape
            New shape for the haplotypes and all annotations. The total number of elements
            must remain the same.
        """
        return RaggedIntervals(
            self.starts.reshape(shape),
            self.ends.reshape(shape),
            self.values.reshape(shape),
        )

    def squeeze(self, axis: int | tuple[int, ...] | None = None) -> RaggedIntervals:
        """Squeeze the haplotypes and all annotations along the specified axis.

        Parameters
        ----------
        axis
            Axis or axes to squeeze. If None, all axes of length 1 are squeezed.
        """
        return RaggedIntervals(
            self.starts.squeeze(axis),  # type: ignore
            self.ends.squeeze(axis),  # type: ignore
            self.values.squeeze(axis),  # type: ignore
        )

    def to_fixed_shape(
        self, shape: tuple[int, ...]
    ) -> tuple[NDArray[np.int32], NDArray[np.int32], NDArray[np.float32]]:
        """If all entries in the ragged array have the same shape, convert to a rectilinear shape.

        Parameters
        ----------
        shape
            Shape to convert to, including the length axis. The total number of elements must remain the same.
        """
        starts = self.starts.data.reshape(shape)
        ends = self.ends.data.reshape(shape)
        values = self.values.data.reshape(shape)
        return starts, ends, values

    def to_packed(self) -> RaggedIntervals:
        """Apply :func:`ak.to_packed` to all arrays."""
        starts = ak.to_packed(self.starts)
        ends = ak.to_packed(self.ends)
        values = ak.to_packed(self.values)
        return RaggedIntervals(starts, ends, values)

    def to_nested_tensor_batch(
        self, device: str | torch.device = "cpu"
    ) -> list[RagItvBatch]:
        out = []
        n_tracks = cast(int, self.values.shape[1])
        for t in range(n_tracks):
            # (batch tracks ... ~itv) -> (batch ... ~itv)
            starts = ak.to_packed(self.starts[:, t])
            ends = ak.to_packed(self.ends[:, t])
            values = ak.to_packed(self.values[:, t])

            offsets = torch.from_numpy(values.offsets.astype(np.int32)).to(device)
            max_len = int(values.lengths.max())

            starts = torch.from_numpy(starts.data).to(device)
            starts = nt_jag(starts, offsets)
            ends = torch.from_numpy(ends.data).to(device)
            ends = nt_jag(ends, offsets)
            values = torch.from_numpy(values.data.astype(np.float32)).to(device)
            values = nt_jag(values, offsets)

            out.append(
                RagItvBatch(starts=starts, ends=ends, values=values, max_seqlen=max_len)
            )

        return out

    def prepend_pad_itv(
        self, start: int = -1, end: int = -1, value: float = 0.0
    ) -> RaggedIntervals:
        """Prepend a pad interval so that every group is guaranteed to have at least 1 interval.

        Parameters
        ----------
        start
            The start position to use for the pad interval
        end
            The end position to use for the pad interval
        value
            The value to use for the pad interval
        """
        b, t, *_ = self.values.shape
        b = cast(int, b)
        t = cast(int, t)

        pad_start = ak.from_numpy(
            np.full((b, t, 1), start, np.int32), regulararray=True
        )
        # (b t ~v)
        new_starts = ak.concatenate([pad_start, self.starts], axis=2)
        pad_end = ak.from_numpy(np.full((b, t, 1), end, np.int32), regulararray=True)
        # (b t ~v)
        new_ends = ak.concatenate([pad_end, self.ends], axis=2)
        pad_value = ak.from_numpy(
            np.full((b, t, 1), value, np.float32), regulararray=True
        )
        # (b t ~v)
        new_values = ak.concatenate([pad_value, self.values], axis=2)

        return RaggedIntervals(Ragged(new_starts), Ragged(new_ends), Ragged(new_values))


class RagItvBatch(TypedDict):
    """Dictionary of nested tensors."""

    starts: torch.Tensor
    ends: torch.Tensor
    values: torch.Tensor
    max_seqlen: int


class RaggedSeqs(
    Ragged[np.bytes_], Phantom, predicate=partial(is_rag_dtype, dtype=np.bytes_)
): ...


class RaggedTracks(
    Ragged[np.float32], Phantom, predicate=partial(is_rag_dtype, dtype=np.float32)
): ...


@define
class RaggedAnnotatedHaps:
    """Ragged version of :class:`AnnotatedHaps`."""

    haps: Ragged[np.bytes_]
    """Haplotypes with dtype S1."""
    var_idxs: Ragged[np.int32]
    """Variant indices for each position in the haplotypes. A value of -1 indicates no variant was applied at the position."""
    ref_coords: Ragged[np.int32]
    """Reference coordinates for each position in haplotypes."""

    @property
    def shape(self):
        """Shape of the haplotypes and all annotations."""
        return self.haps.shape

    def to_padded(self) -> AnnotatedHaps:
        """Convert this Ragged array to a rectilinear array by right-padding each entry with appropriate values.
        The final axis will have the maximum length across all entries."""
        haps = to_padded(self.haps, b"N")
        var_idxs = to_padded(self.var_idxs, -1)
        ref_coords = to_padded(self.ref_coords, np.iinfo(self.ref_coords.dtype).max)
        return AnnotatedHaps(haps, var_idxs, ref_coords)

    def reshape(self, shape: int | tuple[int, ...]) -> RaggedAnnotatedHaps:
        """Reshape the haplotypes and all annotations.

        Parameters
        ----------
        shape
            New shape for the haplotypes and all annotations. The total number of elements
            must remain the same.
        """
        return RaggedAnnotatedHaps(
            self.haps.reshape(shape),
            self.var_idxs.reshape(shape),
            self.ref_coords.reshape(shape),
        )

    def squeeze(self, axis: int | tuple[int, ...] | None = None) -> RaggedAnnotatedHaps:
        """Squeeze the haplotypes and all annotations along the specified axis.

        Parameters
        ----------
        axis
            Axis or axes to squeeze. If None, all axes of length 1 are squeezed.
        """
        return RaggedAnnotatedHaps(
            self.haps.squeeze(axis),
            self.var_idxs.squeeze(axis),
            self.ref_coords.squeeze(axis),
        )

    def to_numpy(self) -> AnnotatedHaps:
        """If all entries in the ragged array have the same shape, convert to a rectilinear shape.

        Parameters
        ----------
        shape
            Shape to convert to, including the length axis. The total number of elements must remain the same.
        """
        haps = self.haps.to_numpy()
        var_idxs = self.var_idxs.to_numpy()
        ref_coords = self.ref_coords.to_numpy()
        return AnnotatedHaps(haps, var_idxs, ref_coords)


def to_padded(rag: Ragged[RDTYPE], pad_value: Any) -> NDArray[RDTYPE]:
    """Convert this Ragged array to a rectilinear array by right-padding each entry with a value.
    The ragged axis will be padded to have the maximum length across all entries.

    Parameters
    ----------
    rag
        Ragged array to pad.
    pad_value
        Value to pad the entries with.

    Returns
    -------
        Padded array.
    """
    length = int(rag.lengths.max())
    if is_rag_dtype(rag, np.bytes_):
        rag = ak_str.rpad(rag, length, pad_value)
    else:
        rag = ak.pad_none(rag, length, clip=True)
        rag = ak.fill_none(rag, pad_value)
    arr = Ragged(rag).to_numpy()
    return arr


_COMP = np.frombuffer(bytes.maketrans(b"ACGT", b"TGCA"), np.uint8)


@nb.vectorize(["u1(u1)"], nopython=True)
def ufunc_comp_dna(seq: NDArray[np.uint8]) -> NDArray[np.uint8]:
    return _COMP[seq]


def _ak_comp_dna_helper(layout, **kwargs):
    if layout.is_numpy:
        return NumpyArray(
            ufunc_comp_dna(layout.data),  # type: ignore
            parameters=layout.parameters,
        )


T = TypeVar("T", bound=ak.Array)


def reverse_complement(arr: T) -> T:
    og_type = type(arr)
    arr = ak.to_packed(arr)
    arr = ak_str.reverse(ak.transform(_ak_comp_dna_helper, arr))
    return og_type(arr)
