from ast import Raise
import numpy as np
import pandas as pd
import pyranges as pr
from typing import Any, Union, Optional, List
from typing import Dict, Iterable, Sequence, Mapping
from os import PathLike
from collections import OrderedDict
from functools import singledispatch
from pandas.api.types import is_string_dtype
from copy import deepcopy
from ._SeqDataset import SeqDataset
Index1D = Union[slice, int, str, np.int64, np.ndarray]


class SeqData:
    """SeqData object used to containerize and store data for EUGENe workflows.
    
    Attributes
    ----------
    seqs : np.ndarray
        Numpy array of sequences.
    names : np.ndarray
        Numpy array of names.
    rev_seqs : np.ndarray
        Numpy array of reverse complement sequences.
    seqs_annot : pr.PyRanges
        PyRanges object or dict of sequences annotations.
    pos_annot : pr.PyRanges
        PyRanges object or dict of positions annotations.
    ohe_seqs : np.ndarray
        Numpy array of one-hot encoded sequences.
    ohe_rev_seqs : np.ndarray
        Numpy array of one-hot encoded reverse complement sequences.
    seqsm : np.ndarray
        Numpy array of sequences or dict of sequences.
    uns : dict
        Dict of additional/unstructured information.
    seqsidx : Index1D
        Index of sequences to use.
    """

    def __init__(
        self,
        seqs: np.ndarray = None,
        names: np.ndarray = None,
        rev_seqs: np.ndarray = None,
        seqs_annot: Optional[Union[pd.DataFrame, Mapping[str, Iterable[Any]]]] = None,
        pos_annot: Union[pr.PyRanges, Dict, str] = None,
        ohe_seqs: np.ndarray = None,
        ohe_rev_seqs: np.ndarray = None,
        seqsm: Optional[Union[np.ndarray, Mapping[str, Sequence[Any]]]] = None,
        uns: Optional[Mapping[str, Any]] = None,
        seqidx: Index1D = None,
    ):

        if seqs is not None:
            self.seqidx = (
                range(seqs.shape[0]) 
                if seqidx is None or len(seqidx) == 0 
                else seqidx
            )
        else:
            self.seqidx = (
                range(ohe_seqs.shape[0])
                if seqidx is None or len(seqidx) == 0
                else seqidx
            )

        # Sequence representations
        self.seqs = np.array(seqs[self.seqidx]) if seqs is not None else None
        self.names = np.array(names[self.seqidx]) if names is not None else None
        self.rev_seqs = np.array(rev_seqs[self.seqidx]) if rev_seqs is not None else None
        self.ohe_seqs = np.array(ohe_seqs[self.seqidx]) if ohe_seqs is not None else None
        self.ohe_rev_seqs = np.array(ohe_rev_seqs[self.seqidx]) if ohe_rev_seqs is not None else None

        # n_obs
        self._n_obs = len(self.seqs) if self.seqs is not None else len(self.ohe_seqs)

        # seq_annot (handled by gen dataframe)
        if isinstance(self.seqidx, slice):
            self.seqs_annot = _gen_dataframe(
                seqs_annot, self._n_obs, ["obs_names", "row_names"]
            )[self.seqidx]
        elif type(self.seqidx[0]) in [bool, np.bool_]:
            self.seqs_annot = _gen_dataframe(
                seqs_annot, self._n_obs, ["obs_names", "row_names"]
            ).loc[self.seqidx]
        else:
            self.seqs_annot = _gen_dataframe(
                seqs_annot, self._n_obs, ["obs_names", "row_names"]
            ).iloc[self.seqidx]

        # pos_annot
        if isinstance(pos_annot, dict):
            self.pos_annot = pr.from_dict(pos_annot)
        elif isinstance(pos_annot, str):
            self.pos_annot = pr.read_bed(pos_annot)
        else:
            self.pos_annot = pos_annot

        # uns
        self.uns = uns or OrderedDict()

        # seqsm TODO: Think about consequences of making obsm a group in hdf
        if seqsm is not None:
            seqsm = seqsm.copy()
            for key in seqsm:
                seqsm[key] = seqsm[key][self.seqidx]
        self.seqsm = convert_to_dict(seqsm)
            

    @property
    def seqs(self) -> np.ndarray:
        """np.ndarray: Numpy array of string representation of sequences."""
        return self._seqs

    @seqs.setter
    def seqs(self, seqs: np.ndarray):
        self._seqs = seqs

    @property
    def names(self) -> np.ndarray:
        """np.ndarray: Numpy array of names or identifiers of sequences."""
        return self._names

    @names.setter
    def names(self, names: np.ndarray):
        self._names = names

    @property
    def n_obs(self) -> int:
        """int: Number of sequences contained in the object."""
        return self._n_obs

    @property
    def rev_seqs(self) -> np.ndarray:
        """np.ndarray: Numpy array of reverse complement sequences."""
        return self._rev_seqs

    @rev_seqs.setter
    def rev_seqs(self, rev_seqs: np.ndarray):
        self._rev_seqs = rev_seqs

    @property
    def ohe_seqs(self) -> np.ndarray:
        """np.ndarray: Numpy array of one-hot encoded sequences."""
        return self._ohe_seqs

    @ohe_seqs.setter
    def ohe_seqs(self, ohe_seqs: np.ndarray):
        self._ohe_seqs = ohe_seqs

    @property
    def seqs_annot(self) -> pd.DataFrame:
        """pd.DataFrame: Pandas dataframe of per sequence annotations."""
        return self._seqs_annot

    @seqs_annot.setter
    def seqs_annot(self, seqs_annot: Union[pd.DataFrame, Mapping[str, Iterable[Any]]]):
        self._seqs_annot = _gen_dataframe(
            seqs_annot, self._n_obs, ["obs_names", "row_names"]
        )

    @property
    def pos_annot(self) -> pr.PyRanges:
        """pr.PyRanges: PyRanges object of per sequence annotations.""" 
        return self._pos_annot

    @pos_annot.setter
    def pos_annot(self, pos_annot: pr.PyRanges):
        self._pos_annot = pos_annot

    @property
    def ohe_rev_seqs(self) -> np.ndarray:
        """np.ndarray: Numpy array of one-hot encoded reverse complement sequences."""
        return self._ohe_rev_seqs

    @ohe_rev_seqs.setter
    def ohe_rev_seqs(self, ohe_rev_seqs: np.ndarray):
        self._ohe_rev_seqs = ohe_rev_seqs

    @property
    def seqsm(self) -> Mapping[str, Sequence[Any]]:
        """Mapping[str, Sequence[Any]]: Dictionary of multidimensional sequence representations."""
        return self._seqsm

    @seqsm.setter
    def seqsm(self, seqsm: Mapping[str, Sequence[Any]]):
        self._seqsm = seqsm

    @property
    def uns(self) -> Mapping[str, Any]:
        """Mapping[str, Any]: Dictionary of unstructured annotations."""
        return self._uns

    @uns.setter
    def uns(self, uns: Mapping[str, Any]):
        self._uns = uns

    def __getitem__(self, index):
        """Get item from data. Defines slicing of object."""
        if isinstance(index, str):
            return self.seqs_annot[index]
        elif isinstance(index, slice):
            index = np.arange(self.n_obs)[index]
            return SeqData(
                seqs=self.seqs,
                names=self.names,
                rev_seqs=self.rev_seqs,
                seqs_annot=self.seqs_annot,
                pos_annot=self.pos_annot,
                ohe_seqs=self.ohe_seqs,
                ohe_rev_seqs=self.ohe_rev_seqs,
                seqsm=self.seqsm,
                uns=self.uns,
                seqidx=index,
            )
        else:
            return SeqData(
                seqs=self.seqs,
                names=self.names,
                rev_seqs=self.rev_seqs,
                seqs_annot=self.seqs_annot,
                pos_annot=self.pos_annot,
                ohe_seqs=self.ohe_seqs,
                ohe_rev_seqs=self.ohe_rev_seqs,
                seqsm=self.seqsm,
                uns=self.uns,
                seqidx=index,
            )

    def __setitem__(self, index, value):
        """Add a column to seqs_annot."""
        if isinstance(index, str):
            self.seqs_annot[index] = value
        else:
            raise ValueError(
                "SeqData only supports setting seq_annot columns with indexing."
            )

    def __repr__(self):
        """Representation of SeqData object."""
        descr = f"SeqData object with = {self._n_obs} seqs"
        for attr in [
            "seqs",
            "names",
            "rev_seqs",
            "ohe_seqs",
            "ohe_rev_seqs",
            "seqs_annot",
            "pos_annot",
            "seqsm",
            "uns",
        ]:
            if attr in ["seqs", "names", "rev_seqs", "ohe_seqs", "ohe_rev_seqs"]:
                if getattr(self, attr) is not None:
                    descr += f"\n{attr} = {getattr(self, attr).shape}"
                else:
                    descr += f"\n{attr} = None"
            elif attr in ["seqs_annot"]:
                keys = getattr(self, attr).keys()
                if len(keys) > 0:
                    descr += f"\n{attr}: {str(list(keys))[1:-1]}"
            elif attr in ["pos_annot"]:
                if getattr(self, attr) is not None:
                    descr += f"\n{attr}: PyRanges object with {len(getattr(self, attr))} features"
                else:
                    descr += f"\n{attr}: None"
            elif attr in ["seqsm", "uns"]:
                if len(getattr(self, attr)) > 0:
                    descr += f"\n{attr}: {str(list(getattr(self, attr).keys()))[1:-1]}"
                else:
                    descr += f"\n{attr}: None"
        return descr

    def copy(self):
        """Return a copy of the SeqData object."""
        return deepcopy(self)

    def write_h5sd(self, path: PathLike, mode: str = "w"):
        """Write SeqData object to h5sd file.

        Parameters
        ----------
        path: PathLike
            Path to h5sd file. If file exists, it will be overwritten.
        mode: str, optional
            Mode to open h5sd file. Default is "w".
        """
        from .._io import write_h5sd

        write_h5sd(self, path, mode)

    def to_dataset(
        self,
        target_keys: Union[str, List[str]] = None,
        seq_transforms: List[str] = None,
        transform_kwargs: dict = {},
    ) -> SeqDataset:
        """Convert SeqData object to SeqDataset PyTorch dataset.

        Parameters
        ----------
        target: str or list of str, optional
            Target to use for training. If None, -1 will be used for all sequences.
        seq_transforms: list of str, optional
            List of sequence transforms to apply.
        transform_kwargs: dict, optional
            Keyword arguments to pass to sequence transforms.

        Returns
        -------
        SeqDataset
            PyTorch dataset class for SeqData object.
        """
        from .._transforms import ReverseComplement, OneHotEncode, ToTensor
        from torchvision import transforms as torch_transforms

        transforms = []
        if target_keys is None:
            targs = None
        else:
            targs = self.seqs_annot[target_keys].values

        if seq_transforms is None:
            print("No transforms given, assuming just need to tensorize.")
            transforms = [ToTensor(**transform_kwargs)]
            return SeqDataset(
                self.ohe_seqs,
                names=self.names,
                targets=targs if target_keys is not None else None,
                rev_seqs=self.ohe_rev_seqs,
                transform=torch_transforms.Compose(transforms),
            )

        if "reverse_complement" in seq_transforms:
            if self.seqs is not None:
                transforms.append(ReverseComplement(**transform_kwargs))
            else:
                Raise(ValueError("Cannot reverse complement sequences if seqs is None"))

        ohe_flag = True
        if "one_hot_encode" in seq_transforms:
            if self.seqs is not None:
                transforms.append(OneHotEncode(**transform_kwargs))
            else:
                Raise(ValueError("Cannot one hot encode sequences if seqs is None"))
            ohe_flag = False

        transforms.append(ToTensor(**transform_kwargs))

        if ohe_flag:
            return SeqDataset(
                self.ohe_seqs,
                names=self.names,
                targets=targs if target_keys is not None else None,
                rev_seqs=self.ohe_rev_seqs,
                transform=torch_transforms.Compose(transforms),
            )
        else:
            return SeqDataset(
                self.seqs,
                names=self.names,
                targets=targs if target_keys is not None else None,
                rev_seqs=self.rev_seqs,
                transform=torch_transforms.Compose(transforms),
            )

    def make_names_unique(self):
        """Make sequence names unique by appending a number to the end of each name."""
        n_digits = len(str(self.n_obs))
        new_index = np.array(
            [
                "seq{num:0{width}}".format(num=i, width=n_digits)
                for i in range(self.n_obs)
            ]
        )
        self.names = new_index
        self.seqs_annot.index = new_index


