# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_types.ipynb.

# %% auto 0
__all__ = ['Genes', 'Cells', 'Gene', 'Cell', 'populate_genes', 'populate_cells', 'Lineage']

# %% ../nbs/03_types.ipynb 4
import io, string, pkg_resources, resource
from itertools import repeat
import numpy as np, pandas as pd

from typing import Union, Optional

from atyp import AnnData, Series, DataFrame, NPArray

from nchr import U1, DOT, NIL, SPACE
from nlit import X_, PHATE_HVG
from chck import notnone, isdf
from ensc import EnsemblClient, GeneInfo

from pstr import safeattr, keymatch, shiftchars, axeschars, splitjoin

# %% ../nbs/03_types.ipynb 5
from .read import read_csv, save_csv
from .meta import Collection, Registry

# %% ../nbs/03_types.ipynb 27
class Gene(metaclass=Registry):
    '''
    A class representing a gene with associated information and functionality.

    Attributes
    ----------
    name : str
        The common name of the gene.
    ensg : str
        The unique Ensembl gene identifier.
    desc : str
        A description of the gene.
    cells : set
        A set of `Cell` objects that mark the presence of the gene.

    Methods
    -------
    id()
        Property that returns the Ensembl gene identifier.
    marks(cell)
        Associate a cell type with the gene, adding the cell to the `cells` set.
    __repr__()
        Compute the "official" string representation of the gene.
    __hash__()
        Return the hash value of the gene based on its name.
    __eq__(other)
        Check for equality between two `Gene` objects or a `Gene` and a string.

    '''
    ensc = EnsemblClient()

    def __init__(self, name, ensg, desc):
        self.name = name
        self.ensg = ensg
        self.desc = desc
        self.cells = set()

    @property
    def id(self): return self.ensg

    def marks(self, cell):
        if isinstance(cell, str):
            cell = Cell.lookup(cell) or Cell(cell)
            
        if cell not in self.cells:
            self.cells.add(cell)
            cell.add(self)
    
    @classmethod
    def search(cls, gene: str, species: str = 'human'): return cls.ensc.lookup(gene, species)
    def fetch(self, species: str = 'human'):
        return self.search(self.name, species)

    def __repr__(self): return f'{self.name}'
    def __hash__(self): return hash(self.name)
    def __eq__(self, other):
        if isinstance(other, Gene):
            return self.name == other.name and self.ensg == other.ensg
        elif isinstance(other, str):
            return bool(same := keymatch(other, self.name, self.ensg))
        return NotImplemented
    
    def __getitem__(self, name):        
        for cell in self.cells:
            if cell == name:  return cell
        raise KeyError(f'No cell named {name} in {self.cells}')

# %% ../nbs/03_types.ipynb 29
class Cell(metaclass=Registry):
    '''
    A class representing a cell type with associated information and functionality.

    Attributes
    ----------
    name : str
        The common name of the cell type.
    abrv : str, optional
        An abbreviation for the cell type.
    desc : str, optional
        A description of the cell type.
    genes : set
        A set of `Gene` objects that are markers for the cell type.

    Methods
    -------
    add(gene)
        Add a gene as a marker for the cell type.
    __repr__()
        Compute the "official" string representation of the cell type.
    __hash__()
        Return the hash value of the cell type based on its name.
    __eq__(other)
        Check for equality between two `Cell` objects or a `Cell` and a string.

    '''
    def __init__(
        self, name, *genes, 
        type = None, subtype = None, abrv = None, species = None, desc = None, 
        **kwargs
    ):
        self.name = name
        self.type = desc
        self.subtype = desc
        self.abrv = abrv
        self.species = desc
        self.desc = desc
        self.genes = set(genes)

    def add(self, gene):
        '''Add a gene to the cell type.'''
        if isinstance(gene, str):
            gene = Gene.lookup(gene) or Gene(gene, None, None)
        if gene not in self.genes:
            self.genes.add(gene)
            gene.marks(self)

    def __repr__(self): return f'{self.name} ({len(self.genes)} marker genes)'
    def __hash__(self): return hash(self.name)
    def __eq__(self, other):
        if isinstance(other, Cell):
            return self.name == other.name
        elif isinstance(other, str):
            return bool(same := keymatch(other, self.name, self.abrv))
        return NotImplemented
    
    def __getitem__(self, name):
        for gene in self.genes:
            if gene == name: return gene
        raise KeyError(f'No gene named {name} in {self.name}')

