#!/usr/bin/env python3

# -*- coding: utf-8 -*-

import scanpy as sc

import matplotlib.pyplot as plt

import os

import sys

import pysam

import collections

import argparse

import gzip

import pickle

import matplotlib

import numpy as np

import singlecellmultiomics

import singlecellmultiomics.molecule

import singlecellmultiomics.fragment

import singlecellmultiomics.features

import pysamiterators.iterators

import pysam

import pandas as pd

import scipy.sparse

import gzip

from singlecellmultiomics.molecule import MoleculeIterator

from singlecellmultiomics.alleleTools import alleleTools

import multiprocessing

from singlecellmultiomics.bamProcessing.bamFunctions import sort_and_index

matplotlib.use('Agg')

matplotlib.rcParams['figure.dpi'] = 160

def get_gene_id_to_gene_name_conversion_table(annotation_path_exons,

                                              featureTypes=['gene_name']):

    """Create a dictionary converting a gene id to other gene features,

        such as gene_name/gene_biotype etc.

    Arguments:

        annotation_path_exons(str) : path to GTF file (can be gzipped)

        featureTypes(list) : list of features to convert to, for example ['gene_name','gene_biotype']

    Returns:

        conversion_dict(dict) : { gene_id : 'firstFeature_secondFeature'}

        """

    conversion_table = {}

    with (gzip.open(annotation_path_exons, 'rt') if annotation_path_exons.endswith('.gz') else open(annotation_path_exons, 'r')) as t:

        for i, line in enumerate(t):

            parts = line.rstrip().split(None, 8)

            keyValues = {}

            for part in parts[-1].split(';'):

                kv = part.strip().split()

                if len(kv) == 2:

                    key = kv[0]

                    value = kv[1].replace('"', '')

                    keyValues[key] = value

            # determine the conversion name:

            if 'gene_id' in keyValues and any(

                    [feat in keyValues for feat in featureTypes]):

                conversion_table[keyValues['gene_id']] = '_'.join([

                    keyValues.get(feature, 'None')

                    for feature in featureTypes])

    return conversion_table