# TODO: move this to utils or something
@singledispatch
def _gen_dataframe(anno, length, index_names):
    if anno is None or len(anno) == 0:
        return pd.DataFrame(index=pd.RangeIndex(0, length, name=None).astype(str))
    for index_name in index_names:
        if index_name in anno:
            return pd.DataFrame(
                anno,
                index=anno[index_name],
                columns=[k for k in anno.keys() if k != index_name],
            )
    return pd.DataFrame(anno, index=pd.RangeIndex(0, length, name=None).astype(str))


@_gen_dataframe.register(pd.DataFrame)
def _(anno, length, index_names):
    anno = anno.copy(deep=False)
    if not is_string_dtype(anno.index):
        # warnings.warn("Transforming to str index.", ImplicitModificationWarning)
        anno.index = anno.index.astype(str)
    return anno


@_gen_dataframe.register(pd.Series)
@_gen_dataframe.register(pd.Index)
def _(anno, length, index_names):
    raise ValueError(f"Cannot convert {type(anno)} to DataFrame")


@singledispatch
def convert_to_dict(obj) -> dict:
    return dict(obj)


@convert_to_dict.register(dict)
def convert_to_dict_dict(obj: dict):
    return obj


@convert_to_dict.register(np.ndarray)
def convert_to_dict_ndarray(obj: np.ndarray):
    if obj.dtype.fields is None:
        raise TypeError(
            "Can only convert np.ndarray with compound dtypes to dict, "
            f"passed array had “{obj.dtype}”."
        )
    return {k: obj[k] for k in obj.dtype.fields.keys()}


@convert_to_dict.register(type(None))
def convert_to_dict_nonetype(obj: None):
    return dict()
