import argparse
import pandas as pd
import os
import numpy as np
import scanpy as sc
import muon as mu
from matplotlib import use
use('Agg')
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use("Agg")
import seaborn as sns
import panpipes.funcs as pnp
from scipy.stats import rankdata

import sys
import logging
L = logging.getLogger()
L.setLevel(logging.INFO)
log_handler = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter('%(asctime)s: %(levelname)s - %(message)s')
log_handler.setFormatter(formatter)
L.addHandler(log_handler)
# parse arguments
parser = argparse.ArgumentParser()


parser.add_argument('--channel_col',
                    default="sample_id",
                    help='')
parser.add_argument('--filtered_mudata',
                    default=None,
                    help='')  
parser.add_argument('--bg_mudata',
                    default=None,
                    help='')    
parser.add_argument('--figpath',
                    default="./figures/background",
                    help='')  


args, opt = parser.parse_known_args()
L.info("Running with params: %s", args)

sc.settings.figdir = args.figpath
# load filtered data - if use_umon is True this will return a mudata object, else returns an mudata object

try:
    L.info("Reading in MuData from '%s'" % args.filtered_mudata)
    mdata = mu.read(args.filtered_mudata)
    # need to delete rep else we lose too many cells in intersect_obs later
    if "rep" in mdata.mod.keys():
        L.info("Deleting modality 'rep' from MuData")
        del mdata.mod["rep"]
except FileNotFoundError:
    L.error("MuData object with path '%s' could not be found" % args.filtered_mudata)
    sys.exit("MuData object with path '%s' could not be found" % args.filtered_mudata)


# mu.pp.intersect_obs(mdata)
# load bg data


try:
    L.info("Reading in bg MuData from '%s'" % args.bg_mudata)
    mdata_bg = mu.read(args.bg_mudata)
except FileNotFoundError:
    L.error("bg_mudata object with path '%s' could not be found" % args.bg_mudata)
    sys.exit("bg_mudata object with path '%s' could not be found" % args.bg_mudata)

# barcode rank plots
L.info("Plotting barcode ranks")
df = mdata_bg['rna'].obs.copy()
df['total_counts'] = np.ravel(mdata_bg['rna'].X.sum(axis=1))
df['log_total_counts'] = np.log(df['total_counts'])
df['ranks'] = df.groupby('sample_id', as_index=False)['total_counts'].transform(lambda x: rankdata(x*-1, method="ordinal"))
df['log_ranks'] = np.log(df['ranks'])
df['is_cell'] = df.index.isin(mdata['rna'].obs_names.tolist())
df['is_cell'] = df['is_cell'].astype('category')
df['is_cell'] = df['is_cell'].cat.set_categories([True, False])

fig, ax= plt.subplots(figsize=(6,6))
sns.scatterplot(data=df, x="log_ranks", y="log_total_counts", 
                hue=args.channel_col, size="is_cell",  sizes=(2, 20),
                ax=ax, linewidth=0)
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
plt.savefig(os.path.join(args.figpath,'rna_barcode_ranks.png'), bbox_inches="tight")
plt.clf()

n_samples = len(df[args.channel_col].unique())
if n_samples > 1:
    g = sns.FacetGrid(df, col=args.channel_col, 
                    height=2.5,  
                    col_wrap=np.min([n_samples,8]))
    g.map(sns.scatterplot, "log_ranks", "log_total_counts", linewidth=0, s=1.5)
    g.set_titles(col_template="{col_name}")
    plt.savefig(os.path.join(args.figpath,'rna_facet_barcode_ranks.png'))
    plt.clf()



if 'rna' in mdata.mod.keys():
    # remove cells that passed cellranger filtering from bg
    mdata_bg['rna'].obs['is_empty'] = ~mdata_bg['rna'].obs_names.isin(mdata['rna'].obs_names)
    mu.pp.filter_obs(mdata_bg['rna'], 'is_empty')

# mdata_bg.update()
# mu.pp.intersect_obs(mdata_bg)

# calculate qc metrics
if 'rna' in mdata.mod.keys():
    # run stadard qc metrics
    L.info("Computing standard QC metrics for RNA")
    sc.pp.calculate_qc_metrics(mdata['rna'], inplace=True, percent_top=None, log1p=True)
    sc.pp.calculate_qc_metrics(mdata_bg['rna'], inplace=True, percent_top=None, log1p=True)



if 'prot' in mdata.mod.keys():
    L.info("Computing standard QC metrics for prot")
    sc.pp.calculate_qc_metrics(mdata['prot'], inplace=True, percent_top=None, log1p=True)
    sc.pp.calculate_qc_metrics(mdata_bg['prot'], inplace=True, percent_top=None, log1p=True)


mdata.update()
mdata_bg.update()

if 'rna' in mdata.mod.keys():
    n_samples_rna=len(mdata['rna'].obs['sample_id'].unique())
if "prot" in mdata.mod.keys():
    n_samples_prot=len(mdata['prot'].obs['sample_id'].unique())

# quantifying the top background features
if os.path.exists(args.figpath) is False:
    os.mkdirs(args.figpath)

