#!/usr/bin/env python

import argparse

parser = argparse.ArgumentParser(description='Prepare cell2location reference signatures')

parser.add_argument("infile", type=str,default=None,help='input h5ad file with reference dataset')

parser.add_argument("output", type=str,default=None,help='folder to write output')

parser.add_argument("labels_key", type=str,default=None,help='column in adata.obs to be used as cell type label')

parser.add_argument("--batch_key", type=str,default=None,help='column in adata.obs to be used as bacth (single 10x reaction)')

parser.add_argument("--categorical_covariate_key",default=None, action='append',type=str,help='column in adata.obs to be used as categrical covariates - donor, 3/5, etc (no covariates by default). Multiple columns can be supplied by repetitive usage of this option.')

parser.add_argument("--continuous_covariate_key",default=None, action='append',type=str,help='column in adata.obs to be used as categrical covariates (no covariates by default). Multiple columns can be supplied by repetitive usage of this option.')

parser.add_argument("--gene_id", type=str,default=None,help='column in adata.var to be used as gene id')

parser.add_argument("--cell_count_cutoff", type=int,default=5,help='Gene filtering parameter: All genes detected in less than cell_count_cutoff cells will be excluded.')

parser.add_argument("--cell_percentage_cutoff2", type=float,default=0.03,help='Gene filtering parameter: All genes detected in at least this percentage of cells will be included.')

parser.add_argument("--nonz_mean_cutoff", type=float,default=1.12,help='Gene filtering parameter: genes detected in the number of cells between the above mentioned cutoffs are selected only when their average expression in non-zero cells is above this cutoff.')

parser.add_argument("--max_epochs", type=int,default=250,help='max_epochs for training')

parser.add_argument("--remove_genes_column", type=str,default=None,help='logical column in adata.var to be used to remove genes, for example mitochonrial. All genes with True in the column will be removed. None (defualt) mean to remove nothing.')

parser.add_argument("--seed", type=int,default=1,help='scvi seed value')

args = parser.parse_args()

import sys

import os

import scanpy as sc

from scipy.sparse import issparse

import anndata

import pandas as pd

import numpy as np

import matplotlib.pyplot as plt

import matplotlib as mpl

import cell2location

from cell2location.utils.filtering import filter_genes

from cell2location.models import RegressionModel

import torch

import scvi

from scvi import REGISTRY_KEYS

#######################

# create output folder

os.mkdir(args.output)

sys.stdout = open(args.output+"/c2l.ref.log", "w")

print(args)

print("cuda avaliable: "+str(torch.cuda.is_available()))

scvi.settings.seed = args.seed

# read data

ref = sc.read_h5ad(args.infile)

mtcnt = np.sum([gene.startswith('MT-') for gene in ref.var.index])

if mtcnt > 0:

    print('There are ' + str(mtcnt) + 'MT genes! Consider to remove them!')

if args.gene_id is not None:

  ref.var[args.gene_id] = ref.var[args.gene_id].astype('string')

  ref.var=ref.var.set_index(args.gene_id)

  print('Raw: cells = '+str(ref.shape[0])+"; genes = " + str(ref.shape[1]))

# filter genes

if args.remove_genes_column != None:

  print('Remove genes by "'+args.remove_genes_column+'". Following genes were removed:')

  print(ref.var[ref.var[args.remove_genes_column]])

  ref = ref[:,~ref.var[args.remove_genes_column]]

# filter genes 

selected = filter_genes(ref,

                        cell_count_cutoff=args.cell_count_cutoff,

                        cell_percentage_cutoff2=args.cell_percentage_cutoff2,

                        nonz_mean_cutoff=args.nonz_mean_cutoff)

plt.savefig(args.output+'/gene.filter.pdf') 

print('Before filtering: cells = '+str(ref.shape[0])+"; genes = " + str(ref.shape[1]))

ref = ref[:, selected].copy()

print('After filtering: cells = '+str(ref.shape[0])+"; genes = " + str(ref.shape[1]))

# remove slashes from celltype names

ref.obs[args.labels_key] = ref.obs[args.labels_key].astype(str).str.replace('/','_')

# train