def count_transcripts(cargs):

    args, contig = cargs

    if args.alleles is not None:

        allele_resolver = alleleTools.AlleleResolver(

            args.alleles, lazyLoad=(not args.loadAllelesToMem))

    else:

        allele_resolver = None

    contig_mapping = None

    if args.contigmapping == 'danio':

        contig_mapping = {

            '1': 'CM002885.2',

            '2': 'CM002886.2',

            '3': 'CM002887.2',

            '4': 'CM002888.2',

            '5': 'CM002889.2',

            '6': 'CM002890.2',

            '7': 'CM002891.2',

            '8': 'CM002892.2',

            '9': 'CM002893.2',

            '10': 'CM002894.2',

            '11': 'CM002895.2',

            '12': 'CM002896.2',

            '13': 'CM002897.2',

            '14': 'CM002898.2',

            '15': 'CM002899.2',

            '16': 'CM002900.2',

            '17': 'CM002901.2',

            '18': 'CM002902.2',

            '19': 'CM002903.2',

            '20': 'CM002904.2',

            '21': 'CM002905.2',

            '22': 'CM002906.2',

            '23': 'CM002907.2',

            '24': 'CM002908.2',

            '25': 'CM002909.2',

        }

    # Load features

    contig_mapping = None

    #conversion_table = get_gene_id_to_gene_name_conversion_table(args.gtfexon)

    features = singlecellmultiomics.features.FeatureContainer()

    if contig_mapping is not None:

        features.remapKeys = contig_mapping

    features.loadGTF(

        args.gtfexon,

        select_feature_type=['exon'],

        identifierFields=(

            'exon_id',

            'transcript_id'),

        store_all=True,

        head=args.hf,

        contig=contig)

    features.loadGTF(

        args.gtfintron,

        select_feature_type=['intron'],

        identifierFields=['transcript_id'],

        store_all=True,

        head=args.hf,

        contig=contig)

    # What is used for assignment of molecules?

    if args.method == 'nla':

        molecule_class = singlecellmultiomics.molecule.AnnotatedNLAIIIMolecule

        fragment_class = singlecellmultiomics.fragment.NlaIIIFragment

        pooling_method = 1  # all data from the same cell can be dealt with separately

        stranded = None  # data is not stranded

    elif args.method == 'vasa' or args.method == 'cs':

        molecule_class = singlecellmultiomics.molecule.VASA

        fragment_class = singlecellmultiomics.fragment.SingleEndTranscriptFragment

        pooling_method = 1

        stranded = 1  # data is stranded, mapping to other strand

    else:

        raise ValueError("Supply a valid method")

    # COUNT:

    exon_counts_per_cell = collections.defaultdict(

        collections.Counter)  # cell->gene->umiCount

    intron_counts_per_cell = collections.defaultdict(

        collections.Counter)  # cell->gene->umiCount

    junction_counts_per_cell = collections.defaultdict(

        collections.Counter)  # cell->gene->umiCount

    gene_counts_per_cell = collections.defaultdict(

        collections.Counter)  # cell->gene->umiCount

    gene_set = set()

    sample_set = set()

    annotated_molecules = 0

    read_molecules = 0

    if args.producebam:

        bam_path_produced = f'{args.o}/output_bam_{contig}.unsorted.bam'

        with pysam.AlignmentFile(args.alignmentfiles[0]) as alignments:

            output_bam = pysam.AlignmentFile(

                bam_path_produced, "wb", header=alignments.header)

    ref = None

    if args.ref is not None:

        ref = pysamiterators.iterators.CachedFasta(pysam.FastaFile(args.ref))

    for alignmentfile_path in args.alignmentfiles:

        i = 0

        with pysam.AlignmentFile(alignmentfile_path) as alignments:

            molecule_iterator = MoleculeIterator(

                alignments=alignments,

                check_eject_every=5000,

                molecule_class=molecule_class,

                molecule_class_args={

                    'features': features,

                    'stranded': stranded,

                    'min_max_mapping_quality': args.minmq,

                    'reference': ref,

                    'allele_resolver': allele_resolver

                },

                fragment_class=fragment_class,

                fragment_class_args={

                    'umi_hamming_distance': args.umi_hamming_distance,

                    'features':features

                    },

                perform_qflag=True,

                # when the reads have not been tagged yet, this flag is very

                # much required

                pooling_method=pooling_method,

                contig=contig

            )

            for i, molecule in enumerate(molecule_iterator):

                if not molecule.is_valid():

                    if args.producebam:

                        molecule.write_tags()

                        molecule.write_pysam(output_bam)

                    continue

                molecule.annotate(args.annotmethod)

                molecule.set_intron_exon_features()

                if args.producebam:

                    molecule.write_tags()

                    molecule.write_pysam(output_bam)

                allele = None

                if allele_resolver is not None:

                    allele = molecule.allele

                    if allele is None:

                        allele = 'noAllele'

                # Obtain total count introns/exons reduce it so the sum of the

                # count will be 1:

                # len(molecule.introns.union( molecule.exons).difference(molecule.junctions))+len(molecule.junctions)

                total_count_for_molecule = len(molecule.genes)

                if total_count_for_molecule == 0:

                    continue  # we didn't find  any gene counts

                # Distibute count over amount of gene hits:

                count_to_add = 1 / total_count_for_molecule

                for gene in molecule.genes:

                    if allele is not None:

                        gene = f'{allele}_{gene}'

                    gene_counts_per_cell[molecule.sample][gene] += count_to_add

                    gene_set.add(gene)

                    sample_set.add(molecule.get_sample())

                # Obtain introns/exons/splice junction information:

                for intron in molecule.introns:

                    gene = intron

                    if allele is not None:

                        gene = f'{allele}_{intron}'

                    intron_counts_per_cell[molecule.sample][gene] += count_to_add

                    gene_set.add(gene)

                for exon in molecule.exons:

                    gene = exon

                    if allele is not None:

                        gene = f'{allele}_{exon}'

                    exon_counts_per_cell[molecule.sample][gene] += count_to_add

                    gene_set.add(gene)

                for junction in molecule.junctions:

                    gene = junction

                    if allele is not None:

                        gene = f'{allele}_{junction}'

                    junction_counts_per_cell[molecule.sample][gene] += count_to_add

                    gene_set.add(gene)

                annotated_molecules += 1

                if args.head and (i + 1) > args.head:

                    print(

                        f"-head was supplied, {i} molecules discovered, stopping")

                    break

        read_molecules += i

    if args.producebam:

        output_bam.close()

        final_bam_path = bam_path_produced.replace('.unsorted', '')

        sort_and_index(bam_path_produced, final_bam_path, remove_unsorted=True)

    return (

        gene_set,

        sample_set,

        gene_counts_per_cell,

        junction_counts_per_cell,

        exon_counts_per_cell,

        intron_counts_per_cell,

        annotated_molecules,

        read_molecules,

        contig

    )