# %% ../nbs/03_types.ipynb 31
# Utilize the registry directly
Genes = Gene._registry
Cells = Cell._registry

# %% ../nbs/03_types.ipynb 33
def populate_genes(df_gene):    
    for _, row in df_gene.iterrows():
        Gene(row['gene'], row['id'], row['desc'])    

def populate_cells(df_cell):
    for _, row in df_cell.iterrows():
        drow = dict(row)
        gene = drow.pop('gene')
        cell = Cell(**drow)
        cell.add(gene)

# %% ../nbs/03_types.ipynb 34
# Populate genes and cells
populate_genes(read_csv('genes'))
populate_cells(read_csv('cells'))

# %% ../nbs/03_types.ipynb 45
class Lineage:
    def __init__(
        self, 
        name: str, 
        zsc_suffix: str = '_zscore', 
        obs_suffix: str = '_cells', 
        **kwargs
    ):
        self.name = name
        self.zsc_suffix = zsc_suffix or ''
        self.obs_suffix = obs_suffix or ''

    def __str__(self): return self.name
    def __hash__(self): return hash(self.name)
    def __eq__(self, other):
        if isinstance(other, str): return self.name == other 
        return super().__eq__(other)

    @property
    def score(self) -> str: return f'{self.name}{self.zsc_suffix}'
    @property
    def cells(self) -> str: return f'{self.name}{self.obs_suffix}'
    @property
    def strs(self) -> tuple: return (self.name, self.score, self.cells)
    @classmethod
    def from_str(cls, name: str): return cls(name)
    def idx(self, adata: AnnData) -> Series:
        '''Where the lineage is present in the data.
        
        Examples
        --------
        Note that the following:
        >>> Lineage('esc').idx(adata)
        
        is equivalent to:
        >>> adata.obs['esc_cells'] == 'esc'
        '''
        return adata.obs[self.cells] == self.name

    def subset(self, adata: AnnData) -> AnnData:        
        bools = self.idx(adata)
        return adata[bools]

    def scores(self, adata: AnnData) -> Series:
        return adata.obs[self.score]
    
    @classmethod
    def extract(cls, a: str, *lins: 'Lineage') -> list:
        '''Extract the attribute `a` from each lineage in `lins`.'''
        return list(getattr(l, a) for l in lins)

    @classmethod
    def obscols(cls, adata: AnnData, *lins: 'Lineage') -> DataFrame:
        names, cells = cls.extract('name'), cls.extract('cells')
        return adata.obs[cells].isin(names)

    @classmethod
    def make_df(cls, adata: AnnData, *lins: 'Lineage') -> DataFrame:
        '''Make a dataframe of lineages from an AnnData object.'''
        df_lineages = cls.obscols(adata, *lins)
        boolean_idx = np.any(df_lineages.values, axis=1)
        return df_lineages[boolean_idx]
    
    @classmethod
    def plot_lineages(cls, adata: AnnData, *lins, **kwargs):
        try:
            import scanpy as sc   
            scores, cells = cls.extract('score'), cls.extract('cells')
            if len(lins) > 1: kwargs.setdefault('ncols', len(lins))
            sc.pl.embedding(adata, color=[*scores, *cells], ncols=len(lins), **kwargs)
        except ImportError: return

    def plot(self, adata: AnnData, **kwargs):
        type(self).plot_lineages(adata, self, **kwargs)
