<reponame>flo-compbio/monet
# Author: <NAME> <<EMAIL>>
# Copyright (c) 2021 <NAME>
#
# This file is part of Monet.

from typing import Tuple
import logging
import os
import gc

import plotly.graph_objs as go

from monet import ExpMatrix
from monet import preprocess as pp
from monet.util import annotations as ann
from monet.velocyto import VelocytoData

_LOGGER = logging.getLogger(__name__)


def preprocess_matrix(
        matrix: ExpMatrix, ensembl_annotation_file: str,
        output_file: str = None,
        species='human',
        gene_min_cells_expressed: int = 3,
        max_mito_frac=0.15, min_transcripts=2000,
        title=None) \
            -> Tuple[ExpMatrix, go.Figure]:
    """Preprocess raw scRNA-Seq expresison matrix."""

    # read Ensembl GTF file
    gene_table = ann.get_ensembl_genes(ensembl_annotation_file)
    _LOGGER.info('Number of Ensembl genes: %d', gene_table.shape[0])

    matrix, _, _, fig = pp.preprocess_data(
        matrix, gene_table, species=species,
        gene_min_cells_expressed=gene_min_cells_expressed,
        max_mito_frac=max_mito_frac,
        min_transcripts=min_transcripts,
        title=title)

    if output_file is not None:
        # write matrix to disk
        matrix.save(output_file)

    return matrix, fig


def preprocess_10x_v1(
        cellranger_file: str, ensembl_annotation_file: str,
        output_file: str = None,
        prefix='filtered_gene_bc_matrices/hg19',
        *args, **kwargs) -> Tuple[ExpMatrix, go.Figure]:
    """Preprocess data generated by 10x Genomics' CellRanger v3 pipeline."""

    # read CellRanger file
    matrix = ExpMatrix.load_10x_v1(cellranger_file, prefix)

    return preprocess_matrix(matrix, ensembl_annotation_file, output_file,
                             *args, **kwargs)


def preprocess_10x_v3(
        cellranger_file: str, ensembl_annotation_file: str,
        output_file: str = None,
        prefix='filtered_feature_bc_matrix',
        *args, **kwargs) -> Tuple[ExpMatrix, go.Figure]:
    """Preprocess data generated by 10x Genomics' CellRanger v3 pipeline."""

    # read CellRanger file
    matrix = ExpMatrix.load_10x_v3(cellranger_file, prefix)

    return preprocess_matrix(matrix, ensembl_annotation_file, output_file,
                             *args, **kwargs)


def preprocess_10x_v3ts(
        cellranger_file: str, ensembl_annotation_file: str,
        expression_output_file: str = None, antibody_output_file: str = None,
        prefix='filtered_feature_bc_matrix',
        *args, **kwargs) -> Tuple[ExpMatrix, ExpMatrix, go.Figure]:
    """Preprocess data generated by 10x Genomics' CellRanger v3 pipeline."""

    # read CellRanger file
    matrix = ExpMatrix.load_10x_v3(cellranger_file, prefix)    
    print('Full matrix:', matrix.shape)

    sel = matrix.index.str.startswith('ENSG')
    antibody_matrix = matrix.loc[~sel].copy()
    matrix = matrix.loc[sel].copy()
    gc.collect()
    
    matrix, fig = preprocess_matrix(
        matrix, ensembl_annotation_file, expression_output_file,
        *args, **kwargs)

    antibody_matrix = antibody_matrix.loc[:, matrix.cells]

    if antibody_output_file is not None:
        # write antibody matrix to disk
        antibody_matrix.save(antibody_output_file)

    return matrix, antibody_matrix, fig


def preprocess_starsolo(
        expression_dir: str, ensembl_annotation_file: str,
        output_file: str = None,
        prefix='filtered_feature_bc_matrix',
        *args, **kwargs) -> Tuple[ExpMatrix, go.Figure]:
    """Preprocess data generated by STARsolo."""

    # read STARsolo output
    matrix = ExpMatrix.load_starsolo(expression_dir)

    return preprocess_matrix(matrix, ensembl_annotation_file, output_file,
                             *args, **kwargs)


def preprocess_starsolo_velocyto(
        starsolo_dir, genome_annotation_file, species='human',
        min_transcripts: int = 2000, max_mito_frac: float = 0.10,
        title=None, seed: int = 0) -> VelocytoData:
    """Preprocess Velcoyto data generated by STARsolo."""

    gene_table = ann.get_ensembl_genes(genome_annotation_file)
    print('Number of Ensembl genes:', gene_table.shape[0])    

    if title is None:
        title = ''

    gene_dir = os.path.join(starsolo_dir, 'Gene/raw')
    velo_dir = os.path.join(starsolo_dir, 'Velocyto/raw')

    # preprocess expression data
    matrix = ExpMatrix.load_starsolo(gene_dir)
    print(matrix.shape)

    matrix, _, _, fig = pp.preprocess_data(
        matrix, gene_table, species=species,
        max_mito_frac=max_mito_frac, min_transcripts=min_transcripts,
        title='%s (expression matrix)' % title)

    # preprocess spliced data
    spliced_matrix = ExpMatrix.load_starsolo(
        velo_dir, matrix_file_name='spliced.mtx')

    spliced_matrix, _, _, spliced_fig = pp.preprocess_data(
        spliced_matrix, gene_table, species=species,
        sel_cells=matrix.cells, sel_genes=matrix.genes,
        title='%s (spliced matrix)' % title)

    # preprocess unspliced data
    unspliced_matrix = ExpMatrix.load_starsolo(
        velo_dir, matrix_file_name='unspliced.mtx')

    unspliced_matrix, _, _, unspliced_fig = pp.preprocess_data(
        unspliced_matrix, gene_table, species=species,
        sel_cells=matrix.cells, sel_genes=matrix.genes,
        title='%s (unspliced matrix)' % title)

    data = VelocytoData(matrix, spliced_matrix, unspliced_matrix)

    return data, fig, spliced_fig, unspliced_fig
