import os,sys,numpy,pandas,scipy.stats
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn
from optparse import OptionParser
opts = OptionParser()
usage = "Clustering by accesson\nusage: %prog -p project02"
opts = OptionParser(usage=usage, version="%prog 1.0")
opts.add_option("-p", help="The project02 folder.")
options, arguments = opts.parse_args()
#
#
def calculate_fisher(markers, all_genes, accesson_dict):
    cluster_markers = {}
    for cluster in accesson_dict.keys():
        accfile = accesson_dict[cluster]
        accesson_df = pandas.read_csv(accfile, sep='\t', index_col=0)
        marker_genes = []
        for acc in accesson_df.index.values:
            genes = accesson_df.loc[acc, 'relevant_genes']
            if genes!='NONE':
                marker_genes.extend(genes.split(';'))
        cluster_markers[cluster] = list(set(marker_genes))
    clusters, ctypes = list(cluster_markers.keys()), list(markers.keys())
    fisher = numpy.zeros((len(clusters), len(ctypes)))
    for icluster,cluster in enumerate(clusters):
        in_cluster = cluster_markers[cluster]
        out_cluster = list(set(all_genes) - set(in_cluster))
        matrix = numpy.zeros((2,2), dtype=int)
        for ictype,ctype in enumerate(ctypes):
            in_ctype = markers[ctype]
            out_ctype = list(set(all_genes)-set(in_ctype))
            matrix[0,0] = len(list(set(in_cluster).intersection(set(in_ctype))))
            matrix[0,1] = len(list(set(in_cluster).intersection(set(out_ctype))))
            matrix[1,0] = len(list(set(out_cluster).intersection(set(in_ctype))))
            matrix[1,1] = len(list(set(out_cluster).intersection(set(out_ctype))))
            odd, pvalue = scipy.stats.fisher_exact(matrix)
            fisher[icluster, ictype] = abs(numpy.log10(pvalue))
    fisher_df = pandas.DataFrame(fisher, index=clusters, columns=ctypes)
    cluster_order = ['c_2', 'c_3', 'c_5', 'c_6', 'c_8', 'c_0', 'c_1', 'c_10', 'c_11', 'c_12']
    ctype_order = ['ExcL23', 'ExcL5_3', 'ExcL4', 'ExcL5_1', 'ExcL6', 'ExcL5_2',
               'Int_Vip', 'Int_Npy', 'Int_Cck', 'Int_Sst_2', 'Int_Sst_1', 'Int_Pv']
    fisher_df = fisher_df.loc[cluster_order, ctype_order]
    fisher_df.to_csv('fisher_test_between_ATAC_RNA.csv', sep='\t')
    return
#
#
#### These are marker genes generated by Hrvatin et al. from scRNA-seq (inDrop-seq) data
#
markers = {}
markers['ExcL23'] = ['Arpp21', 'Lamp5', 'Camk2n1', 'Calb1', 'Nrgn', 'Enc1', 'Bdnf', 'Nefm', 'Nefl', 'Atp2b4',
                  'Nptxr', 'Car10', 'B230216N24Rik', 'Nwd2', 'Rasgrf2', 'Camk2a', 'Pcsk2', 'Pdp1',
                  '2010300C02Rik', 'Ptk2b']
markers['ExcL4'] = ['Camk4', 'Nrn1', 'Rgs4', 'Arpp21', 'Stx1a', 'Nrsn1', 'Arc', 'Egr1', 'Tmem145', 'Rorb',
                 'Krt12', 'Car10', 'Coch', 'Nrgn', 'Brinp1', 'Kalrn', 'Nptxr', 'Gpr158', 'Camk2n1', 'R3hdm1']
markers['ExcL5_1'] = ['Cck', 'C1ql3', 'Lmo4', 'Bdnf', 'Col6a1', 'Nrn1', 'Enc1', 'Tesc', 'Hs3st2', 'Nrgn',
                  'Atp2b4', 'Slc17a7', 'Rasl10a', 'Sv2b', 'Bmp3', 'Slc8a2', 'Scg2', 'Lingo1', 'Nptxr', 'Nptx1']
