"""Data input and output handling for ontologies
"""


import warnings
import numpy as np
from typing import Optional
from .accession_types import change_accessions
import gzip


class GeneSets:
    """Container for a Set of Pathways

    Attributes
    ==========
    genes: np.ndarray
        the sorted list of genes found in the pathways
    pathways: np.ndarray
        the sorted list of pathways found in the index
    n_genes: int
        the number of genes found
    n_pathways: int
        the number of pathways found
    bool_array: np.ndarray
        a (n_pathways, n_genes) bit array representing the bin-position of each gene
    pathway_sizes: np.ndarray
        a (n_pathways, ) array where each value represents the number of genes in that pathway index
    avg_p_size: float
        the mean number of genes across pathways
    membership: np.ndarray
        array representing the membership of each gene in an annotation
    n_bins: int
        number of bins to use when binning membership array
    bin_array: np.ndarray
        binned membership array

    Methods
    -------
    get_gene_subset:
        returns a subset of the `bool_array` for the provided gene list
    filter_pathways:
        filters pathways based on membership
    """
    def __init__(
            self,
            genes: Optional[np.ndarray] = None,
            pathways: Optional[np.ndarray] = None,
            ann_file: Optional[str] = None,
            n_bins: Optional[int] = 3,
            first_col_is_genes: Optional[bool] = False):
        """
        Parameters
        ==========
        genes: np.ndarray
            an array of gene names
        pathways: np.ndarray
            an array associated pathways
        """
        if ann_file:
            self._read_annotation_file(ann_file, first_col_is_genes)
        else:
            self._validate_inputs(genes, pathways)
            self._load_genes(genes)
            self._load_pathways(pathways)
            self._build_bool_array(genes, pathways)
        self._make_membership_profile()
        self.n_bins = n_bins

    def _read_annotation_file(self,
                              ann_file: str,
                              first_col_is_genes: Optional[bool] = False):
        """
         Reads annotation files which are in a descriptive format,
         i.e, "Gene1 Pathway1 Pathway2..." or "Pathway1 Gene1 Gene2..."

         Parameters
         ----------
         ann_file
             a file name of the annotation
         first_is_gene
             a parameter which specifies whether a gene is
         Returns
         -------
         pd.DataFrame
             a dataframe in a long format
         """
        is_gz = ann_file.endswith('.gz')

        row_names = []
        column_names = set()
        opener = gzip.open(ann_file, 'r') if is_gz else open(ann_file)
        with opener as f:
            for line in f:
                if is_gz:
                    line = line.decode('ASCII')
                els = line.rstrip().split('\t')
                row_names.append(els[0])
                els.pop(0)
                if 'http://' in els[0]:
                    els.pop(0)
                column_names |= set(els)
        column_names = list(column_names)

        positions = dict(zip(column_names, np.arange(len(column_names))))
        db_profiles = np.zeros((len(row_names), len(column_names)), dtype=int)

        opener = gzip.open(ann_file, 'r') if is_gz else open(ann_file)
        with opener as f:
            i = 0
            for line in f:
                if is_gz:
                    line = line.decode('ASCII')
                els = line.rstrip().split('\t')[1:]
                if 'http://' in els[0]:
                    els.pop(0)
                indices = [positions[el] for el in els]
                db_profiles[i, indices] = 1
                i += 1

        if first_col_is_genes:
            db_profiles = db_profiles.T
            db_genes = row_names
            db_names = column_names
        else:
            db_genes = column_names
            db_names = row_names
        self.pathways = np.array(db_names)
        self.genes = np.array(db_genes)
        self.bool_array = db_profiles
        self.n_genes = len(self.genes)
        self.n_pathways = len(self.pathways)

    def _validate_inputs(
            self,
            x: np.ndarray,
            y: np.ndarray):
        """validates inputs are as expected
        """
        if x.size == 0:
            raise ValueError("provided array must not be empty")
        if x.size != y.size:
            raise ValueError("genes and pathway arrays must be equal sized")
        if x.shape != y.shape:
            raise ValueError("genes and pathway arrays must be equally shaped")

    def _load_genes(
            self,
            genes: np.ndarray):
        """load genes and associated indices
        """
        self._gene_indices = {n: idx for idx, n in enumerate(np.sort(np.unique(genes)))}
        self.genes = np.array(list(self._gene_indices.keys()))
        self.n_genes = self.genes.size

    def _load_pathways(
            self,
            pathways: np.ndarray):
        """load pathways and associated indices
        """
        self._pathway_indices = {n: idx for idx, n in enumerate(np.sort(np.unique(pathways)))}
        self.pathways = np.array(list(self._pathway_indices.keys()))
        self.n_pathways = self.pathways.size

    def _calculate_pathway_sizes(self):
        """calculates and sets pathway sizes
        """
        self.pathway_sizes = self.bool_array.sum(axis=1)
        self.avg_p_size = self.pathway_sizes.mean()

    def _build_bool_array(
            self,
            genes: np.ndarray,
            pathways: np.ndarray):
        """create the bool array of genes/pathway interactions
        """
        self.bool_array = np.zeros((self.n_pathways, self.n_genes), dtype=int)
        for g, p in zip(genes, pathways):
            self.bool_array[self._pathway_indices[p]][self._gene_indices[g]] += 1

        self._calculate_pathway_sizes()

        # these may change with filtering so better to remove them
        del self._pathway_indices
        del self._gene_indices

    def _build_bin_split(
            self,
            membership: np.ndarray,
            n_bins: int) -> np.ndarray:
        """converts membership array to binned data using equivlanet split method
        """
        argidx = np.argsort(membership)
        max_size = membership.size
        bin_size = int(max_size / n_bins)

        bin_identities = np.zeros(max_size, dtype=int)
        self.bin_sizes = np.zeros(n_bins, dtype=int)
        self.bin_ranges = np.zeros(n_bins)
        self.bin_ranges[-1] = membership.max()

        for i in np.arange(0, n_bins):
            lower_bound = membership[argidx[bin_size * i]]

            if i < n_bins - 1:
                upper_bound = membership[argidx[bin_size * (i + 1)]]
                mask = (membership >= lower_bound) & (membership < upper_bound)

            # put remaining into last bin
            else:
                mask = (membership >= lower_bound)

            self.bin_sizes[i] = mask.sum()
            self.bin_ranges[i] = lower_bound
            bin_identities[mask] = i

        return bin_identities

    def _make_membership_profile(self) -> np.ndarray:
        """create a gene membership array"""
        self.membership = self.bool_array.sum(0)

    def get_gene_subset(
            self,
            gene_subset: np.ndarray) -> np.ndarray:
        """
        Index the bool-array for the required gene subset.
        Expects the subset to be sorted

        Parameters
        ==========
        gene_subset: np.ndarray
            a list of genes to subset the bin_array to

        Returns
        =======
        np.ndarray
            the bool_array subsetted to the indices of the `gene_subset`
        """
        gene_to_idx = {g: i for i, g in enumerate(self.genes)}
        missing = [g for g in gene_subset if g not in gene_to_idx]
        if missing:
            raise ValueError(f"Genes not found in gene sets: {missing[:5]}")
        idxs = [gene_to_idx[gene] for gene in gene_subset]
        self.sub_bool_array = self.bool_array[:, idxs]
        return self.sub_bool_array

    def get_membership_subset(
            self,
            gene_subset: np.ndarray) -> np.ndarray:
        """
        Index the bool-array for the required gene membership subset and bin it.
        Expects the subset to be sorted

        Parameters
        ==========
        gene_subset: np.ndarray
            a list of genes to subset the bin_array to

        Returns
        =======
        np.ndarray
            the bool_array subsetted to the indices of the `gene_subset`
        """
        gene_to_idx = {g: i for i, g in enumerate(self.genes)}
        missing = [g for g in gene_subset if g not in gene_to_idx]
        if missing:
            raise ValueError(f"Genes not found in gene sets: {missing[:5]}")
        idxs = [gene_to_idx[gene] for gene in gene_subset]
        sub_membership = self.membership[idxs]
        sub_membership_binned = self._build_bin_split(sub_membership, self.n_bins)
        return sub_membership_binned

    def filter_pathways(
            self,
            min_size: Optional[int] = None,
            max_size: Optional[int] = None):
        """
        Filters pathways to those within a specified range

        Parameters
        ----------
        min_size: Optional[int]
            the minimum size of the pathways, defaults to 0
        max_size: Optional[int]
            the maximum size of the pathways, default to current maximum
        """
        if not min_size and not max_size:
            warnings.warn("No minimum or maximum size provided. Doing nothing.")
            return
        if not min_size:
            min_size = 0
        if not max_size:
            max_size = np.max(self.pathway_sizes)
        if min_size < 0:
            raise ValueError("Provided minimum must be >= 0")
        if max_size <= min_size:
            raise ValueError(f"Provided maximum must be > min_size: {min_size}")

        # determine pathway-level mask
        p_mask = (self.pathway_sizes >= min_size) & (self.pathway_sizes <= max_size)

        # filter pathways
        self.bool_array = self.bool_array[p_mask]
        self.pathway_sizes = self.pathway_sizes[p_mask]
        self.pathways = self.pathways[p_mask]

        # determine gene level mask
        g_mask = self.bool_array.sum(axis=0) != 0

        # filter genes with no pathways
        self.bool_array = self.bool_array[:, g_mask]
        self.genes = self.genes[g_mask]

    @classmethod
    def from_gmt(cls, gmt_file, n_bins=3, min_size=None, max_size=None):
        """Load gene sets from a GMT file (e.g., MSigDB).

        Parameters
        ----------
        gmt_file : str
            Path to .gmt or .gmt.gz file.
        n_bins : int
            Number of bins for membership binning (default: 3).
        min_size : int, optional
            Minimum pathway size (filter after loading).
        max_size : int, optional
            Maximum pathway size (filter after loading).

        Returns
        -------
        GeneSets
        """
        is_gz = gmt_file.endswith('.gz')
        opener = gzip.open(gmt_file, 'rt') if is_gz else open(gmt_file)

        genes_list = []
        pathways_list = []
        descriptions = {}

        with opener as f:
            for line in f:
                parts = line.rstrip('\n\r').split('\t')
                if len(parts) < 3:
                    continue
                pathway_name = parts[0]
                description = parts[1]
                gene_names = [g for g in parts[2:] if g]

                if not gene_names:
                    continue

                descriptions[pathway_name] = description
                for gene in gene_names:
                    genes_list.append(gene)
                    pathways_list.append(pathway_name)

        obj = cls(
            genes=np.array(genes_list),
            pathways=np.array(pathways_list),
            n_bins=n_bins,
        )
        obj.descriptions = descriptions

        if min_size is not None or max_size is not None:
            kwargs = {}
            if min_size is not None:
                kwargs['min_size'] = min_size
            if max_size is not None:
                kwargs['max_size'] = max_size
            obj.filter_pathways(**kwargs)

        return obj

    def to_gmt(self, output_file, descriptions=None):
        """Export gene sets to GMT format.

        Parameters
        ----------
        output_file : str
            Path to output .gmt or .gmt.gz file.
        descriptions : dict, optional
            Mapping of pathway name to description string.
            Falls back to self.descriptions if available, otherwise 'na'.
        """
        if descriptions is None:
            descriptions = getattr(self, 'descriptions', {})

        is_gz = output_file.endswith('.gz')
        opener = gzip.open(output_file, 'wt') if is_gz else open(output_file, 'w')

        with opener as f:
            for i, pathway in enumerate(self.pathways):
                gene_mask = self.bool_array[i] > 0
                pathway_genes = self.genes[gene_mask]
                desc = descriptions.get(pathway, 'na')
                line = '\t'.join([pathway, desc] + list(pathway_genes))
                f.write(line + '\n')

    def map_genes(self, mapper, from_type='ensg', to_type='symbol'):
        """Convert gene IDs using a GeneMapper.

        Genes that can't be mapped are dropped. Updates self.genes,
        self.bool_array, and self.membership in place.

        Parameters
        ----------
        mapper : GeneMapper
            A GeneMapper instance with a cached mapping table.
        from_type : str
            Source ID type ('ensg', 'symbol', or 'entrez').
        to_type : str
            Target ID type ('ensg', 'symbol', or 'entrez').
        """
        converted, unmapped = mapper.convert(self.genes, from_type=from_type, to_type=to_type)

        # Build mask: keep genes that were successfully mapped
        keep_mask = np.array([c is not None for c in converted])

        if not np.any(keep_mask):
            raise ValueError("No genes could be mapped. Check from_type and to_type.")

        n_dropped = np.sum(~keep_mask)
        if n_dropped > 0:
            warnings.warn(f"Dropped {n_dropped} unmapped genes out of {len(self.genes)}.")

        self.genes = np.array([c for c, k in zip(converted, keep_mask) if k])
        self.bool_array = self.bool_array[:, keep_mask]
        self.n_genes = len(self.genes)
        self._make_membership_profile()

        # Remove pathways that now have zero genes
        pathway_sizes = self.bool_array.sum(axis=1)
        p_mask = pathway_sizes > 0
        if not np.all(p_mask):
            self.bool_array = self.bool_array[p_mask]
            self.pathways = self.pathways[p_mask]
            self.n_pathways = len(self.pathways)

    def convert_from_to(self,
                        input_format: str,
                        output_format: str,
                        species: Optional[str] = 'human'):
        """
        A function which changes accessions
        Parameters
        ----------
        input_format
            input accession type, takes 'enst', 'ensg', 'refseq', 'entrez', 'gs', 'ext'
        output_format
            output accession type, takes 'enst', 'ensg', 'refseq', 'entrez', 'gs', 'ext'
        species
            analyzed species, takes either 'human' or 'mouse'
        """
        if input_format in ['ensg', 'enst', 'refseq']:  # for ex., remove '.1' from 'ENSG00000128016.1'
            self.genes = np.array([gene.split('.')[0] for gene in self.genes])
        self.genes = change_accessions(self.genes,
                                       input_format,
                                       output_format,
                                       species)

    def __repr__(self) -> str:
        """
        """
        s = ""
        s += "Gene Ontology\n"
        s += f">> num_genes: {self.n_genes}\n"
        s += f">> num_pathways: {self.n_pathways}\n"
        return s


