# import brie
import numpy as np
import pandas as pd
from brie.utils import fetch_reads
import os
import pickle
from functools import reduce
import anndata as ad
import multiprocessing 
import pysam
# from sklearn.cluster import AgglomerativeClustering
# from scipy.optimize import linear_sum_assignment
import time
# from sklearn.cluster import KMeans
# from collections import defaultdict
from os import getpid
import random
import pickle
# import numpy.lib.recfunctions as rfn
# import numpy.ma as ma
# from itertools import compress
# import statistics
import editdistance
import warnings
# from collections import ChainMap
# from sklearn.preprocessing import MinMaxScaler
# from sklearn import linear_model
# from pathlib import Path
import sys
from sklearn.neighbors import NearestNeighbors



warnings.simplefilter(action='ignore', category=FutureWarning)
warnings.filterwarnings("ignore", category=Warning)







global CACHE_CHROM
global CACHE_SAMFILE
CACHE_CHROM = None
CACHE_SAMFILE = None

def check_pysam_chrom(samFile,chrom=None):
    """
    check samFile is a file name or a pysam object, and if chrom format.
    """
    global CACHE_CHROM
    global CACHE_SAMFILE

    if CACHE_CHROM is not None:
        if (samFile==CACHE_SAMFILE) and (chrom==CACHE_CHROM):
            return CACHE_SAMFILE,CACHE_CHROM

    
    if type(samFile)==str or type(samFile)==np.str_:
        ftype=samFile.split(".")[-1]
        if ftype != "bam" and ftype!="sam" and ftype!="cram":
            print("Error: file type need suffix of bam, sam or cram")
            sys.exit(1)
        if ftype == "cram":
            samFile=pysam.AlignmentFile(samFile,'rc')
        elif ftype=="bam":
            samFile=pysam.AlignmentFile(samFile,'rb')
        else:
            samFile=pysam.AlignmentFile(samFile,'r')

    if chrom is not None:
        if chrom not in samFile.references:
            if chrom.startswith("chr"):
                chrom=chrom.split('chr')[1]
            else:
                chrom="chr"+chrom
        if chrom not in samFile.references:
            print("Can't find reference %s in samFile"%chrom)
            return samFile,None

    CACHE_CHROM=chrom
    CACHE_SAMFILE=samFile
    return samFile,chrom



def get_fastq_file(fastqFilePath):
    fastqFile=pysam.FastaFile(fastqFilePath)
    return fastqFile