markers['ExcL5_2'] = ['Fam19a1', 'Lamp5', 'Fezf2', 'Crym', 'Ina', 'C1ql3', 'Gap43', 'Slc6a7', 'Nos1ap',
                  'Bcl11b', 'Parm1', 'Ntng1', 'Gda', 'Etv1', 'Fam131a', 'Nrgn', 'Cds1', 'Miat', 'Vat1l', 'Rnf112']
markers['ExcL5_3'] = ['Deptor', 'Hs3st2', 'Gpr88', 'Pde1a', 'Pcp4', 'Dkk3', 'Nrn1', 'Stx1a', 'Nefm', 'Egr1',
                  'Pgm2l1', 'Tbr1', 'Slc6a7', 'Elavl4', 'Fam163b', 'Lmo4', 'Slc17a7', 'Sv2b', 'Tmsb10', 'Hrh3']
markers['ExcL6'] = ['Rprm', 'Pcp4', '3110035E14Rik', 'Nptx1', 'Col5a1', 'Trbc2', 'Crym', 'Baiap2', 'Garnl3',
                 'Nxph3', 'Foxp2', 'Tbr1', 'Homer1', 'Nos1ap', 'Ina', 'Nefm', 'Syt7', 'Arc', 'Adcy1', 'Ttc9b']
markers['Int_Cck'] = ['Cnr1', 'Htr3a', 'Cxcl14', 'Cadps2', 'Tac2', 'Rgs12', 'Necab1', 'Tnfaip8l3', 'Fxyd6',
                      'Kcnip1', 'Cck', 'Nrip3', 'Rab3c', 'Cplx2', 'Gad1', 'Shisa9', 'Slc32a1', 'Slc7a14',
                      'Dner', 'Nr2f2']
markers['Int_Pv'] = ['Pvalb', 'Syt2', 'Cplx1', 'Kcnc1', 'Kcnc2', 'Tac1', 'Kcnip2', 'Kcnab3', 'Gad1', 'Hcn1',
                     'Gabra1', 'Vamp1', 'Erbb4', 'Slc32a1', 'Btbd11', 'Ppargc1a', 'Cox6a2', 'Slc6a1', 'Fndc5',
                     'Slc24a2']
markers['Int_Vip'] = ['Vip', 'Penk', 'Crh', 'Calb2', 'Tac2', 'Synpr', 'Pthlh', 'Igf1', 'Cxcl14', 'Dlx1', 'Ap1s2',
                      'Zcchc12', 'Asic4', 'Prox1', 'Cbln2', 'Kit', 'Adra1b', 'Slc24a3', 'Gad1', 'Vgf']
markers['Int_Npy'] = ['Npy', 'Pnoc', 'Gad2', 'Kit', 'Gad1', 'Cplx3', 'Pde11a', 'Slc6a1', 'Lamp5', 'Nyap2',
                      'Dner', 'Rab3c', 'Hapln1', 'Fgf13', 'Resp18', 'Slc2a13', 'Caln1', 'Dlx6os1', 'Slc32a1', 'Id2']
markers['Int_Sst_1'] = ['Sst', 'Lypd6', 'Rab3b', 'Grin3a', 'Trhde', 'Rbp4', 'Satb1', 'Grm1', 'Crhbp', 'Calb1',
                        'Synpr', 'Cdh13', 'Zcchc12', 'Arx', 'Lypd6b', 'Mgat4c', 'Pdyn', 'Lhx6', 'Elfn1', 'Elavl2']
markers['Int_Sst_2'] = ['Pdyn', 'Sst', 'Crhbp', 'Cdh13', 'Calb1', 'Grin3a', 'Synpr', 'Rab3b', 'Vgf', 'Cbln4',
                        'Elavl2', 'Arx', 'Serpini1', 'Elfn1', 'Rpp25', 'Rgs17', 'Col19a1', 'Clmp', 'Calb2', 'Grm1']
