import multiprocessing
import os
import shutil
import subprocess
from abc import ABC, abstractmethod
from typing import List, Union

import numpy as np
import scanpy as sc
from gtars.tokenizers import MetaTokenizer as GMetaTokenizer
from gtars.tokenizers import Region as GRegion
from gtars.tokenizers import TreeTokenizer as GTreeTokenizer
from huggingface_hub import hf_hub_download
from rich.progress import track

from geniml.io import Region, RegionSet
from geniml.tokenization.split_file import split_file

from .hard_tokenization_batch import main as hard_tokenization
from .utils import Timer, time_str


class Tokenizer(ABC):
    @abstractmethod
    def tokenize(self, *args, **kwargs):
        raise NotImplementedError


class Namespace:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)


class TreeTokenizer(Tokenizer):
    """
    A fast, in memory, tokenizer that uses `gtokenizers` - a rust based tokenizer.

    This should be used to tokenize bulk data, like BED files. If you need to tokenize single cell data, use
    the `AnnDataTokenizer` instead.
    """

    @classmethod
    def from_pretrained(cls, model_path: str, **kwargs):
        """
        Create a new tokenizer from a pretrained model's vocabulary.

        Usage:
        ```
        tokenizer = TreeTokenizer.from_pretrained("path/to/universe.bed")
        ```

        :param str model_path: The path to the pretrained model on huggingface.
        """
        universe_file_path = hf_hub_download(model_path, "universe.bed")
        return cls(universe_file_path, **kwargs)

    def __init__(self, universe: str) -> GTreeTokenizer:
        """
        Create a new tokenizer.

        This tokenizer only accepts a path to a BED file containing regions.

        :param str universe: The universe to use for tokenization.
        """
        self._tokenizer = GTreeTokenizer(universe)

    @property
    def universe(self):
        return self._tokenizer.universe

    def tokenize(self, query: Union[str, RegionSet]) -> List[List[Region]]:
        """
        Tokenize a Region or RegionSet into the universe

        :param Union[Region, RegionSet] query: The query to tokenize.
        """
        if isinstance(query, sc.AnnData) or isinstance(query, RegionSet):
            result = self._tokenizer(query)
            return result.to_regions()
        else:
            raise ValueError(
                f"Please pass a RegionSet object or a path to a BED file. You passed: {type(query)}"
            )

    def encode(self, query: sc.AnnData) -> List[List[int]]:
        """
        Tokenize an AnnData object to IDs.

        :param sc.AnnData query: The query to tokenize.
        """
        if isinstance(query, sc.AnnData) or isinstance(query, RegionSet):
            result = self._tokenizer(query)
            return result.to_ids()
        else:
            raise ValueError(
                f"Please pass a RegionSet object or a path to a BED file. You passed: {type(query)}"
            )

    def decode(self, query: List[List[int]]) -> List[List[Region]]:
        """
        Decode a list of IDs back to regions.

        :param List[List[int]] query: The query to decode.
        """
        return [self._tokenizer.decode(ids) for ids in query]

    def padding_token(self) -> GRegion:
        return self._tokenizer.padding_token

    def padding_token_id(self) -> int:
        return self._tokenizer.padding_token_id

    def unknown_token(self) -> GRegion:
        return self._tokenizer.unknown_token

    def unknown_token_id(self) -> int:
        return self._tokenizer.unknown_token_id

    def mask_token(self) -> GRegion:
        return self._tokenizer.mask_token

    def mask_token_id(self) -> int:
        return self._tokenizer.mask_token_id

    def cls_token(self) -> GRegion:
        return self._tokenizer.cls_token

    def cls_token_id(self) -> int:
        return self._tokenizer.cls_token_id

    def bos_token(self) -> GRegion:
        return self._tokenizer.bos_token

    def bos_token_id(self) -> int:
        return self._tokenizer.bos_token_id

    def eos_token(self) -> GRegion:
        return self._tokenizer.eos_token

    def eos_token_id(self) -> int:
        return self._tokenizer.eos_token_id

    def sep_token(self) -> GRegion:
        return self._tokenizer.sep_token

    def sep_token_id(self) -> int:
        return self._tokenizer.sep_token_id

    def export(self, path: str):
        return self._tokenizer.export(path)

    def __len__(self):
        return len(self._tokenizer)

    def __call__(self, query: Union[str, RegionSet]) -> List[List[Region]]:
        if isinstance(query, str) or isinstance(query, RegionSet) or isinstance(query, list):
            result = self._tokenizer(query)
            return result
        else:
            raise NotImplementedError(
                "Only RegionSet, region, or str is supported for this tokenizer"
            )

    def __repr__(self):
        return "TreeTokenizer()"

    def __str__(self):
        return "TreeTokenizer()"