class get_old_TSS_count():
    def __init__(self,generefPath,filterTssPath,bamfilePath,fastqFilePath,outdir,cellBarcodePath,nproc,minCount,maxReadCount,clusterDistance):
        self.generefdf=pd.read_csv(generefPath,delimiter='\t')
        self.generefdf['len']=self.generefdf['End']-self.generefdf['Start']
        self.tssrefdf=pd.read_csv(filterTssPath,delimiter='\t')
        self.bamfilePath=bamfilePath
        self.count_out_dir=str(outdir)+'/count/'
        if not os.path.exists(self.count_out_dir):
            os.mkdir(self.count_out_dir)
        self.minCount=minCount
        self.cellBarcode=pd.read_csv(cellBarcodePath,delimiter='\t')['cell_id'].values
        self.nproc=nproc
        self.maxReadCount=maxReadCount
        self.clusterDistance=clusterDistance
        self.fastqFilePath=fastqFilePath


        

    def _getreads(self,bamfilePath,fastqFilePath,geneid):
        #fetch reads1 in gene 
        samFile, _chrom = check_pysam_chrom(bamfilePath,'chr'+str(self.generefdf.loc[geneid]['Chromosome']))
        reads = fetch_reads(samFile, _chrom,  self.generefdf.loc[geneid]['Start'] , self.generefdf.loc[geneid]['End'],  trimLen_max=100)
        reads1_umi = reads["reads1"]



        #select according to GX tag and CB (filter according to user owned cell)
        reads1_umi=[r for r in reads1_umi if r.get_tag('GX')==geneid]
        reads1_umi=[r for r in reads1_umi if r.get_tag('CB') in self.cellBarcode]

        # #filter strand invasion
        fastqFile=get_fastq_file(fastqFilePath)

        reads1_umi=[r for r in reads1_umi if editdistance.eval(fastqFile.fetch(start=r.reference_start-14, end=r.reference_start-1, region='chr'+str(self.generefdf.loc[geneid]['Chromosome'])),'TTTCTTATATGGG') >3 ]


        # #filter according to the cateria of SCAFE
        if self.generefdf.loc[geneid]['Strand']=='+':
            reads1_umi=[r for r in reads1_umi if r.is_reverse==False]
            reads1_umi=[r for r in reads1_umi if editdistance.eval(r.query_sequence[9:14],'ATGGG')<=4]
            reads1_umi=[r for r in reads1_umi if len(r.cigartuples)>=2]
            reads1_umi=[r for r in reads1_umi if (r.cigartuples[0][0]==4)&(r.cigartuples[0][1]>6)&(r.cigartuples[0][1]<20)&(r.cigartuples[1][0]==0)&(r.cigartuples[1][1]>5)]
        
        elif self.generefdf.loc[geneid]['Strand']=='-':
            reads1_umi=[r for r in reads1_umi if r.is_reverse==True]
            reads1_umi=[r for r in reads1_umi if editdistance.eval(r.query_sequence[-13:-8],'CCCAT')<=4]
            reads1_umi=[r for r in reads1_umi if len(r.cigartuples)>=2]
            reads1_umi=[r for r in reads1_umi if (r.cigartuples[0][0]==0)&(r.cigartuples[0][1]>5)&(r.cigartuples[1][0]==4)&(r.cigartuples[1][1]>6)&(r.cigartuples[1][1]<20)]


        #store start of reads and CB 
        reads_info=[]
        reads_info=[(r.positions[0],r.get_tag('CB'),r.cigarstring) for r in reads1_umi]
        #,r.query_sequence[:]
        

        return reads_info
    

        
    def _get_gene_reads(self):

        pool = multiprocessing.Pool(processes=self.nproc)

        self.generefdf.set_index('gene_id',inplace=True)
        bamfilePath=self.bamfilePath
        fastqFilePath=self.fastqFilePath

        readinfodict={}
        results=[]

        #get reads because pysam object cannot be used for multiprocessing so inputting bam file path 
        for i in self.generefdf.index:
            results.append(pool.apply_async(self._getreads,(bamfilePath,fastqFilePath,i)))
        pool.close()
        pool.join()


        results=[res.get() for res in results]

        for geneid,resls in zip(self.generefdf.index,results):
            readinfodict[geneid]=resls  


        #delete gene whose reads length is larger than maxReadCount
        for i in list(readinfodict.keys()):
            if len(readinfodict[i])>self.maxReadCount:
                readinfodict[i]=random.sample(readinfodict[i],self.maxReadCount)
            if len(readinfodict[i])<2:
                del readinfodict[i] 

        #store reads fetched
        outfilename=self.count_out_dir+'fetch_reads.pkl'
        with open(outfilename,'wb') as f:
            pickle.dump(readinfodict,f)

        return readinfodict



    def distribution(self,geneID,readinfo):
        singletranscriptdict={}
        try:
            generefdf=self.tssrefdf[self.tssrefdf['gene_id']==geneID]
            generefdf.reset_index(inplace=True,drop=True)
            trainX=generefdf['TSS'].values.reshape(-1,1)
            posiarray=np.array([t[0] for t in readinfo]).reshape(-1,1)
            CBarray=np.array([t[1] for t in readinfo]).reshape(-1,1)
            nbrs = NearestNeighbors(n_neighbors=1, algorithm='ball_tree').fit(trainX)
            distances, indices = nbrs.kneighbors(posiarray)
            select=indices[distances<=5]
            selectCB=CBarray[distances<=5]
            label=np.unique(select)
            for i in range(0,len(label)):
                keyid=generefdf.loc[label[i]]['gene_id']+'_'+generefdf.loc[label[i]]['transcript_id']+'@'+generefdf.loc[label[i]]['TSS'].astype('str')
                if len(select[select==label[i]])>=self.minCount:
                    singletranscriptdict[keyid]=(len(select[select==label[i]]),selectCB[select==label[i]])

            if len(singletranscriptdict)>=2:
                return singletranscriptdict
            else:
                return {}
            
        except ValueError:
            return {}
            



    def get_transcriptdict(self,):

        start_time=time.time()
        pool = multiprocessing.Pool(processes=self.nproc)
        readinfodict=self._get_gene_reads()
        transcriptdictls=[]

        for geneID in readinfodict.keys():
            transcriptdictls.append(pool.apply_async(self.distribution,(geneID,readinfodict[geneID])))
        pool.close()
        pool.join()

        transcriptdictls=[res.get() for res in transcriptdictls]

        #print(transcriptdictls)
        transcriptdict = {}
        for d in transcriptdictls:
            transcriptdict.update(d)


        

        print('do annotation',int(time.time()-start_time),'seconds.')


        #store reads fetched
        outfilename=self.count_out_dir+'transcript_store.pkl'
        with open(outfilename,'wb') as f:
            pickle.dump(transcriptdict,f)

        return transcriptdict



    def produce_sclevel(self):
        ctime=time.time()
        transcriptdict=self.get_transcriptdict()

        cellIDls=[]
        for transcriptID in transcriptdict.keys():
            cellID=np.unique(transcriptdict[transcriptID][1])
            cellIDls.append(list(cellID))
        cellIDset = set([item for sublist in cellIDls for item in sublist])
        finaldf=pd.DataFrame(index=cellIDset)


        for transcriptID in transcriptdict.keys():
            cellID,count=np.unique(transcriptdict[transcriptID][1],return_counts=True)
            transcriptdf=pd.DataFrame({'cell_id':cellID,transcriptID:count})
            transcriptdf.set_index('cell_id',inplace=True)
            finaldf[transcriptID]=finaldf.index.map(transcriptdf[transcriptID])
        finaldf.fillna(0,inplace=True)
        adata=ad.AnnData(finaldf) 

        print('we get adata')    
        
        vardf=pd.DataFrame(adata.var.copy())
        vardf.reset_index(inplace=True)
        vardf['gene_id']=vardf['index'].str.split('_',expand=True)[0]
        vardf=vardf.merge(self.generefdf,on='gene_id')
        vardf.set_index('index',drop=True,inplace=True)
        adata.var=vardf
        sc_output_h5ad=self.count_out_dir+'sc_TSS_count.h5ad'
        adata.write(sc_output_h5ad)
        print('produce h5ad Time elapsed',int(time.time()-ctime),'seconds.')


        return adata