if 'rna' in mdata.mod.keys():
    L.info("Quantifying the top background features for RNA")
    sc.pl.highest_expr_genes(mdata_bg['rna'],n_top=30, save="_rna_background.png")
    top_genes = pnp.scmethods.get_top_expressed_features(mdata_bg['rna'], n_top=30, group_by=args.channel_col)
    bg_df = pnp.scmethods.get_mean_background_fraction(mdata_bg['rna'], top_background=top_genes, group_by=args.channel_col)
    if bg_df.shape[0] >1:
        fig, ax= plt.subplots(figsize=(12,8))
        sns.heatmap(bg_df, ax=ax)
        fig.tight_layout()
        plt.savefig(os.path.join(args.figpath,"heatmap_background_" + args.channel_col + "_rna_top_expressed.png"))
    else:
        fig, ax= plt.subplots(figsize=(12,8))
        plt_df = bg_df.stack().reset_index().rename(columns={'level_1':'feature', 0:'exprs'})
        sns.barplot(data=plt_df, x='feature', y='exprs', ax=ax)
        pnp.plotting.adjust_x_axis(ax)
        fig.tight_layout()
        plt.savefig(os.path.join(args.figpath,"barplot_background_" + args.channel_col + "_rna_top_expressed.png"))


# quantifying the top background features
## Repeat for protein (if it exists)
if 'prot' in mdata.mod.keys():
    # this time we'll just use all the prot vars.
    L.info("Quantifying the top background features for prot")
    top_genes = list(mdata_bg['prot'].var_names)
    bg_df = pnp.scmethods.get_mean_background_fraction(mdata_bg['prot'], top_background=top_genes, group_by=args.channel_col)
    if bg_df.shape[0] >1:
        if len(top_genes) > 50:
            # split the plot into two
            split_int = int(np.ceil(137/2))
            fig, ax = plt.subplots(nrows=2,ncols=1,figsize=(24,10), facecolor="white")
            sns.heatmap(bg_df.iloc[:,1:split_int], ax=ax[0])
            sns.heatmap(bg_df.iloc[:,split_int:len(top_genes)], ax=ax[1])
            fig.suptitle("mean exprs (raw counts) of PROT in empty drops")  
            fig.tight_layout()
        else:
            fig, ax = plt.subplots(figsize=(12,10), facecolor="white")
            sns.heatmap(bg_df, ax=ax)
            fig.tight_layout()  
        plt.savefig(os.path.join(args.figpath,"heatmap_background_" + args.channel_col + "_prot_top_expressed.png"))
    else:
        plt_df = bg_df.stack().reset_index().rename(columns={'level_1':'feature', 0:'exprs'})
        if len(top_genes) > 50:
            split_int = int(np.ceil(137/2))
            fig, ax = plt.subplots(nrows=2,ncols=1,figsize=(24,10), facecolor="white")
            sns.heatmap(plt_df.iloc[1:split_int,:], ax=ax[0])
            pnp.plotting.adjust_x_axis(ax[0])
            sns.heatmap(plt_df.iloc[split_int:len(top_genes),:], ax=ax[1])
            pnp.plotting.adjust_x_axis(ax[1])
            fig.suptitle("mean exprs (raw counts) of PROT in empty drops")  
            fig.tight_layout()
        else:
            fig, ax= plt.subplots(figsize=(12,8))
            sns.barplot(data=plt_df, x='feature', y='exprs', ax=ax)
            pnp.plotting.adjust_x_axis(ax)
            fig.tight_layout()
        plt.savefig(os.path.join(args.figpath,"barplot_background_" + args.channel_col + "_prot_top_expressed.png"))


## QC for rna and protein comparing foreground and background

if 'rna' in mdata.mod.keys():
    L.info("Comparing foreground and background for RNA")
    pnp.pl.scatter_fg_vs_bg(mdata, mdata_bg,x="rna:log1p_n_genes_by_counts", y="rna:log1p_total_counts", facet_row=args.channel_col)
    plt.subplots_adjust(right=0.7)
    plt.savefig(os.path.join(args.figpath,"scatter_bg_fg_rna_nGene_rna_nUMI.png"),transparent=False)
    if 'prot' in mdata.mod.keys():
        L.info("Comparing foreground and background for RNA and prot")
        if n_samples_rna != n_samples_prot:
            L.warning("n_samples are not equal in rna and prot,taking intersect for mdata and mdata_bg")
            mu.pp.intersect_obs(mdata_bg)
            mu.pp.intersect_obs(mdata)
            pnp.pl.scatter_fg_vs_bg(mdata, mdata_bg,x="prot:log1p_total_counts", y="rna:log1p_n_genes_by_counts", facet_row=args.channel_col)
            plt.subplots_adjust(right=0.7)
            plt.savefig(os.path.join(args.figpath, "scatter_bg_fg_prot_nUMI_rna_nGene.png"),transparent=False)
            pnp.pl.scatter_fg_vs_bg(mdata, mdata_bg,x="rna:log1p_total_counts", y="prot:log1p_total_counts", facet_row=args.channel_col)
            plt.subplots_adjust(right=0.7)
            plt.savefig(os.path.join(args.figpath,"scatter_bg_fg_rna_nUMI_prot_nUMI.png"),transparent=False)
        else:
            L.info("n_samples are equal in RNA and prot")
            pnp.pl.scatter_fg_vs_bg(mdata, mdata_bg,x="prot:log1p_total_counts", y="rna:log1p_n_genes_by_counts", facet_row=args.channel_col)
            plt.subplots_adjust(right=0.7)
            plt.savefig(os.path.join(args.figpath, "scatter_bg_fg_prot_nUMI_rna_nGene.png"),transparent=False)
            pnp.pl.scatter_fg_vs_bg(mdata, mdata_bg,x="rna:log1p_total_counts", y="prot:log1p_total_counts", facet_row=args.channel_col)
            plt.subplots_adjust(right=0.7)
            plt.savefig(os.path.join(args.figpath,"scatter_bg_fg_rna_nUMI_prot_nUMI.png"),transparent=False)



plt.clf()


L.info("Done")