cell2location.models.RegressionModel.setup_anndata(adata=ref,

                        # 10X reaction / sample / batch

                        batch_key=args.batch_key,

                        # cell type, covariate used for constructing signatures

                        labels_key=args.labels_key,

                        # multiplicative technical effects (platform, 3' vs 5', donor effect)

                        categorical_covariate_keys=args.categorical_covariate_key,

                        continuous_covariate_keys=args.continuous_covariate_key

                       )

mod = RegressionModel(ref)

mod.view_anndata_setup()

#mod.train(max_epochs=args.max_epochs,use_gpu=True,progress_bar_refresh_rate=0)

mod.train(max_epochs=args.max_epochs,progress_bar_refresh_rate=0)

# plot ELBO loss history during training, removing first 20 epochs from the plot

fig, ax = plt.subplots()

mod.plot_history(20)

plt.savefig(args.output+'/train.history.pdf')

ref = mod.export_posterior(

    ref, sample_kwargs={'num_samples': 1000, 'batch_size': 2500, 'use_gpu': True}

)

mod.save(args.output+"/rsignatures", overwrite=True)

# most likely I do not need this file

ref.write(args.output+"/rsignatures/sc.h5ad")

# save signatures

inf_aver = ref.varm['means_per_cluster_mu_fg'][[f'means_per_cluster_mu_fg_{i}' for i in ref.uns['mod']['factor_names']]].copy()

inf_aver.columns = ref.uns['mod']['factor_names']

inf_aver.to_csv(args.output+'/rsignatures/inf_aver.csv')

# function to plot QCs into file

def plot_QC1(m,plot,summary_name: str = "means",use_n_obs: int = 1000):

  if use_n_obs is not None:

    ind_x = np.random.choice(m.adata_manager.adata.n_obs, np.min((use_n_obs, m.adata.n_obs)), replace=False)

  else:

    ind_x = None

  m.expected_nb_param = m.module.model.compute_expected(

    m.samples[f"post_sample_{summary_name}"], m.adata_manager, ind_x=ind_x

  )

  x_data = m.adata_manager.get_from_registry(REGISTRY_KEYS.X_KEY)[ind_x, :]

  if issparse(x_data):

    x_data = np.asarray(x_data.toarray())

  mu = m.expected_nb_param["mu"]

  data_node = x_data

  plot.hist2d(np.log10(data_node.flatten()+1), np.log10(mu.flatten()+1), bins=50, norm=mpl.colors.LogNorm())

  plot.set_title("Reconstruction accuracy")

  plot.set(xlabel="Data, log10", ylabel="Posterior sample, values, log10")

def plot_QC2(m,plot,summary_name: str = "means",use_n_obs: int = 1000,scale_average_detection: bool = True):

  inf_aver = m.samples[f"post_sample_{summary_name}"]["per_cluster_mu_fg"].T

  if scale_average_detection and ("detection_y_c" in list(m.samples[f"post_sample_{summary_name}"].keys())):

    inf_aver = inf_aver * m.samples[f"post_sample_{summary_name}"]["detection_y_c"].mean()

  aver = m._compute_cluster_averages(key=REGISTRY_KEYS.LABELS_KEY)

  aver = aver[m.factor_names_]

  plot.hist2d(

    np.log10(aver.values.flatten() + 1),

    np.log10(inf_aver.flatten() + 1),

    bins=50,

    norm=mpl.colors.LogNorm(),)

  plot.set(xlabel="Mean expression for every gene in every cluster", ylabel="Estimated expression for every gene in every cluster")

# unfortunatelly it may not work specifically in case of underpopulated covariates/cell_types. It cannot be fixed on this level, so I'll use "try"

# see https://github.com/BayraktarLab/cell2location/issues/74

fig, (ax1,ax2) = plt.subplots(1,2)

try:

    plot_QC1(mod,plot=ax1,use_n_obs=10000)

except Exception as e:

    print(e)

try:

    plot_QC2(mod,plot=ax2)

except Exception as e:

    print(e)

plt.tight_layout()

plt.savefig(args.output+'/train.QC.pdf')

cell2location.utils.list_imported_modules()

sys.stdout.close()