class AnnDataTokenizer(Tokenizer):
    """
    A fast, in memory, tokenizer that uses `gtokenizers` - a rust based tokenizer.

    This is actually a wrapper around the core TreeTokenizer. Ideally, we'd tokenize the
    AnnData natively in Rust, but that's a bit more involved. So, we'll just use the
    Python wrapper for now.
    """

    @classmethod
    def from_pretrained(cls, model_path: str, **kwargs):
        """
        Create a new tokenizer from a pretrained model's vocabulary.

        Usage:
        ```
        tokenizer = AnnDataTokenizer.from_pretrained("path/to/universe.bed")
        ```

        :param str model_path: The path to the pretrained model on huggingface.
        """
        universe_file_path = hf_hub_download(model_path, "universe.bed")
        return cls(universe_file_path, **kwargs)

    def __init__(self, universe: str = None, verbose: bool = False, tokenizer_type: str = "tree"):
        """
        Create a new tokenizer.

        This tokenizer only accepts a path to a BED file containing regions.

        :param str universe: The universe to use for tokenization.
        """
        self.verbose = verbose
        self._tokenizer: Union[GTreeTokenizer, GMetaTokenizer]
        if universe is not None:
            if tokenizer_type == "meta":
                self._tokenizer = GMetaTokenizer(universe)
            else:
                self._tokenizer = GTreeTokenizer(universe)
        else:
            self._tokenizer = None

    def _tokenize_anndata(self, adata: sc.AnnData) -> List[List[Region]]:
        """
        Tokenize an AnnData object. This is more involved, so it gets its own function.

        :param sc.AnnData query: The query to tokenize.
        """
        # extract regions from AnnData
        # its weird because of how numpy handle Intervals, the parent class of Region,
        # see here:
        # https://stackoverflow.com/a/43722306/13175187
        adata_features = [
            Region(chr, int(start), int(end))
            for chr, start, end in track(
                zip(adata.var["chr"], adata.var["start"], adata.var["end"]),
                total=adata.var.shape[0],
                description="Extracting regions from AnnData",
                disable=not self.verbose,
            )
        ]
        features = np.ndarray(len(adata_features), dtype=object)
        for i, region in enumerate(adata_features):
            features[i] = region
        del adata_features

        # tokenize
        tokenized = []
        for row in track(
            range(adata.shape[0]),
            total=adata.shape[0],
            description="Tokenizing",
            disable=not self.verbose,
        ):
            _, non_zeros = adata.X[row].nonzero()
            regions = features[non_zeros]
            tokenized.append(self._tokenizer(regions))

        return tokenized

    @property
    def universe(self):
        return self._tokenizer.universe

    def tokenize(self, query: sc.AnnData) -> List[List[Region]]:
        """
        Tokenize a Region or RegionSet into the universe

        :param Union[Region, RegionSet, sc.AnnData] query: The query to tokenize.
        :param bool ids_only: Whether to return only the IDs or the full TokenizedRegionSet
        :param bool as_strings: Whether to return the IDs as strings or ints
        """
        if isinstance(query, sc.AnnData):
            adata = query
        elif isinstance(query, str):
            adata = sc.read_h5ad(query)
        else:
            raise NotImplementedError("Only AnnData is supported right now.")

        adata_features = [
            Region(chr, int(start), int(end))
            for chr, start, end in track(
                zip(adata.var["chr"], adata.var["start"], adata.var["end"]),
                total=adata.var.shape[0],
                description="Extracting regions from AnnData",
                disable=not self.verbose,
            )
        ]
        features = np.ndarray(len(adata_features), dtype=object)
        for i, region in enumerate(adata_features):
            features[i] = region
        del adata_features

        # tokenize
        tokenized = []
        for row in track(
            range(adata.shape[0]),
            total=adata.shape[0],
            description="Tokenizing",
            disable=not self.verbose,
        ):
            _, non_zeros = adata.X[row].nonzero()
            regions = features[non_zeros]
            tokenized.append(self._tokenizer.tokenize(regions))

        return tokenized

    def encode(self, query: Union[sc.AnnData, str]) -> List[List[int]]:
        """
        Tokenize an AnnData object to IDs.

        :param sc.AnnData query: The query to tokenize.
        """
        if isinstance(query, sc.AnnData):
            adata = query
        elif isinstance(query, str):
            adata = sc.read_h5ad(query)
        else:
            raise NotImplementedError("Only AnnData is supported right now.")

        adata_features = [
            Region(chr, int(start), int(end))
            for chr, start, end in track(
                zip(adata.var["chr"], adata.var["start"], adata.var["end"]),
                total=adata.var.shape[0],
                description="Extracting regions from AnnData",
                disable=not self.verbose,
            )
        ]
        features = np.ndarray(len(adata_features), dtype=object)
        for i, region in enumerate(adata_features):
            features[i] = region
        del adata_features

        # tokenize
        tokenized = []
        for row in track(
            range(adata.shape[0]),
            total=adata.shape[0],
            description="Tokenizing",
            disable=not self.verbose,
        ):
            _, non_zeros = adata.X[row].nonzero()
            regions = features[non_zeros]
            tokenized.append(self._tokenizer.encode(regions))

        return tokenized

    def decode(self, query: List[List[int]]) -> List[List[Region]]:
        """
        Decode a list of IDs back to regions.

        :param List[List[int]] query: The query to decode.
        """
        return [self._tokenizer.decode(ids) for ids in query]

    def padding_token(self) -> GRegion:
        return self._tokenizer.padding_token

    def padding_token_id(self) -> int:
        return self._tokenizer.padding_token_id

    def unknown_token(self) -> GRegion:
        return self._tokenizer.unknown_token

    def unknown_token_id(self) -> int:
        return self._tokenizer.unknown_token_id

    def mask_token(self) -> GRegion:
        return self._tokenizer.mask_token

    def mask_token_id(self) -> int:
        return self._tokenizer.mask_token_id

    def cls_token(self) -> GRegion:
        return self._tokenizer.cls_token

    def cls_token_id(self) -> int:
        return self._tokenizer.cls_token_id

    def bos_token(self) -> GRegion:
        return self._tokenizer.bos_token

    def bos_token_id(self) -> int:
        return self._tokenizer.bos_token_id

    def eos_token(self) -> GRegion:
        return self._tokenizer.eos_token

    def eos_token_id(self) -> int:
        return self._tokenizer.eos_token_id

    def sep_token(self) -> GRegion:
        return self._tokenizer.sep_token

    def sep_token_id(self) -> int:
        return self._tokenizer.sep_token_id

    def export(self, path: str):
        return self._tokenizer.export(path)

    def __len__(self):
        return len(self.universe.regions)

    def __call__(self, query: Union[sc.AnnData, str]) -> List[List[Region]]:
        if isinstance(query, sc.AnnData):
            result = self._tokenize_anndata(query)
            return result
        elif isinstance(query, str):
            query = sc.read_h5ad(query)
            result = self._tokenize_anndata(query)
            return result
        else:
            raise NotImplementedError("Only AnnData is supported for this tokenizer.")

    def __repr__(self):
        return "AnnDataTokenizer()"

    def __str__(self):
        return "AnnDataTokenizer()"