if __name__ == '__main__':

    argparser = argparse.ArgumentParser(

        formatter_class=argparse.ArgumentDefaultsHelpFormatter,

        description='Create count tables from BAM file.')

    argparser.add_argument(

        '-o',

        type=str,

        help="output data folder",

        default='./rna_counts/')

    argparser.add_argument('alignmentfiles', type=str, nargs='+')

    argparser.add_argument(

        '-gtfexon',

        type=str,

        required=True,

        help="exon GTF file containing the features to plot")

    argparser.add_argument(

        '-gtfintron',

        type=str,

        required=True,

        help="intron GTF file containing the features to plot")

    argparser.add_argument('-umi_hamming_distance', type=int, default=1)

    argparser.add_argument(

        '-contigmapping',

        type=str,

        help="Use this when the GTF chromosome names do not match the ones in you bam file")

    argparser.add_argument(

        '-minmq',

        type=int,

        help="Minimum molcule mapping quality",

        default=20)

    argparser.add_argument(

        '-method',

        type=str,

        help="Data type: vasa,nla,cs",

        required=True)

    argparser.add_argument(

        '-head',

        type=int,

        help="Process this amount of molecules and export tables, also set -hf to be really fast")

    argparser.add_argument(

        '-hf',

        type=int,

        help="headfeatures Process this amount features and then continue, for a quick test set this to 1000 or so.")

    argparser.add_argument('-ref', type=str, help="Reference file (FASTA)")

    argparser.add_argument('-alleles', type=str, help="Allele file (VCF)")

    argparser.add_argument(

        '--loadAllelesToMem',

        action='store_true',

        help='Load allele data completely into memory')

    argparser.add_argument(

        '--producebam',

        action='store_true',

        help='Produce bam file with counts tagged')

    argparser.add_argument(

        '--ignoreMT',

        action='store_true',

        help='Ignore mitochondria')

    argparser.add_argument(

        '-t',

        type=int,

        default=8,

        help="Amount of chromosomes processed in parallel")

    argparser.add_argument(

        '-annotmethod',

        type=int,

        default=1,

        help="Annotation resolving method. 0: molecule consensus aligned blocks. 1: per read per aligned base")

    #argparser.add_argument('-tagged_bam_out',  type=str, help="Output bam file" )

    args = argparser.parse_args()

    workers = multiprocessing.Pool(args.t)

    if not os.path.exists(args.o):

        os.makedirs(args.o)

    jobs = []

    contigs_todo = []

    with pysam.AlignmentFile(args.alignmentfiles[0]) as g:

        # sort by size and do big ones first.. this will be faster

        for _, chrom in sorted(

                list(zip(g.lengths, g.references)), reverse=True):

            if chrom.startswith('ERCC') or chrom.startswith('chrUn') or chrom.endswith(

                    '_random') or chrom.startswith('GL') or chrom.startswith('JH'):

                continue

            if chrom.startswith('KN') or chrom.startswith('KZ') or chrom.startswith(

                    'chrUn') or chrom.endswith('_random') or 'ERCC' in chrom:

                continue

            if args.ignoreMT and chrom in ('mt', 'çhrMT', 'MT'):

                print("Ignoring mitochondria")

                continue

            jobs.append((args, chrom))

            contigs_todo.append(chrom)

    gene_counts_per_cell = collections.defaultdict(

        collections.Counter)  # cell->gene->umiCount

    exon_counts_per_cell = collections.defaultdict(

        collections.Counter)  # cell->gene->umiCount

    intron_counts_per_cell = collections.defaultdict(

        collections.Counter)  # cell->gene->umiCount

    junction_counts_per_cell = collections.defaultdict(

        collections.Counter)  # cell->gene->umiCount

    gene_set = set()

    sample_set = set()

    read_molecules = 0

    annotated_molecules = 0

    for (

        result_gene_set,

        result_sample_set,

        result_gene_counts_per_cell,

        result_junction_counts_per_cell,

        result_exon_counts_per_cell,

        result_intron_counts_per_cell,

        result_annotated_molecules,

        result_read_molecules,

        result_contig

    ) in workers.imap_unordered(count_transcripts, jobs):

        # Update all:

        gene_set.update(result_gene_set)

        sample_set.update(result_sample_set)

        for cell, counts in result_gene_counts_per_cell.items():

            gene_counts_per_cell[cell].update(counts)

        for cell, counts in result_junction_counts_per_cell.items():

            junction_counts_per_cell[cell].update(counts)

        for cell, counts in result_exon_counts_per_cell.items():

            exon_counts_per_cell[cell].update(counts)

        for cell, counts in result_intron_counts_per_cell.items():

            intron_counts_per_cell[cell].update(counts)

        read_molecules += result_read_molecules

        annotated_molecules += result_annotated_molecules

        # Now we finished counting

        contigs_todo = [x for x in contigs_todo if x != result_contig]

        print(

            f'Finished {result_contig}, so far found {read_molecules} molecules, annotated {annotated_molecules}, {len(sample_set)} samples')

        print(f"Remaining contigs:{','.join(contigs_todo)}")

        print('writing current matrices')

        # freeze order of samples and genes:

        sample_order = sorted(list(sample_set))

        gene_order = sorted(list(gene_set))

        # Construct the sparse matrices:

        sparse_gene_matrix = scipy.sparse.dok_matrix(

            (len(sample_set), len(gene_set)), dtype=np.int64)

        # Construct the sparse matrices:

        sparse_intron_matrix = scipy.sparse.dok_matrix(

            (len(sample_set), len(gene_set)), dtype=np.int64)

        # sparse_intron_matrix.setdefault(0)

        sparse_exon_matrix = scipy.sparse.dok_matrix(

            (len(sample_set), len(gene_set)), dtype=np.int64)

        # sparse_exon_matrix.setdefault(0)

        sparse_junction_matrix = scipy.sparse.dok_matrix(

            (len(sample_set), len(gene_set)), dtype=np.int64)

        for sample_idx, sample in enumerate(sample_order):

            if sample in gene_counts_per_cell:

                for gene, counts in gene_counts_per_cell[sample].items():

                    gene_idx = gene_order.index(gene)

                    sparse_gene_matrix[sample_idx, gene_idx] = counts

            if sample in exon_counts_per_cell:

                for gene, counts in exon_counts_per_cell[sample].items():

                    gene_idx = gene_order.index(gene)

                    sparse_exon_matrix[sample_idx, gene_idx] = counts

            if sample in intron_counts_per_cell:

                for gene, counts in intron_counts_per_cell[sample].items():

                    gene_idx = gene_order.index(gene)

                    sparse_intron_matrix[sample_idx, gene_idx] = counts

            if sample in junction_counts_per_cell:

                for gene, counts in junction_counts_per_cell[sample].items():

                    gene_idx = gene_order.index(gene)

                    sparse_junction_matrix[sample_idx, gene_idx] = counts

        # Write matrices to disk

        sparse_gene_matrix = sparse_gene_matrix.tocsc()

        sparse_intron_matrix = sparse_intron_matrix.tocsc()

        sparse_exon_matrix = sparse_exon_matrix.tocsc()

        sparse_junction_matrix = sparse_junction_matrix.tocsc()

        scipy.sparse.save_npz(

            f'{args.o}/sparse_gene_matrix.npz',

            sparse_gene_matrix)

        scipy.sparse.save_npz(

            f'{args.o}/sparse_intron_matrix.npz',

            sparse_intron_matrix)

        scipy.sparse.save_npz(

            f'{args.o}/sparse_exon_matrix.npz',

            sparse_exon_matrix)

        scipy.sparse.save_npz(

            f'{args.o}/sparse_junction_matrix.npz',

            sparse_junction_matrix)

        try:

            # Write scanpy file vanilla

            adata = sc.AnnData(

                sparse_gene_matrix.todense()

            )

            adata.var_names = gene_order

            adata.obs_names = sample_order

            adata.write(f'{args.o}/scanpy_vanilla.h5ad')

            # Write scanpy file, with introns

            adata = sc.AnnData(

                sparse_gene_matrix,

                layers={

                    'spliced': sparse_junction_matrix,

                    'unspliced': sparse_intron_matrix,

                    'exon': sparse_exon_matrix

                }

            )

            adata.var_names = gene_order

            adata.obs_names = sample_order

            adata.write(f'{args.o}/scanpy_complete.h5ad')

        except Exception as e:

            print("Could not (yet?) write the scanpy files, error below")

            print(e)

    print("Writing final tables to dense csv files")

    pd.DataFrame(sparse_gene_matrix.todense(), columns=gene_order,

                 index=sample_order).to_csv(f'{args.o}/genes.csv.gz')

    pd.DataFrame(sparse_intron_matrix.todense(), columns=gene_order,

                 index=sample_order).to_csv(f'{args.o}/introns.csv.gz')

    pd.DataFrame(sparse_exon_matrix.todense(), columns=gene_order,

                 index=sample_order).to_csv(f'{args.o}/exons.csv.gz')

    pd.DataFrame(

        sparse_junction_matrix.todense(),

        columns=gene_order,

        index=sample_order).to_csv(f'{args.o}/junctions.csv.gz')

    # Write as plaintext:

    adata.to_df().to_csv(f'{args.o}/counts.csv')