import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

def reconcile_Ma_Sestan_BICCN_schemes(dataset):
    root_path = "/path/to/all/files/"
    Ma_Sestan_name = f"{root_path}{dataset}_Azimuth_predictions_Ma_Sestan.csv"
    BICCN_name = f"{root_path}{dataset}_Azimuth_predictions_BICCN.csv"

    # read and rename annotations
    # 1) Read in the two csv files as dataframes.
    # 2) Change the column names to avoid confusion later.
    df_Ma_Sestan = pd.read_csv(Ma_Sestan_name, header=0, index_col=0)
    df_Ma_Sestan = df_Ma_Sestan.rename(mapper={'x':'Ma_Sestan'}, axis=1)
    df_BICCN = pd.read_csv(BICCN_name, header=0, index_col=0)
    df_BICCN = df_BICCN.rename(mapper={'x':'BICCN'}, axis=1)


    #3) Merge the two dataframes into a single study-specific dataframe.
    df_Ma_Sestan_BICCN_merge = pd.concat([df_Ma_Sestan,df_BICCN],axis=1)

    #4) Look at the cell type label for each annotation scheme.
    df_Ma_Sestan_BICCN_merge['Ma_Sestan_class'] = df_Ma_Sestan_BICCN_merge['Ma_Sestan'].apply(lambda x:ma_map[x])
    df_Ma_Sestan_BICCN_merge['BICCN_class'] = df_Ma_Sestan_BICCN_merge['BICCN'].apply(lambda x:BICCN_map[x])

    # 5) If the cell classes are same for the BICCN and the Ma_Sestan schemes, keep the cell barcode.
    #If the cell classes are different, remove the cell barcode. Please store all removed cell barcodes in a separate file for analysis.
    df_Ma_Sestan_BICCN_merge_kept = df_Ma_Sestan_BICCN_merge[df_Ma_Sestan_BICCN_merge['Ma_Sestan_class'] == df_Ma_Sestan_BICCN_merge['BICCN_class']]
    df_Ma_Sestan_BICCN_merge_removed = df_Ma_Sestan_BICCN_merge[df_Ma_Sestan_BICCN_merge['Ma_Sestan_class'] != df_Ma_Sestan_BICCN_merge['BICCN_class']]

    #However, I want to make sure that you keep the BICCN subclass names for the neurons and the Ma names for the Glial cells. Could you modify your code to do that? Just discard the other name in those cases.
    df_Ma_Sestan_BICCN_merge_kept_E_I = df_Ma_Sestan_BICCN_merge_kept[(df_Ma_Sestan_BICCN_merge_kept['BICCN_class'] == 'Exc') | (df_Ma_Sestan_BICCN_merge_kept['BICCN_class'] == 'Inh')]
    df_Ma_Sestan_BICCN_merge_kept_G = df_Ma_Sestan_BICCN_merge_kept[(df_Ma_Sestan_BICCN_merge_kept['Ma_Sestan_class'] == 'Glial')]

    df_Ma_Sestan_BICCN_merge_kept_E_I_x = df_Ma_Sestan_BICCN_merge_kept_E_I[['BICCN']].rename(mapper={'BICCN':'x'}, axis=1)
    df_Ma_Sestan_BICCN_merge_kept_G_x = df_Ma_Sestan_BICCN_merge_kept_G[['Ma_Sestan']].rename(mapper={'Ma_Sestan':'x'}, axis=1)

    #merge all
    df_Ma_Sestan_BICCN_merge_kept_x = pd.concat([df_Ma_Sestan_BICCN_merge_kept_E_I_x,df_Ma_Sestan_BICCN_merge_kept_G_x], axis=0)

    # output file
    df_Ma_Sestan_BICCN_merge_kept_x.to_csv(f"{root_path}{dataset}_Azimuth_predictions_Ma_Sestan_BICCN_reconcile.csv")
    df_Ma_Sestan_BICCN_merge_removed.to_csv(f"{root_path}{dataset}_Azimuth_predictions_Ma_Sestan_BICCN_inconsistent.csv")



if __name__=="__main__":
    BICCN_map = {'L6 IT':'Exc', 'L5/6 NP':'Exc', 'L2/3 IT':'Exc', 'L5 IT':'Exc', 'L6b':'Exc', 'L5 ET':'Exc','L4 IT':'Exc','L6 IT Car3':'Exc','L6 CT':'Exc','Chandelier':'Inh', 'Vip':'Inh','Lamp5':'Inh', 'Lamp5 Lhx6':'Inh', 'Sst Chodl':'Inh', 'Sst':'Inh', 'Pvalb':'Inh','Sncg':'Inh', 'Pax6':'Inh','Astro':'Glial', 'Oligo':'Glial',  'Micro/PVM':'Glial', 'Endo':'Glial', 'OPC':'Glial',  'VLMC':'Glial'}

    ma_map = {'L5-6 NP':'Exc', 'L6 IT-2':'Exc','L6B':'Exc','L3-5 IT-3':'Exc', 'L5 ET':'Exc', 'L3-5 IT-2':'Exc', 'L6 IT-1':'Exc','L3-5 IT-1':'Exc','L2-3 IT':'Exc','L6 CT':'Exc','ADARB2 KCNG1':'Inh','SST NPY':'Inh','VIP':'Inh', 'PVALB ChC':'Inh','SST':'Inh', 'LAMP5 RELN':'Inh', 'SST HGF':'Inh', 'PVALB':'Inh', 'LAMP5 LHX6':'Inh','PC':'Glial', 'RB':'Glial','SMC':'Glial','Micro':'Glial','Immune':'Glial','Astro':'Glial', 'Oligo':'Glial', 'Endo':'Glial', 'OPC':'Glial','VLMC':'Glial'}

    #List of all the dataset prefixes for which the reconciliation between the two schemes needs to be done
    dataset_list = []

    for dataset in dataset_list:
        reconcile_Ma_Sestan_BICCN_schemes(dataset)