#
#
#### Please run script_for_project02.py first !!!!
#
#
gene_score_file = options.p+'/matrix/genes_scored_by_TSS_peaks.csv'
#
if not os.path.exists(gene_score_file):
    print('Error !!!!')
    print('Please run script_for_project02.py first !!!!')
    sys.exit()
#
#
expr_matrix = pandas.read_csv(gene_score_file, sep=',', index_col=0)
all_genes = list(set(expr_matrix.columns.values))
for ctype in markers.keys():
    markers[ctype] = list(set(markers[ctype]).intersection(set(all_genes)))
#
#### These are differential accessons generated by script_for_project02.py
#
accesson_files = [options.p+'/result/differential_accesson_of_cluster_0_vs_1_10_11_12.csv',
                  options.p+'/result/differential_accesson_of_cluster_10_vs_0_1_11_12.csv',
                  options.p+'/result/differential_accesson_of_cluster_11_vs_0_1_10_12.csv',
                  options.p+'/result/differential_accesson_of_cluster_12_vs_0_1_10_11.csv',
                  options.p+'/result/differential_accesson_of_cluster_1_vs_0_10_11_12.csv',
                  options.p+'/result/differential_accesson_of_cluster_2_vs_3_5_6_8.csv',
                  options.p+'/result/differential_accesson_of_cluster_3_vs_2_5_6_8.csv',
                  options.p+'/result/differential_accesson_of_cluster_5_vs_2_3_6_8.csv',
                  options.p+'/result/differential_accesson_of_cluster_6_vs_2_3_5_8.csv',
                  options.p+'/result/differential_accesson_of_cluster_8_vs_2_3_5_6.csv']
accesson_dict = {'c_0': accesson_files[0], 'c_10': accesson_files[1], 'c_11': accesson_files[2],
                 'c_12': accesson_files[3], 'c_1': accesson_files[4], 'c_2': accesson_files[5],
                 'c_3': accesson_files[6], 'c_5': accesson_files[7], 'c_6': accesson_files[8],
                 'c_8': accesson_files[9]}
calculate_fisher(markers, all_genes, accesson_dict)
#
fisher_df = pandas.read_csv('fisher_test_between_ATAC_RNA.csv', sep='\t', index_col=0)
#
fisher1 = fisher_df.loc[['c_6', 'c_5', 'c_2', 'c_3', 'c_8'],
                        ['ExcL23', 'ExcL4', 'ExcL5_1', 'ExcL5_2', 'ExcL5_3', 'ExcL6']]
fisher11 = scipy.stats.zscore(fisher1,axis=1)
fisher11 = scipy.stats.zscore(fisher11,axis=0)
fisher1 = pandas.DataFrame(fisher11, index=fisher1.index, columns=fisher1.columns)
fig1 = plt.figure()
seaborn.heatmap(fisher1, cmap='Reds', vmin=0, vmax=2)
plt.savefig('clusters_vs_Excl_types.png', bbox_inches='tight')
plt.close()
#
fisher2 = fisher_df.loc[['c_0', 'c_10', 'c_1', 'c_12', 'c_11'],
                        ['Int_Vip', 'Int_Pv', 'Int_Npy', 'Int_Cck', 'Int_Sst_2', 'Int_Sst_1']]
fisher22 = scipy.stats.zscore(fisher2,axis=1)
fisher22 = scipy.stats.zscore(fisher22,axis=0)
fisher2 = pandas.DataFrame(fisher22, index=fisher2.index, columns=fisher2.columns)
fig2 = plt.figure()
seaborn.heatmap(fisher2, cmap='Blues', vmin=0, vmax=2)
plt.savefig('clusters_vs_Int_types.png', bbox_inches='tight')
plt.close()
#
print('output figures:')
print('clusters_vs_Excl_types.png')
print('clusters_vs_Int_types.png')
#
#