def hard_tokenization_main(
    src_folder: str,
    dst_folder: str,
    universe_file: str,
    fraction: float = 1e-9,
    file_list: List[str] = None,
    num_workers: int = 10,
    bedtools_path: str = "bedtools",
) -> int:
    """Tokenizes raw BED files in parallel.

    This is the main function for hard tokenization. It uses multiple processes to
    speed up the tokenization process.

    Args:
        src_folder (str): The folder where raw BED files reside.
        dst_folder (str): The folder to store tokenized BED files.
        universe_file (str): The path to a universe file.
        fraction (float, optional): A parameter for bedtools.intersect.
            Defaults to 1e-9.
        file_list (list[str], optional): A list of files (just names not full
            paths) that need to be tokenized. Defaults to None and uses all BED
            files in src_folder.
        num_workers (int, optional): Number of processes used. Defaults to 10.
        bedtools_path (str, optional): The path to a bedtools binary. Defaults
            to "bedtools".

    Raises:
        Exception: No bedtools executable found

    Returns:
        int: 0 when the dst_folder folder has incomplete list of tokenized BED
            files which should be removed first; 1 when the dst_folder folder
            has the complete tokenized BED files or the tokenization process
            succeeds.
    """
    timer = Timer()
    start_time = timer.t()

    file_list_path = os.path.join(dst_folder, "file_list.txt")
    files = os.listdir(src_folder)
    file_count = len(files)
    if file_count == 0:
        print(f"No files in {src_folder}")
        return 0

    os.makedirs(dst_folder, exist_ok=True)
    if file_list is None:  # use all bed files in data_folder
        # generate a file list
        file_list = files
        print(f"Use all ({file_count}) bed files in {src_folder}")
    else:
        file_number = len(file_list)
        print(f"{file_count} bed files in total, use {file_number} of them")

    # check whether all files in file_list are tokenized
    number = -1
    if os.path.exists(dst_folder):
        all_set = set([f.strip() for f in file_list])
        existing_set = set(os.listdir(dst_folder))
        not_covered = all_set - existing_set
        number = len(not_covered)
    if number == 0 and len(existing_set) == len(all_set):
        print("Skip tokenization. Using the existing tokenization files")
        return 1
    elif len(existing_set) > 0:
        print(
            f"Folder {dst_folder} exists with incomplete tokenized files. Please empty/delete the folder first"
        )
        return 0

    with open(file_list_path, "w") as f:
        for file in file_list:
            f.write(file)
            f.write("\n")

    if bedtools_path == "bedtools":
        try:
            rval = subprocess.call([bedtools_path, "--version"])
        except Exception:
            raise Exception("No bedtools executable found")
        if rval != 0:
            raise Exception("No bedtools executable found")

    print(f"Tokenizing {len(file_list)} bed files ...")

    file_count = len(file_list)
    # split the file_list into several subsets for each worker to process in parallel
    nworkers = min(int(np.ceil(file_count / 20)), num_workers)
    if nworkers <= 1:
        tokenization_args = Namespace(
            data_folder=src_folder,
            file_list=file_list_path,
            token_folder=dst_folder,
            universe=universe_file,
            bedtools_path=bedtools_path,
            fraction=fraction,
        )
        hard_tokenization(tokenization_args)

    else:  # multiprocessing
        dest_folder = os.path.join(dst_folder, "splits")
        split_file(file_list_path, dest_folder, nworkers)
        args_arr = []
        for n in range(nworkers):
            temp_token_folder = os.path.join(dst_folder, f"batch_{n}")
            tokenization_args = Namespace(
                data_folder=src_folder,
                file_list=os.path.join(dest_folder, f"split_{n}.txt"),
                token_folder=temp_token_folder,
                universe=universe_file,
                bedtools_path=bedtools_path,
                fraction=fraction,
            )
            args_arr.append(tokenization_args)
        with multiprocessing.Pool(nworkers) as pool:
            processes = [pool.apply_async(hard_tokenization, args=(param,)) for param in args_arr]
            _ = [r.get() for r in processes]
        # move tokenized files in different folders to expr_tokens
        shutil.rmtree(dest_folder)
        for param in args_arr:
            allfiles = os.listdir(param.token_folder)
            for f in allfiles:
                shutil.move(
                    os.path.join(param.token_folder, f),
                    os.path.join(dst_folder, f),
                )
            shutil.rmtree(param.token_folder)
    os.remove(file_list_path)
    print(f"Tokenization complete {len(os.listdir(dst_folder))}/{file_count} bed files")
    elapsed_time = timer.t() - start_time
    print(f"[Tokenization] {time_str(elapsed_time)}/{time_str(timer.t())}")
    return 1
