import anndata
from Bio.Seq import Seq
from collections import Counter, defaultdict
import csv
import itertools
from more_itertools import flatten, pairwise
import numpy as np
import os
import pandas as pd
import re
import scanpy as sc
import scipy.sparse as sp_sparse
from scipy.stats import chi2_contingency, spearmanr
from shutil import copyfile
from sklearn.cluster import KMeans, SpectralClustering
from sklearn.decomposition import PCA, NMF
from sklearn.feature_selection import RFECV, SelectFromModel
from sklearn.linear_model import LogisticRegression, LogisticRegressionCV
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
from sklearn.model_selection import train_test_split, cross_val_score, StratifiedKFold, StratifiedGroupKFold
from statsmodels.stats.multitest import multipletests
import umap.umap_ as umap
UMAP = umap.UMAP
from textwrap import wrap
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import ticker
import seaborn as sns

from scquint.differential_splicing import run_differential_splicing
#from scquint.dimensionality_reduction import run_pca, run_vae
from scquint.utils import (filter_min_cells_per_cluster,
                            filter_min_cells_per_feature, filter_singletons,
                            group_normalize, relabel,
                            run_differential_expression, recluster)


matplotlib.use('pdf')
sns.set(style="white")

configfile: 'config.yaml'



genome_fasta_path = config["genome_fasta_path"]
full_gtf_path = config["gtf_path"]
chrom_sizes_path = config["chrom_sizes_path"]
encode_blacklist_path = config["encode_blacklist_path"]
sjdb_path = config["sjdb_path"]  # maybe should be created in this workflow
groupings = ["nontransitive", "transitive", "gene"]

sample_info = pd.read_csv("obs.txt.gz", "\t", index_col=0)
sample_ids = sample_info.index.values
print("len(sample_ids): ", len(sample_ids))
sample_info.loc[:, "cell_type"] = sample_info.cell_ontology_class.str.replace(" ", "_").str.replace("-", "_")
sample_info["plate_id"] = sample_info.index.str.split("_").str[1]

functional_class_info = pd.read_csv("input/cell_ontology_class_functional_annotation.073020.tsv", sep="\t", index_col=[0, 1]).filter(items=["cell category"])
functional_class_info.rename(columns={"cell category": "functional_class"}, inplace=True)
functional_class_info.functional_class.replace({
    "immune;stem cell/progenitor": "immune",
    "parenchymal;epithelial": "epithelial",
    "parenchymal": "other",
    "parenchymal;muscle cell": "other",
}, inplace=True)
sample_info = sample_info.merge(functional_class_info, how="left", left_on=["tissue", "cell_ontology_class"], right_index=True)
#sample_info.loc[sample_info.functional_class=="Other", "functional_class"] = sample_info.loc[sample_info.functional_class=="Other", "tissue"]
sample_info.loc[(sample_info.functional_class=="other") & (sample_info.tissue.isin(["Brain_Non-Myeloid", "Pancreas"])), "functional_class"] = sample_info.loc[(sample_info.functional_class=="other") & (sample_info.tissue.isin(["Brain_Non-Myeloid", "Pancreas"])), "tissue"]


sample_info["tct"] = sample_info.tissue.astype(str) + "." + sample_info.cell_type.astype(str)
tct_counts = sample_info.tct.value_counts()
tct_functional_class_map = sample_info.groupby("tct").functional_class.first()

functional_classes = ["endothelial", "epithelial"]
tcts_functional_class = defaultdict(list)

for tct in sample_info.tct.unique():
    for functional_class in functional_classes:
        if tct_counts.loc[tct] >= 100 and functional_class in tct_functional_class_map.loc[tct]:
            tcts_functional_class[functional_class].append(tct)
#print(tcts_functional_class)


sample_ids_subset = sorted(sample_info[sample_info.tissue=="Mammary_Gland"].index.values)
#print("len(sample_ids_subset): ", len(sample_ids_subset))

sample_ids_leafcutter = sample_ids_subset

# zcat GSE109774_list_of_SRR_accessions_and_raw_filenames.txt.gz | cut -f 2,3 | cut -d "-" -f 1-2 | sed 's/-/_/g' > srr_cell_pairs.txt
sample_srr = pd.read_csv("srr_cell_pairs.txt", '\t', header=None, names=['srr', 'cell'])
sample_srr = sample_srr[sample_srr.cell.isin(sample_ids)]

genes_bam_merge = {
    "Foxp1": {
        "region": ["chr6", 98888459, 99713014],
        "cell_types": ["late_pro_B_cell", "precursor_B_cell", "immature_B_cell", "naive_B_cell"],
    },
    "Smarca4": {
        "region": ["chr9", 21615842, 21633577],
        "cell_types": ["late_pro_B_cell", "precursor_B_cell", "immature_B_cell", "naive_B_cell"],
    },
    "Echdc2": {
        "region": ["chr4", 108164887, 108179634],
        "cell_types": ["hepatocyte_cluster_0", "hepatocyte_cluster_1"],
    },
    "Khk": {
        "region": ["chr5", 30924665, 30928620],
        "cell_types": ["kidney_collecting_duct_epithelial_cell", "bladder_urothelial_cell", "epithelial_cell_of_large_intestine"],
    },
    "Itpr1": {
        "region": ["chr6", 108482356, 108509805],
        "cell_types": ["bladder_urothelial_cell", "secretory_cell"],
    },
}


quantifications_mg_individual = ["gene-expression", "kallisto", "bins-nmf", "introns-shared-acceptor"]
#quantifications_mg_individual_expanded = ["gene-expression", "kallisto", "bins", "exons", "introns-gene", "introns-transitive", "SE", "introns-shared-acceptor"]
#quantifications_mg_individual_expanded = ["gene-expression", "kallisto", "bins", "exons", "introns-gene", "leafcutter", "introns-transitive"]
quantifications_mg_individual_expanded = ["leafcutter", "introns-transitive", "introns-shared-acceptor"]
#quantifications_mg_individual_expanded = ["gene-expression", "exons", "introns-gene", "introns-transitive", "SE", "introns-shared-acceptor"]
#quantifications_mg_individual_expanded = ["gene-expression", "kallisto", "bins-nmf"]

motif1 = "AAGCAGTGGTATCAACGCAGAGT"
motif2 = "ACTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT"
motif3 = "AAGCAGTGGTATCAACGCAGAGTACGGG"
motif1_rc = str(Seq(motif1).reverse_complement())
motif2_rc = str(Seq(motif2).reverse_complement())
motif3_rc = str(Seq(motif3).reverse_complement())

sample_info["cell_type"] = sample_info.cell_ontology_class.str.replace(" ", "_").str.replace("-", "_")
sample_info["Cell type"] = sample_info.cell_ontology_class


tissues = sample_info.tissue.unique()
cell_types = sample_info.cell_type.unique()
tissue_cell_type_pairs = []
tissue_cell_type_pairs_including_singletons = []
for tissue in tissues:
    tissue_cell_type_counts = sample_info[sample_info.tissue==tissue].cell_type.value_counts()
    tissue_cell_types = [ct for ct in sample_info[sample_info.tissue==tissue].cell_type.unique()
                         if tissue_cell_type_counts[ct] >= 100]
    tissue_cell_type_pairs_including_singletons += [[tissue, cell_type]
                                                    for cell_type in tissue_cell_types]
    if len(tissue_cell_types) < 2:
        continue
    tissue_cell_type_pairs += [[tissue, cell_type]
                               for cell_type in tissue_cell_types]
#print(len(tissue_cell_type_pairs), len(tissue_cell_type_pairs_including_singletons))

tissue_cell_type_pairs = sorted(tissue_cell_type_pairs)


flatten = lambda l: [item for sublist in l for item in sublist]



labels = ["Cell type", "mouse.id", "sex", "plate_id", "tissue"]


introns_to_plot = [
    "chr6:99260420-99266347",  # Foxp1_11076
    "chr9:21626013-21632805",  # Smarca4_28720
]

introns_to_plot_epithelial = [
    "chr6:108493894-108499108",  # Itpr1_26257
    "chr5:30927157-30928443",  # Khk_24896
]

genes_to_plot = [
    "ENSMUSG00000030067",
]


cell_type_order = ["(1) pro-B", "(2) pre-B", "(3) immature B", "(4) naive B"]


rule all:
    input:
        #"output/plots/marrow_b_n_events.svg",
        #"output/plots/tissues_per_donor.pdf",
        #expand("output/plots/{intron}.pdf", intron=introns_to_plot),
        expand("output/plots/epithelial/{intron}.pdf", intron=introns_to_plot_epithelial),
        #"output/dimensionality_reduction/marrow/classification_score.svg",
        #expand('output/differential_splicing/functional_class/{functional_class}/plot_marker_intron_coordinates.tsv', functional_class=functional_classes),
        #expand('output/differential_splicing/functional_class/{functional_class}/marker_introns.svg', functional_class=functional_classes),
        #expand('output/differential_splicing/functional_class/endothelial/{tct}/splicing.clusters.tsv', tct=tcts_functional_class["endothelial"]),
        #expand('output/differential_splicing/functional_class/epithelial/{tct}/splicing.clusters.tsv', tct=tcts_functional_class["epithelial"]),
        #expand("output/comparison/functional_class/{functional_class}/tissue_cell_type.svg", functional_class=["epithelial", "endothelial"]),
        #expand('output/differential_splicing/tissue_cell_type/{tissue}/marker_introns.svg', tissue=[t for t,ct in tissue_cell_type_pairs if t != "Kidney"]),
        #'output/differential_splicing/tissue_cell_type/Pancreas/marker_introns.svg',
        #'output/differential_splicing/tissue_cell_type/Marrow/marker_introns.svg',
        #'output/differential_splicing/tissue_cell_type/Tongue/marker_introns.svg',
        #'output/differential_splicing/tissue_cell_type/Brain_Non-Myeloid/marker_introns.svg',
        #'output/differential_splicing/tissue_cell_type/Heart/marker_introns.svg',
        #'output/differential_splicing/tissue_cell_type/Large_Intestine/marker_introns.svg',
        #'output/differential_splicing/tissue_cell_type/SCAT/marker_introns.svg',
        #"output/intron_coordinates.txt",
        #"output/intron_coordinates.tsv",
        #"output/dendrogram/exp.pdf",
        #"output/dendrogram/spl.pdf",
        #"output/cell_type_classification/Mammary_Gland/basal_cell/variables.tsv",
        #"output/dendrogram/gene-expression/pca_40_vs_introns-shared-acceptor/vae_hyperopt/entanglement.txt",
        #expand("output/bam_regions/Khk/{cell_type}/merged.bam", cell_type=genes_bam_merge["Khk"]["cell_types"]),
        #expand("output/bam_regions/Itpr1/{cell_type}/merged.bam", cell_type=genes_bam_merge["Itpr1"]["cell_types"]),
        #"output/differential_splicing/marrow_b/tf_Foxp1_trajectory.svg",
        #"output/differential_splicing/marrow_b/tf_Smarca4_trajectory.svg",
#        expand("output/bam_regions/Smarca4/{cell_type}/merged.bam", cell_type=genes_bam_merge["Smarca4"]["cell_types"]),
        #"output/differential_splicing/marrow_b/marker_introns_trajectory_spl.svg",
        #expand('output/differential_splicing/marrow_b/{cell_type}/marker_introns.tsv', cell_type=cell_type_order),
        #expand("output/bam_regions/Echdc2/{cell_type}/merged.bam", cell_type=genes_bam_merge["Echdc2"]["cell_types"]),
        #expand("output/comparison/tissue/Liver/{label}.svg", label=labels),
        #"output/adata_cellxgene_exp_spl.h5ad",
        #expand("output/comparison/pca_vae/{tissue}/cell_type.svg", tissue=["Diaphragm", "Mammary_Gland"]),
        #"output/dimensionality_reduction/all/introns-shared-acceptor/vae_hyperopt/latent.txt",
        #"output/adata_cellxgene_spl_pca.h5ad",
        #"output/comparison/Liver-hepatocyte/",
        #"output/comparison/well/Liver-hepatocyte/",
        #"output/comparison/Heart-monocyte/",
        #"output/comparison/Marrow-hematopoietic_stem_cell/",
        #"output/comparison/all",
        #"output/dimensionality_reduction/all/introns-shared-acceptor/vae_frequency-smoothed_False_100/latent.txt",
        #expand("output/dimensionality_reduction/tissue/{tissue}/introns-shared-acceptor/vae_frequency-smoothed_False/latent.txt", tissue=["Marrow", "Heart", "Brain_Myeloid"]),
        #expand("output/dimensionality_reduction/mg/{quantification}/vae_log_True/latent.txt", quantification=["leafcutter", "introns-transitive", "introns-shared-acceptor", "SE", "SE-shared-acceptor", "SE-shared-donor"]),
        #"output/quantification/introns/output/introns-shared-acceptor/adata.h5ad",
        #expand("output/mapping/filtered_bams/{sample_id}.bam.bai", sample_id=sample_ids),
        #"output/quantification/leafcutter/adata_annotated.h5ad",
        #expand("output/comparison/tissue/{tissue}/cell_ontology_class.svg", tissue=["Marrow", "Liver"]),
        #"output/dimensionality_reduction/marrow/classification_score.svg",
        #"output/quantification/SE/adata_annotated.h5ad",
        #"output/comparison/mg_3_38_F/classification_results_plots/",
        #"output/comparison/cortex_3_9_M/classification_results_plots/",
        #expand("output/differential_splicing/tissue_cell_type/{tct[0]}/{tct[1]}/splicing.clusters.significant.tsv", tct=tissue_cell_type_pairs),
        #expand("output/differential_splicing/tissue_cell_type/{tct[0]}/{tct[1]}/summary.tsv", tct=tissue_cell_type_pairs),
        #"output/adata_cellxgene.h5ad",
        #expand("output/coverage_track/tct/{tct[0]}/{tct[1]}/coverage.bw", tct=tissue_cell_type_pairs_including_singletons),
        #"output/differential_splicing/tissue_cell_type/Marrow/merged_significant_all_genes.svg",
        #"output/dimensionality_reduction/marrow/classification_score.svg",
        #expand("output/plots/{gene}.svg", gene=genes_to_plot),
        #expand("output/plots/{intron}.svg", intron=introns_to_plot),
        #expand("output/comparison/marrow_b/{label}.svg", label=labels),
        #"output/comparison/cortex_3_9_M/cell_ontology_class.svg",
        #"output/comparison/mg_3_38_F/cell_ontology_class.svg",
        #"output/comparison/mg_3_39_F/cell_ontology_class.svg",
        #expand("output/differential_splicing/tissue_cell_type/{tissue}/merged_significant_all.svg", tissue=pd.unique([t for t, ct in tissue_cell_type_pairs])),
        #expand("output/coverage_track/mg_basal_individual_plate/3_38_F/{plate}/coverage.bw", plate=["B002433", "B002432", "B002438"]),
        #expand(
        #    'output/differential_splicing/mg_basal_individual_plate/{individual}/{quantification}/splicing.clusters.csv',
        #    individual=["3_38_F"],
        #    quantification=["gene-expression"],
        #)
        #"output/coverage_track/mg_basal_individual_plate/3_56_F/B000167/coverage.bw",
        #"output/coverage_track/mg_basal_individual_plate/3_56_F/B000166/coverage.bw",
        #expand("output/bam_regions/Foxp1/{cell_type}/merged.bam", cell_type=genes_bam_merge["Foxp1"]["cell_types"]),


rule download:
    output: temp('output/srr/{srr}.sra')
    priority: 100
    shell: """
    prefetch -o {output} {wildcards.srr}
    """


rule fastq_dump:
    input:
        'output/srr/{srr}.sra'
    output:
        temp('output/srr/{srr}_1.fastq.gz'),
        temp('output/srr/{srr}_2.fastq.gz')
    shell: """
    cd output/srr && (fastq-dump --gzip --split-files {wildcards.srr}.sra || rm {wildcards.srr}.sra)
    """


def get_input_merge_fastq(wildcards):
    return expand('output/srr/{{sra}}_{number}.fastq.gz'.format(number=wildcards.number),
                  sra=sample_srr.srr[sample_srr.cell==wildcards.sample_id].values)


rule merge_fastq:
   input: get_input_merge_fastq
   output: temp('output/fastq_raw/{sample_id}_{number}.fastq.gz')
   shell: 'cat {input} > {output}'


rule trim_adapters:
    input:
        "output/fastq_raw/{sample_id}_1.fastq.gz",
        "output/fastq_raw/{sample_id}_2.fastq.gz",
    output:
        "output/fastq/{sample_id}_R1.fastq.gz",
        "output/fastq/{sample_id}_R2.fastq.gz",
    shell:
        "cutadapt -g {motif1} -g {motif2} -g {motif3} -a {motif1_rc} -a {motif2_rc} -a {motif3_rc} -G {motif1} -G {motif2} -G {motif3} -A {motif1_rc} -A {motif2_rc} -A {motif3_rc} -m30 -n 4 -o {output[0]} -p {output[1]} {input[0]} {input[1]}"


rule make_fastq_paths:
    input:
        expand("output/fastq/{sample_id}_R{pair}.fastq.gz", sample_id=sample_ids, pair=[1, 2]),
    output:
        "output/fastq_paths.txt"
    run:
        df = pd.DataFrame(sample_ids, columns=["sample_id"])
        base_path = os.path.join(os.getcwd(), "output/fastq/")
        df["fastq_1"] = base_path + df.sample_id + "_R1.fastq.gz"
        df["fastq_2"] = base_path + df.sample_id + "_R2.fastq.gz"
        df.to_csv(output[0], "\t", index=False, header=False)


rule read_mapping:
    input:
        "output/fastq_paths.txt",
        full_gtf_path
    threads: workflow.cores
    priority: 100
    output:
        expand("output/mapping/mapping_second_pass/{sample_id}/Aligned.rmdup.out.bam", sample_id=sample_ids)
    shell:
        "python -m scquint.quantification.run read_mapping/Snakefile --cores {threads} -d output/mapping/ --config min_cells_per_intron=30 fastq_paths=../fastq_paths.txt fasta_path={genome_fasta_path} gtf_path={full_gtf_path} sjdb_overhang=99"


rule make_fastq_paths_subset:
    input:
        expand("output/fastq/{sample_id}_R1.fastq.gz", sample_id=sample_ids_subset),
    output:
        "output/fastq_paths_subset.txt"
    run:
        df = pd.DataFrame(sample_ids_subset, columns=["sample_id"])
        base_path = os.path.join(os.getcwd(), "output/fastq/")
        df["fastq_1"] = base_path + df.sample_id + "_R1.fastq.gz"
        df["fastq_2"] = base_path + df.sample_id + "_R2.fastq.gz"
        df.to_csv(output[0], "\t", index=False, header=False)


rule kallisto_quantification:
    input:
        "output/fastq_paths_subset.txt",
        full_gtf_path,
    output:
        "output/quantification/kallisto/adata.h5ad"
    threads: workflow.cores
    shell:
        "python -m scquint.quantification.run kallisto/Snakefile --cores {threads} -d output/quantification/kallisto/ --config fastq_paths=../../fastq_paths_subset.txt genome_fasta_path={genome_fasta_path} gtf_path={full_gtf_path} min_cells_per_isoform=30"


rule bins_quantification:
    input:
        "output/bam_paths_subset.txt",
        full_gtf_path,
    threads: workflow.cores
    output:
        "output/quantification/bins/adata.h5ad"
    shell:
        "python -m scquint.quantification.run bins/Snakefile --cores all -d output/quantification/bins/ --config min_cells_per_bin=100 gtf_path={full_gtf_path} bam_paths=../../bam_paths_subset.txt"


rule transform_adata_with_NMF:
    input:
        "output/quantification/bins/adata_annotated.h5ad"
    output:
        "output/quantification/bins-nmf/adata_annotated.h5ad"
    run:
        original_adata = anndata.read_h5ad(input[0])
        cluster_gene_id_map = original_adata.var.groupby("cluster").gene_id.first()
        clusters = []
        gene_ids = []
        Xs = []
        n_clusters = len(pd.unique(original_adata.var.cluster))
        print("n_clusters: ", n_clusters)
        new_cluster = 0
        for cluster in pd.unique(original_adata.var.cluster):
            print(cluster)
            gene_id = cluster_gene_id_map.loc[cluster]
            idx_features = np.where(original_adata.var.cluster==cluster)[0]
            X = original_adata.X[:, idx_features].toarray()
            for n_components in [2, 5, 10]:
                X_NMF = NMF(n_components=n_components, max_iter=10000, solver="mu", beta_loss="frobenius").fit_transform(X)
                n_cells, n_features = X_NMF.shape
                Xs.append(X_NMF)
                clusters.append(np.full(n_features, new_cluster))
                new_cluster += 1
                gene_ids.append(np.full(n_features, gene_id))
        X = np.hstack(Xs)
        clusters = np.concatenate(clusters)
        gene_ids = np.concatenate(gene_ids)
        var = pd.DataFrame(dict(cluster=clusters,gene_id=gene_ids))
        print(var)
        adata = anndata.AnnData(X=X, var=var, obs=original_adata.obs)
        print(adata.shape)
        print("new n_clusters: ", adata.var.cluster.unique().shape)
        adata.write(output[0])


rule process_encode_blacklist:
    input:
        encode_blacklist_path
    output:
        "output/encode_blacklist.bed"
    shell:
        "set +o pipefail; cut -f1-3 {input} | bedtools sort -i stdin | uniq > {output}"


rule filter_bam:
    input:
        "output/mapping/mapping_second_pass/{sample_id}/Aligned.rmdup.out.bam",
        "output/encode_blacklist.bed"
    output:
        protected("output/mapping/filtered_bams/{sample_id}.bam"),
    shell:
        "bedtools intersect -split -sorted -a {input[0]} -b {input[1]} -v > {output}"


rule index_bam:
    input:
        "output/mapping/filtered_bams/{sample_id}.bam",
    output:
        "output/mapping/filtered_bams/{sample_id}.bam.bai"
    shell:
        "samtools index {input}"


rule make_bam_paths:
    input:
        expand("output/mapping/filtered_bams/{sample_id}.bam.bai", sample_id=sample_ids),
    output:
        "output/bam_paths.txt",
    run:
        cwd = os.getcwd()
        pd.DataFrame(dict(
            sample_id=sample_ids,
            bam_path=[f"{cwd}/output/mapping/filtered_bams/{sample_id}.bam"
                      for sample_id in sample_ids])
        ).to_csv(output[0], "\t", index=False, header=False)


rule make_bam_paths_subset:
    input:
        expand("output/mapping/filtered_bams/{sample_id}.bam", sample_id=sample_ids_subset),
    output:
        "output/bam_paths_subset.txt",
    run:
        cwd = os.getcwd()
        pd.DataFrame(dict(
            sample_id=sample_ids_subset,
            bam_path=[f"{cwd}/output/mapping/filtered_bams/{sample_id}.bam"
                      for sample_id in sample_ids_subset])
        ).to_csv(output[0], "\t", index=False, header=False)


rule prepare_chromosomes_file:
    input:
        full_gtf_path
    output:
        "output/chromosomes.txt"
    shell:
        "cut -f1 {input} | grep -v \# | grep -v GL | grep -v JH | sort | uniq > {output}"


rule add_metadata:
    input:
        "output/quantification/{quantification}/adata.h5ad",
    output:
        "output/quantification/{quantification,gene-expression|leafcutter|kallisto-SE|exons|bins|kallisto|introns-gene|introns-shared-acceptor|introns-nontransitive|introns-transitive}/adata_annotated.h5ad",
    run:
        adata = anndata.read_h5ad(input[0])
        print(adata.obs)
        adata.obs = sample_info.loc[adata.obs.index.values]
        print(adata.obs)
        adata.write(output[0], compression="gzip")


rule make_bam_paths_plate:
    input:
        expand("output/mapping/filtered_bams/{sample_id}.bam", sample_id=sample_ids)
    output:
        "output/bam_paths_{plate}.txt",
    run:
        s_ids = sample_ids[np.where(sample_info.plate_id==wildcards["plate"])[0]]
        print(wildcards["plate"], len(s_ids))
        cwd = os.getcwd()
        pd.DataFrame(dict(sample_id=s_ids, bam_path=[f"{cwd}/output/mapping/filtered_bams/{sample_id}.bam" for sample_id in s_ids])).to_csv(output[0], "\t", index=False, header=False)


rule make_coverage_track_bigwig_plate:
    input:
        "output/bam_paths_{plate}.txt",
    output:
        "output/coverage_track/plate/{plate}/coverage.bw",
    threads: workflow.cores // 4
    shell:
        "python -m scquint.quantification.run coverage_track/Snakefile --cores {threads} -d output/coverage_track/plate/{wildcards.plate}/ --config bam_paths=../../../bam_paths_{wildcards.plate}.txt chrom_sizes_path={chrom_sizes_path}"


rule intron_quantification:
    input:
        "output/bam_paths.txt",
        "output/chromosomes.txt",
        full_gtf_path,
        sjdb_path,
        chrom_sizes_path
    threads: workflow.cores
    output:
        "output/quantification/introns/output/introns-shared-acceptor/adata.h5ad",
    shell:
        "python -m scquint.quantification.run introns/Snakefile -q --cores {threads} -d output/quantification/introns/ --config min_cells_per_intron=100 bam_paths=../../bam_paths.txt chromosomes_path=../../chromosomes.txt fasta_path={genome_fasta_path} gtf_path={full_gtf_path} chrom_sizes_path={chrom_sizes_path} sjdb_path={sjdb_path}"


rule extract_intron_quantification:
    input:
        "output/quantification/introns/output/introns-{grouping}/adata.h5ad"
    output:
        "output/quantification/introns-{grouping}/adata.h5ad"
    shell:
        "cp {input} {output}"


rule gene_expression_quantification:
    input:
        "output/bam_paths.txt",
        full_gtf_path,
    threads: workflow.cores
    output:
        "output/quantification/gene-expression/adata.h5ad"
    shell:
        "python -m scquint.quantification.run genes/Snakefile --cores all -q -d output/quantification/gene-expression/ --config min_cells_per_gene=100 gtf_path={full_gtf_path} bam_paths=../../bam_paths.txt"


rule differential_test_mg_basal_individual_plate:
    input:
        "output/quantification/{quantification}/adata_annotated.h5ad",
    output:
        'output/differential_splicing/mg_basal_individual_plate/{individual}/{quantification}/splicing.clusters.csv',
        'output/differential_splicing/mg_basal_individual_plate/{individual}/{quantification}/splicing.introns.csv',
    run:
        print("threads: ", threads)
        obs = sample_info[
            (sample_info.tissue=="Mammary_Gland") &
            (sample_info.cell_ontology_class=="basal cell") &
            (sample_info["mouse.id"]==wildcards["individual"])
        ].sort_index()
        adata = anndata.read_h5ad(input[0])
        adata = adata[obs.index.values]
        plates = obs.plate_id.unique()
        print(plates)
        #assert(len(plates)==2)
        #cell_idx_a = np.where(
        #    (obs.plate_id==plates[0])
        #)[0]
        #cell_idx_b = np.where(
        #    (obs.plate_id!=plates[0])
        #)[0]
        print("hardcoding B002438")
        cell_idx_a = np.where(
            (obs.plate_id=="B002438")
        )[0]
        cell_idx_b = np.where(
            (obs.plate_id!="B002438")
        )[0]

        permute = False
        if permute:
            cell_idx_all = np.concatenate([cell_idx_a, cell_idx_b])
            cell_idx_all_p = np.random.permutation(cell_idx_all)
            cell_idx_a_p = cell_idx_all_p[:len(cell_idx_a)]
            cell_idx_b_p = cell_idx_all_p[len(cell_idx_a):]
            cell_idx_a = cell_idx_a_p
            cell_idx_b = cell_idx_b_p

        if wildcards["quantification"] != "gene-expression":
            adata.var["original_cluster"] = adata.var.cluster
            diff_spl_clusters, diff_spl_introns = run_differential_splicing(
                adata,
                "permutation-Euclidean",
                cell_idx_a,
                cell_idx_b,
                min_cells_per_cluster=30 if wildcards["quantification"] != "bins-nmf" else None,
                min_total_cells_per_intron=30 if wildcards["quantification"] != "bins-nmf" else None,
                device="cuda:0",
                #device="cpu",
                n_permutations=100000,
               )
            diff_spl_clusters.to_csv(output[0], '\t')
            diff_spl_introns.to_csv(output[1], '\t')
        else:
            diff_exp = run_differential_expression(adata, cell_idx_a, cell_idx_b, 30)
            diff_exp.to_csv(output[0], '\t', index=False)
            diff_exp.to_csv(output[1], '\t', index=False)


rule make_bam_paths_mg_basal_individual_plate:
    input:
        "output/bam_paths.txt",
    output:
        "output/bam_paths/mg_basal_individual_{individual}_{plate}.txt",
    run:
        s_ids = sample_ids[
            np.where(
                (sample_info.tissue=="Mammary_Gland") &
                (sample_info.cell_ontology_class=="basal cell") &
                (sample_info["mouse.id"]==wildcards["individual"]) &
                (sample_info.plate_id==wildcards["plate"])
            )[0]
        ]
        print(wildcards["plate"], len(s_ids))
        cwd = os.getcwd()
        pd.DataFrame(dict(sample_id=s_ids, bam_path=[f"{cwd}/output/mapping/filtered_bams/{sample_id}.bam" for sample_id in s_ids])).to_csv(output[0], "\t", index=False, header=False)


rule make_coverage_track_bigwig_mg_basal_individual_plate:
    input:
        "output/bam_paths/mg_basal_individual_{individual}_{plate}.txt",
    output:
        "output/coverage_track/mg_basal_individual_plate/{individual}/{plate}/coverage.bw",
    threads: workflow.cores // 4
    shell:
        "python -m scquint.quantification.run coverage_track/Snakefile --cores {threads} -d output/coverage_track/mg_basal_individual_plate/{wildcards.individual}/{wildcards.plate}/ --config bam_paths=../../../../bam_paths/mg_basal_individual_{wildcards.individual}_{wildcards.plate}.txt chrom_sizes_path={chrom_sizes_path}"


rule dimensionality_reduction_mg_pca:
    input:
        "output/quantification/{quantification}/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/mg/{quantification}/pca_{min_cells}_{k}/latent.txt",
    wildcard_constraints: quantification='.+(?<!gene-expression)'
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue=="Mammary_Gland")]
        sorted_index = sorted(adata.obs.index)
        adata = adata[sorted_index]
        if wildcards["quantification"] != "bins-nmf":
            adata = filter_min_cells_per_feature(adata, int(wildcards["min_cells"]))
        print(adata.shape)
        if wildcards["quantification"] == "introns-transitive" or wildcards["quantification"] == "leafcutter":
            adata = recluster(adata)
        print(adata.shape)
        latent = run_pca(adata, int(wildcards["k"]))
        np.savetxt(output[0], latent)


rule dimensionality_reduction_mg_vae:
    input:
        "output/quantification/{quantification}/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/mg/{quantification}/vae_{n_cells}/latent.txt",
    wildcard_constraints: quantification='.+(?<!gene-expression)'
    resources:
        gpu=1
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue=="Mammary_Gland")]
        sorted_index = sorted(adata.obs.index)
        adata = adata[sorted_index]
        if wildcards["quantification"] != "bins-nmf":
            adata = filter_min_cells_per_feature(adata, int(wildcards["n_cells"]))
        if np.isin(wildcards["quantification"], ["introns-transitive", "introns-transitive-shuffled", "leafcutter", "introns-subset-lc", "introns-subset-lc-nonfiltered"]):
            adata = recluster(adata)
        print(wildcards["quantification"], adata.shape)
        feature_addition = group_normalize(adata.X.sum(axis=0), adata.var.cluster, smooth=False).A1.ravel()

        latent, model = run_vae(
            adata, n_epochs_kl_warmup=30, regularization_gaussian_std=10.0,
            use_cuda=True, input_transform="frequency-smoothed",
            feature_addition=feature_addition, sample=False,
        )
        np.savetxt(output[0], latent)

rule dimensionality_reduction_mg_expression:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/mg/gene-expression/pca_{min_cells}_{k}/latent.txt",
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue=="Mammary_Gland")]
        sorted_index = sorted(adata.obs.index)
        adata = adata[sorted_index]
        sc.pp.filter_genes(adata, min_cells=int(wildcards["min_cells"]))
        sc.pp.normalize_total(adata, target_sum=1e4)
        sc.pp.log1p(adata)
        X = adata.X.toarray()
        latent = PCA(n_components=int(wildcards["k"])).fit_transform(X)
        np.savetxt(output[0], latent)


rule dimensionality_reduction_mg_3_38_F_pca:
    input:
        "output/quantification/{quantification}/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/mg_3_38_F/{quantification}/pca_{k}/latent.txt",
    wildcard_constraints: quantification='.+(?<!gene-expression)'
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue=="Mammary_Gland") & (adata.obs["mouse.id"]=="3_38_F")]
        sorted_index = sorted(adata.obs.index)
        adata = adata[sorted_index]
        if wildcards["quantification"] != "bins-nmf":
            adata = filter_min_cells_per_feature(adata, 50)
        print(adata.shape)
        latent = run_pca(adata, int(wildcards["k"]))
        np.savetxt(output[0], latent)


rule dimensionality_reduction_mg_3_38_F_expression:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/mg_3_38_F/gene-expression/pca_{k}/latent.txt",
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue=="Mammary_Gland") & (adata.obs["mouse.id"]=="3_38_F")]
        sorted_index = sorted(adata.obs.index)
        adata = adata[sorted_index]
        sc.pp.filter_genes(adata, min_cells=50)
        sc.pp.normalize_total(adata, target_sum=1e4)
        sc.pp.log1p(adata)
        X = adata.X.toarray()
        latent = PCA(n_components=int(wildcards["k"])).fit_transform(X)
        np.savetxt(output[0], latent)


rule dimensionality_reduction_all_pca:
    input:
        "output/quantification/{quantification}/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/all/{quantification}/pca_{n_cells}_{k}/latent.txt",
    wildcard_constraints: quantification='.+(?<!gene-expression)'
    run:
        adata = anndata.read_h5ad(input[0])
        if wildcards["quantification"] != "bins-nmf":
            adata = filter_min_cells_per_feature(adata, int(wildcards["n_cells"]))
        print(adata.shape)
        if wildcards["quantification"] == "introns-transitive" or wildcards["quantification"] == "leafcutter":
            adata = recluster(adata)
        print(adata.shape)
        latent = run_pca(adata, int(wildcards["k"]))
        print("latent.shape: ", latent.shape)
        np.savetxt(output[0], latent)


rule dimensionality_reduction_all_vae:
    input:
        "output/quantification/{quantification}/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/all/{quantification}/vae_{transform}_{sample}_{min_cells}/latent.txt",
    wildcard_constraints: quantification='.+(?<!gene-expression)'
    resources:
        gpu=1
    run:
        adata = anndata.read_h5ad(input[0])
        if wildcards["quantification"] != "bins-nmf":
            adata = filter_min_cells_per_feature(adata, int(wildcards["min_cells"]))
        if np.isin(wildcards["quantification"], ["introns-transitive", "introns-transitive-shuffled", "leafcutter", "introns-subset-lc", "introns-subset-lc-nonfiltered"]):
            adata = recluster(adata)
        print(wildcards["quantification"], adata.shape)
        input_transform = wildcards["transform"]
        if input_transform == "frequency-smoothed":
            feature_addition = group_normalize(adata.X.sum(axis=0), adata.var.cluster, smooth=False).A1.ravel()
            print(feature_addition.shape)
        else:
            feature_addition = None
        latent, model = run_vae(
            adata, n_epochs_kl_warmup=30, regularization_gaussian_std=None,
            use_cuda=True, input_transform=input_transform,
            feature_addition=feature_addition, sample=wildcards["sample"]=="True" or wildcards["sample"]=="true",
        )
        np.savetxt(output[0], latent)


rule dimensionality_reduction_all_vae_hyperopt:
    input:
        "output/quantification/{quantification}/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/all/{quantification}/vae_hyperopt/latent.txt",
    wildcard_constraints: quantification='.+(?<!gene-expression)'
    resources:
        gpu=1
    run:
        adata = anndata.read_h5ad(input[0])
        if wildcards["quantification"] != "bins-nmf":
            adata = filter_min_cells_per_feature(adata, 100)
        if np.isin(wildcards["quantification"], ["introns-transitive", "introns-transitive-shuffled", "leafcutter", "introns-subset-lc", "introns-subset-lc-nonfiltered"]):
            adata = recluster(adata)
        print(wildcards["quantification"], adata.shape)
        feature_addition = group_normalize(adata.X.sum(axis=0), adata.var.cluster, smooth=False).A1.ravel()

        latent, model = run_vae(
            adata, n_epochs_kl_warmup=20, regularization_gaussian_std=None,
            use_cuda=True, input_transform="frequency-smoothed",
            feature_addition=feature_addition, sample=False, linearity="non-linear",
            n_layers=2, n_latent=34, dropout_rate=0.224,
        )
        np.savetxt(output[0], latent)


rule dimensionality_reduction_all_expression:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/all/gene-expression/pca_{k}/latent.txt",
    run:
        adata = anndata.read_h5ad(input[0])
        print(adata.shape)
        sc.pp.filter_genes(adata, min_cells=100)
        print(adata.shape)
        sc.pp.normalize_total(adata, target_sum=1e4)
        sc.pp.log1p(adata)
        sc.pp.highly_variable_genes(adata)
        adata = adata[:, adata.var.highly_variable]
        print(adata.shape)

        X = adata.X.toarray()
        latent = PCA(n_components=int(wildcards["k"])).fit_transform(X)
        np.savetxt(output[0], latent)


rule filter_bam_to_region:
    input:
        "output/mapping/filtered_bams/{sample_id}.bam"
    output:
        "output/bam_regions/{gene}/{sample_id}.bam"
    run:
        chromosome, start, end = genes_bam_merge[wildcards["gene"]]["region"]
        shell(f"samtools view -b -o {output} {input} {chromosome}:{start}-{end}")


rule make_bam_paths_gene:
    input:
        lambda wildcards: expand(f"output/bam_regions/{wildcards['gene']}/{{sample_id}}.bam", sample_id=sample_ids[(sample_info.cell_type==wildcards["cell_type"])])
    output:
        "output/bam_paths-{gene}-{cell_type}.txt"
    run:
        pd.DataFrame(input).to_csv(output[0], "\t", index=False, header=False)


rule merge_bams:
    input:
        "output/bam_paths-{gene}-{cell_type}.txt"
    output:
        "output/bam_regions/{gene}/{cell_type}/merged.bam"
    threads:
        workflow.cores // 2
    priority: 10
    shell:
        "samtools merge --threads {threads} -b {input} {output}"


rule differential_test_tissue_cell_type:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
    output:
        'output/differential_splicing/tissue_cell_type/{tissue}/{cell_type}/expression.tsv',
        'output/differential_splicing/tissue_cell_type/{tissue}/{cell_type}/splicing.clusters.tsv',
        'output/differential_splicing/tissue_cell_type/{tissue}/{cell_type}/splicing.introns.tsv',
    threads: workflow.cores // 4
    #threads: workflow.cores
    run:
        adata_exp = anndata.read_h5ad(input[0])
        adata_exp.obs.index = adata_exp.obs.index.astype(str)
        adata_spl = anndata.read_h5ad(input[1])
        adata_spl.obs.index = adata_spl.obs.index.astype(str)
        adata_spl.var.index = adata_spl.var.index.astype(str)
        assert((adata_exp.obs.index == adata_spl.obs.index).all())
        obs = adata_exp.obs
        obs.loc[:, "cell_type"] = obs.cell_ontology_class.str.replace(" ", "_").str.replace("-", "_")
        cell_idx_a = np.where((obs.tissue==wildcards["tissue"]) &
                              (obs.cell_type==wildcards["cell_type"]))[0]
        cell_idx_b = np.where((obs.tissue==wildcards["tissue"]) &
                              (obs.cell_type!=wildcards["cell_type"]) &
                              (obs.cell_type.isin([ct for t, ct in tissue_cell_type_pairs if t==wildcards["tissue"]])))[0]
        MIN_FEATURES = 50
        diff_exp = run_differential_expression(adata_exp, cell_idx_a, cell_idx_b, MIN_FEATURES)
        diff_exp.to_csv(output[0], '\t', index=False)
        diff_spl_clusters, diff_spl_introns = run_differential_splicing(
            adata_spl,
            cell_idx_a,
            cell_idx_b,
            min_cells_per_cluster=MIN_FEATURES,
            min_total_cells_per_intron=MIN_FEATURES,
            n_jobs=threads,
            do_recluster=False,
        )
        diff_spl_clusters.to_csv(output[1], '\t')
        diff_spl_introns.to_csv(output[2], '\t')


rule extract_gene_cds:
    input:
        full_gtf_path
    output:
        "output/gene_cds.txt"
    run:
        df = pd.read_csv(
            input[0], '\t', header=None, comment="#",
            names=['chromosome', 'source', 'feature', 'start', 'end', 'score', 'strand', 'frame', 'attribute'],
           )
        print(df.shape)
        df = df[df.feature=="CDS"]
        print(df.shape)
        df['gene_id'] = df.attribute.str.extract(r'gene_id "([^;]*)";')
        res = df.groupby("gene_id").agg({"chromosome": "first", "start": "min", "end": "max", "strand": "first"})
        print(res)
        res.to_csv(output[0], "\t")


rule extract_gene_name:
    input:
        full_gtf_path
    output:
        "output/gene_name.txt"
    run:
        df = pd.read_csv(
            input[0], '\t', header=None, comment="#",
            names=['chromosome', 'source', 'feature', 'start', 'end', 'score', 'strand', 'frame', 'attribute'],
           )
        print(df.shape)
        df = df[df.feature=="gene"]
        print(df.shape)
        df['gene_id'] = df.attribute.str.extract(r'gene_id "([^;]*)";')
        df['gene_name'] = df.attribute.str.extract(r'gene_name "([^;]*)";')
        res = df.groupby("gene_id").gene_name.first()
        print(res)
        res.to_csv(output[0], "\t")


rule filter_diff_spl_significant:
    input:
        'output/{anything}/expression.tsv',
        'output/{anything}/splicing.clusters.tsv',
        'output/{anything}/splicing.introns.tsv',
        "output/gene_cds.txt",
        "output/gene_name.txt",
    output:
        'output/{anything}/splicing.clusters.significant.tsv',
    run:
        diff_exp = pd.read_csv(input[0], "\t", index_col=0)
        diff_spl_cluster = pd.read_csv(input[1], "\t", index_col=0)
        diff_spl_intron = pd.read_csv(input[2], "\t", index_col=0)
        gene_cds = pd.read_csv(input[3], "\t", index_col=0)
        gene_name = pd.read_csv(input[4], "\t", index_col=0)

        assert(set(diff_spl_cluster.index.values) == set(diff_spl_intron.cluster.unique()))
        cluster_original_cluster = diff_spl_intron[["cluster", "original_cluster"]].drop_duplicates()

        print(diff_spl_cluster.shape)
        diff_spl_cluster = diff_spl_cluster[
            ((diff_spl_cluster.p_value_adj <= config["fdr"]) &
             (diff_spl_cluster.max_abs_delta_psi >= config["min_abs_delta_psi"]))
        ]
        print(diff_spl_cluster.shape)
        diff_spl_cluster = diff_spl_cluster.merge(gene_name, how="left", left_on="gene_id", right_index=True)
        #print(diff_spl_cluster.shape)
        def max_abs_lfc_psi_unannotated(row_cluster):
            introns = diff_spl_intron[diff_spl_intron.cluster==row_cluster.name]
            return (introns.abs_lfc_psi * (~introns.annotated).astype(int)).max()
        diff_spl_cluster["max_abs_lfc_psi_unannotated"] = diff_spl_cluster.apply(max_abs_lfc_psi_unannotated, axis=1)
        def max_abs_delta_psi_unannotated(row_cluster):
            introns = diff_spl_intron[diff_spl_intron.cluster==row_cluster.name]
            return (introns.abs_delta_psi * (~introns.annotated).astype(int)).max()
        diff_spl_cluster["max_abs_delta_psi_unannotated"] = diff_spl_cluster.apply(max_abs_delta_psi_unannotated, axis=1)
        diff_spl_cluster["Annotated"] = diff_spl_cluster.max_abs_delta_psi_unannotated < config["min_abs_delta_psi"]

        #coordinates = diff_spl_intron.groupby("cluster").agg({"chromosome": "first", "start": "unique", "end": "unique"})
        #diff_spl_cluster = diff_spl_cluster.merge(coordinates, how="left", left_index=True, right_index=True)
        diff_spl_intron["coordinates"] = diff_spl_intron.chromosome + ":" + diff_spl_intron.start.astype(str) + "-" + diff_spl_intron.end.astype(str)
        coordinates = diff_spl_intron.groupby("cluster").coordinates.unique()
        diff_spl_cluster = diff_spl_cluster.merge(coordinates, how="left", left_index=True, right_index=True)

        print(diff_spl_cluster.Annotated.value_counts())
        def get_diff_exp_rank(row_cluster):
            try:
                return diff_exp.loc[row_cluster.gene_id].ranking
            except KeyError:
                return np.nan
        diff_spl_cluster["diff_exp_rank"] = diff_spl_cluster.apply(get_diff_exp_rank, axis=1)
        def check_region(row_cluster):
            try:
                cds = gene_cds.loc[row_cluster.gene_id]
            except KeyError:
                return "Non-coding"
            introns = diff_spl_intron[diff_spl_intron.cluster==row_cluster.name]
            cluster_start = introns.start.min()
            cluster_end = introns.end.max()
            if cds.strand == "+":
                if cluster_start < cds.start:
                    return "5' UTR"
                if cluster_start < cds.end:
                    return "CDS"
                return "3' UTR"
            elif cds.strand == "-":
                if cluster_start < cds.start:
                    return "3' UTR"
                if cluster_start < cds.end:
                    return "CDS"
                return "5' UTR"
            else:
                raise Exception("strand not implemented")
        diff_spl_cluster["Region"] = diff_spl_cluster.apply(check_region, axis=1)
        diff_spl_cluster.max_abs_delta_psi = diff_spl_cluster.max_abs_delta_psi.round(decimals=3)

        diff_spl_cluster = diff_spl_cluster.merge(cluster_original_cluster, how="left", on="cluster")

        np.set_printoptions(linewidth=100000)
        diff_spl_cluster.to_csv(
            output[0], "\t",
            columns=["original_cluster", "gene_id", "gene_name", "p_value", "p_value_adj", "max_abs_delta_psi", "Annotated", "Region", "diff_exp_rank", "coordinates"],
        )


rule merge_significant:
    input:
        lambda wildcards: expand('output/differential_splicing/tissue_cell_type/{{tissue}}/{cell_type}/splicing.clusters.significant.tsv', cell_type=[ct for t, ct in tissue_cell_type_pairs if t==wildcards["tissue"]])
    output:
        "output/differential_splicing/tissue_cell_type/{tissue}/merged_significant.tsv",
    run:
        dfs = []
        for cell_type, input_path in zip([ct for t, ct in tissue_cell_type_pairs if t==wildcards["tissue"]], input):
            df = pd.read_csv(input_path, "\t")
            df["Cell type"] = cell_type.replace("_", " ")
            dfs.append(df)
        df = pd.concat(dfs, ignore_index=True)
        df = df.sort_values("p_value_adj")
        df.to_csv(output[0], "\t", index=False)


rule plot_significant:
    input:
        "output/differential_splicing/tissue_cell_type/{tissue}/merged_significant.tsv",
    output:
        "output/differential_splicing/tissue_cell_type/{tissue}/merged_significant_all.svg",
        "output/differential_splicing/tissue_cell_type/{tissue}/merged_significant_all_genes.svg",
        "output/differential_splicing/tissue_cell_type/{tissue}/merged_significant_aggregate.svg",
    run:
        df_all = pd.read_csv(input[0], "\t")
        df_all["Type"] = df_all.Annotated
        df_all.Type = df_all.Type.replace(True, "Annotated").replace(False, "Novel")
        df_all.Region = df_all.Region.replace("Non-coding", "Non-coding RNA")
        df_all["Cell type"] = df_all["Cell type"].replace("hematopoietic stem cell", "HSC")
        df_all["Cell type"] = df_all["Cell type"].replace("granulocytopoietic cell", "granulocytop.")

        def plot_counts(df, output_path, ylabel="Diff. spl. events"):
            df_plot = df.groupby(["Type", "Cell type"]).size().reset_index().pivot(columns='Type', index='Cell type', values=0)
            g = df_plot.plot(kind='bar', stacked=True, figsize=(4,3))
            g.set_xticklabels(g.get_xticklabels(), rotation=45, horizontalalignment='right')
            g.set(xlabel="", ylabel=ylabel)
            #g.despine(left=True)
            sns.despine()
            plt.legend(loc='upper right')
            plt.savefig(output_path, bbox_inches='tight')
            plt.close()
        plot_counts(df_all, output[0])
        print(df_all.sort_values("Type", ascending=False).Type)
        plot_counts(df_all.sort_values("Type", ascending=False).drop_duplicates(["Cell type", "gene_id"]), output[1], ylabel="Diff. spl. genes")

        #df = df_all
        df = df_all.drop_duplicates("original_cluster")
        df_plot = df.groupby(["Type", "Region"]).size().reset_index().pivot(columns='Type', index='Region', values=0).loc[["5' UTR", "CDS", "3' UTR", "Non-coding RNA"]]
        g = df_plot.plot(kind='bar', stacked=True, figsize=(3,3))
        g.set_xticklabels(g.get_xticklabels(), rotation=45, horizontalalignment='right')
        g.set(xlabel="", ylabel='Diff. spl. events')
        sns.despine()
        plt.legend(loc='upper right')
        plt.savefig(output[2], bbox_inches='tight')


rule compare_latent_custom_plot:
    input:
        #expand("output/dimensionality_reduction/mg/{quantification}/pca_50_10/latent.txt", quantification=quantifications_mg_individual_expanded),
        "output/dimensionality_reduction/mg/gene-expression/pca_50_10/latent.txt",
        "output/dimensionality_reduction/mg/kallisto/vae_30/latent.txt",
        "output/dimensionality_reduction/mg/bins/vae_30/latent.txt",
        "output/dimensionality_reduction/mg/exons/vae_30/latent.txt",
        "output/dimensionality_reduction/mg/introns-gene/vae_30/latent.txt",
        "output/dimensionality_reduction/mg/leafcutter/vae_30/latent.txt",
        "output/dimensionality_reduction/mg/introns-shared-acceptor/vae_30/latent.txt",
        #"output/dimensionality_reduction/mg/leafcutter/pca_30_10/latent.txt",
        #"output/dimensionality_reduction/mg/leafcutter/vae_30/latent.txt",
        #"output/dimensionality_reduction/mg/leafcutter/vae_50/latent.txt",
        #"output/dimensionality_reduction/mg/introns-shared-acceptor/pca_30_10/latent.txt",
        #"output/dimensionality_reduction/mg/introns-shared-acceptor/vae_30/latent.txt",
        #"output/dimensionality_reduction/mg/introns-shared-acceptor/vae_50/latent.txt",
        #"output/dimensionality_reduction/mg/leafcutter/vae_100/latent.txt",
    output:
        #"output/comparison/mg_3_38_F/cell_ontology_class.svg",
        "output/comparison/mg_{individual}/cell_ontology_class.svg",
        "output/comparison/mg_{individual}/plate.svg",
    run:
        obs = sample_info[(sample_info.tissue=="Mammary_Gland")]
        obs = obs.sort_index()
        idx = np.where((obs["mouse.id"]==wildcards["individual"]) & (obs.cell_ontology_class!="endothelial cell"))[0]
        obs = obs.iloc[idx]
        obs["Cell type"] = obs.cell_ontology_class
        obs["Plate ID"] = obs.plate_id
        dfs = []
        #names = ["Leafcutter-PCA", "LeafCutter-VAE", "LeafCutter-Ours-PCA", "LeafCutter-Ours-VAE", "SharedAcceptor-PCA", "SharedAcceptor-VAE"]
        #names = ["Leafcutter", "SharedAcceptor", "AnnotatedSE", "AnnotatedSE-SharedAcceptor", "AnnotatedSE-SharedDonor"]
        names = ["Gene expression", "kallisto", "ODEGR-NMF", "DEXSeq", "DESJ", "LeafCutter", "scQuint"]
        #names = ["LeafCutter-VAE-30", "LeafCutter-VAE-50", "LeafCutter-VAE-100"]
        #names = ["LeafCutter-VAE-30", "LeafCutter-VAE-50", "scQuint-VAE-30", "scQuint-VAE-50"]
        #names = ["LeafCutter-PCA", "LeafCutter-VAE", "scQuint-PCA", "scQuint-VAE"]
        for name, input_path in zip(names, input):
            print(name)
            #if name == "gene-expression":
            #    #name = "Gene expression \n (featureCounts)"
            #    name = "Gene expression"
            #if name == "kallisto":
            #    pass
            #    #name = "Isoform proportions \n (kallisto)"
            #if name == "bins":
            #    #name = "100-bp bin cov. proportions \n (ODEGR-NMF)"
            #    name = "ODEGR-NMF"
            #if name == "introns-shared-acceptor":
            #    #name = "Alt. intron proportions \n (scQuint)"
            #    name = "scQuint"
            #if name == "introns-transitive":
            #    name = "LeafCutter (ours)"
            #if name == "introns-gene":
            #    name = "DESJ"
            #if name == "SE":
            #    name = "Skipped exons"
            #if name == "exons":
            #    name = "DEXSeq"
            df = obs.copy()
            latent = np.loadtxt(input_path)
            latent = latent[idx]
            proj = UMAP(min_dist=0.5, n_neighbors=15, random_state=42).fit_transform(latent)
            df["UMAP 1"] = proj[:, 0]
            df["UMAP 2"] = proj[:, 1]
            df["Quantification"] = name
            dfs.append(df)
        df = pd.concat(dfs)
        df = df.replace("luminal epithelial cell of mammary gland", "luminal\nepithelial")
        df = df.replace("basal cell", "basal")
        df = df.replace("stromal cell", "stromal")
        g = sns.relplot(
            data=df, x="UMAP 1", y="UMAP 2",
            row="Quantification",
            #row_order=["Gene expression", "kallisto", "DEXSeq", "ODEGR-NMF", "DESJ", "LeafCutter", "Skipped exons", "scQuint"],
            #row_order=["Gene expression", "kallisto", "DEXSeq", "ODEGR-NMF", "DESJ", "LeafCutter"],
            row_order=names,
            hue="Cell type",
            kind="scatter",
            facet_kws={'sharey': False, 'sharex': False},
            height=1.4, palette="tab10", edgecolor="none", s=6,
            aspect=1.4,
        )
        g.set_titles(row_template="{row_name}")
        g.fig.subplots_adjust(hspace=0.7)
        for ax in g.axes.flat:
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_xlabel("UMAP 1")
            ax.set_ylabel("UMAP 2")
        sns.despine()
        #leg = g._legend
        #leg.set_bbox_to_anchor([0.5, 1.0])
        #leg._loc = 8
        plt.savefig(output[0], bbox_inches='tight')
        plt.close()
        g = sns.relplot(
            data=df, x="UMAP 1", y="UMAP 2",
            row="Quantification",
            #row_order=["Gene expression", "kallisto", "DEXSeq", "ODEGR-NMF", "DESJ", "LeafCutter", "Skipped exons", "scQuint"],
            #row_order=["Gene expression", "kallisto", "DEXSeq", "ODEGR-NMF", "DESJ", "LeafCutter"],
            row_order=names,
            hue="Plate ID",
            kind="scatter",
            facet_kws={'sharey': False, 'sharex': False, "margin_titles": False},
            height=1.4, palette=["C3", "C8", "C9"], edgecolor="none", s=6,
            aspect=1.4,
            #legend=False,
        )
        g.set_titles(row_template="{row_name}")
        g.fig.subplots_adjust(hspace=0.7)
        for ax in g.axes.flat:
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_xlabel("UMAP 1")
            ax.set_ylabel("UMAP 2")
        sns.despine()
        #leg = g._legend
        #leg.set_bbox_to_anchor([0.5, 1.0])
        #leg._loc = 8
        #plt.savefig(output[1])
        #plt.savefig(output[1], bbox_extra_artists=(leg,), bbox_inches='tight')
        #plt.savefig(output[1], bbox_extra_artists=(leg,))
        #g.get_figure().savefig(output[1], bbox_inches='tight')
        #plt.subplots_adjust(top=0.95, right=0.95)
        #g.savefig(output[1], bbox_inches='tight', bbox_extra_artists=(leg,))
        plt.savefig(output[1], bbox_inches='tight')
        plt.close()


rule dimensionality_reduction_marrow_b_expression:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/marrow_b/gene-expression/PCA/latent.txt",
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue=="Marrow") & (adata.obs.cell_ontology_class.str.endswith("B cell"))]
        sc.pp.filter_genes(adata, min_cells=30)
        sc.pp.normalize_total(adata, target_sum=1e4)
        sc.pp.log1p(adata)
        X = adata.X.toarray()
        latent = PCA(n_components=20).fit_transform(X)
        np.savetxt(output[0], latent)


rule dimensionality_reduction_marrow_b_pca:
    input:
        "output/quantification/{quantification}/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/marrow_b/{quantification}/PCA_{K}/latent.txt",
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue=="Marrow") & (adata.obs.cell_ontology_class.str.endswith("B cell"))]
        adata = filter_min_cells_per_feature(adata, 30)
        if wildcards["quantification"] == "introns-transitive":
            adata = recluster(adata)
        latent = run_pca(adata, int(wildcards["K"]))
        np.savetxt(output[0], latent)


rule compare_latent_marrow_b:
    input:
        #"output/dimensionality_reduction/marrow_b/gene-expression/PCA/latent.txt",
        #"output/dimensionality_reduction/marrow_b/introns-transitive/PCA_10/latent.txt",
        #"output/dimensionality_reduction/marrow_b/introns-shared-acceptor/PCA_20/latent.txt",
        "output/dimensionality_reduction/all/gene-expression/pca_40/latent.txt",
        "output/dimensionality_reduction/all/introns-shared-acceptor/vae_hyperopt/latent.txt",
    output:
        expand("output/comparison/marrow_b/{label}.svg", label=labels)
    run:
        obs = sample_info.copy()
        idx = np.where((obs.tissue=="Marrow") & (obs.cell_ontology_class.str.endswith("B cell")))[0]
        obs = obs.iloc[idx]

        obs = obs.replace("late pro-B cell", "(1) pro-B")
        obs = obs.replace("precursor B cell", "(2) pre-B")
        obs = obs.replace("immature B cell", "(3) immature B")
        obs = obs.replace("naive B cell", "(4) naive B")

        dfs = []
        for name, input_path in zip(["Expression latent space", "Splicing latent space"], input):
            df = obs.copy()
            latent = np.loadtxt(input_path)
            latent = latent[idx]

            proj = UMAP(min_dist=0.5, n_neighbors=15, random_state=42).fit_transform(latent)
            df["UMAP 1"] = proj[:, 0]
            df["UMAP 2"] = proj[:, 1]
            df["Quantification"] = name
            df = df[df.cell_ontology_class != "early pro-B cell"]
            df = df.sort_values("cell_ontology_class")
            dfs.append(df)
        df = pd.concat(dfs)
        for i, label in enumerate(labels):
            g = sns.relplot(
                data=df, x="UMAP 1", y="UMAP 2",
                col="Quantification", hue=label,
                kind="scatter", facet_kws={'sharey': False, 'sharex': False},
                height=3, palette="tab10", edgecolor="none", s=4,
               )
            g.set_titles(col_template="{col_name}")
            g.fig.subplots_adjust(wspace=0.1)
            for ax in g.axes.flat:
                ax.set_xticks([])
                ax.set_yticks([])
                ax.set_ylabel("UMAP 2")
            sns.despine()
            plt.savefig(output[i], bbox_inches='tight')


rule psi_plot_marrow_b:
    input:
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
    output:
        expand("output/plots/{intron}.pdf", intron=introns_to_plot),
    run:
        adata = anndata.read_h5ad(input[0])
        adata.obs = adata.obs.replace("late pro-B cell", "(1) pro-B")
        adata.obs = adata.obs.replace("precursor B cell", "(2) pre-B")
        adata.obs = adata.obs.replace("immature B cell", "(3) immature B")
        adata.obs = adata.obs.replace("naive B cell", "(4) naive B")
        adata.obs["Cell type"] = adata.obs.cell_ontology_class
        adata = adata[adata.obs["Cell type"].isin(cell_type_order)]
        X = group_normalize(adata.X.toarray(), adata.var.cluster.values, smooth=False)
        var = adata.var.copy()
        var["position"] = np.arange(len(var))
        var["id"] = var.chromosome.astype(str) + ":" + var.start.astype(str) + "-" + var.end.astype(str)
        var = var.set_index("id")
        for i, intron in enumerate(introns_to_plot):
            obs = adata.obs.copy()
            obs["PSI"] = X[:, var.loc[intron].position].ravel()
            eps = 1e-4

            obs["external_donor_name"] = obs["mouse.id"]

            obs_all = obs.copy()
            obs_all.external_donor_name = "All"
            most_common_individuals = obs.external_donor_name.value_counts()[:6].index.tolist()
            obs = obs[obs.external_donor_name.isin(most_common_individuals)]
            obs = pd.concat([obs_all, obs], ignore_index=True)
            most_common_individuals = ["All"] + most_common_individuals

            for x in cell_type_order:
                for y in most_common_individuals:
                    mask = (obs["Cell type"] == x) & (obs.external_donor_name == y)
                    if obs.loc[mask, "PSI"].notna().sum() < 3:
                        obs.loc[mask, "PSI"] = np.nan

            g = sns.FacetGrid(
                obs,
                col="Cell type",
                col_order=cell_type_order,
                row="external_donor_name",
                row_order=most_common_individuals,
                hue="Cell type",
                hue_order=cell_type_order,
                palette="tab10",
                sharex=True,
                sharey=True,
                height=1.5,
                aspect=1.0,
                margin_titles=True,
            )
            g.map_dataframe(
                sns.histplot,
                y="PSI",
                bins=np.linspace(0-eps, 1+eps, 11),
                stat="proportion",
            )
            for ax in g.axes.flat:
                if len(ax.patches) == 0:
                    ax.text(0.5, 0.5, "N/A")
            g.fig.subplots_adjust(wspace=0.25, hspace=0.25)
            g.set_titles(col_template="{col_name}", row_template="Donor:\n{row_name}")
            g.set_xlabels("Freq.")
            g.set_ylabels("PSI")
            g.set(xlim=(0, 1), ylim=(0, 1))
            sns.despine()
            plt.savefig(output[i], bbox_inches='tight')
            plt.close()


rule psi_plot_epithelial:
    input:
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
    output:
        expand("output/plots/epithelial/{intron}.pdf", intron=introns_to_plot_epithelial),
    run:
        adata = anndata.read_h5ad(input[0])

        adata.obs["tct"] = adata.obs.tissue.astype(str) + "." + sample_info.cell_type.astype(str)
        tct_order = ["bladder\nurothelial", "skin\nepidermal", "large intestine\ngoblet", "large intestine\nepithelial", "large intestine\nenterocyte", "mammary gland\nepithelial", "large intestine\nsecretory", "kidney\nepithelial"]
        adata.obs.tct = adata.obs.tct.replace('Bladder.bladder_urothelial_cell', "bladder\nurothelial")
        adata.obs.tct = adata.obs.tct.replace('Skin.epidermal_cell', "skin\nepidermal")
        adata.obs.tct = adata.obs.tct.replace('Large_Intestine.large_intestine_goblet_cell', "large intestine\ngoblet")
        adata.obs.tct = adata.obs.tct.replace('Large_Intestine.epithelial_cell_of_large_intestine', "large intestine\nepithelial")
        adata.obs.tct = adata.obs.tct.replace('Large_Intestine.enterocyte_of_epithelium_of_large_intestine', "large intestine\nenterocyte")
        adata.obs.tct = adata.obs.tct.replace('Mammary_Gland.luminal_epithelial_cell_of_mammary_gland', "mammary gland\nepithelial")
        adata.obs.tct = adata.obs.tct.replace('Large_Intestine.secretory_cell', "large intestine\nsecretory")
        adata.obs.tct = adata.obs.tct.replace('Kidney.kidney_collecting_duct_epithelial_cell', "kidney\nepithelial")

        adata = adata[adata.obs.tct.isin(tct_order)]
        print(adata.obs["mouse.id"].value_counts())

        X = group_normalize(adata.X.toarray(), adata.var.cluster.values, smooth=False)
        var = adata.var.copy()
        var["position"] = np.arange(len(var))
        var["id"] = var.chromosome.astype(str) + ":" + var.start.astype(str) + "-" + var.end.astype(str)
        var = var.set_index("id")
        for i, intron in enumerate(introns_to_plot_epithelial):
            obs = adata.obs.copy()
            obs["PSI"] = X[:, var.loc[intron].position].ravel()
            eps = 1e-4

            obs["external_donor_name"] = obs["mouse.id"]

            obs_all = obs.copy()
            obs_all.external_donor_name = "All"
            most_common_individuals = obs.external_donor_name.value_counts()[:6].index.tolist()
            obs = obs[obs.external_donor_name.isin(most_common_individuals)]
            obs = pd.concat([obs_all, obs], ignore_index=True)
            most_common_individuals = ["All"] + most_common_individuals

            for x in tct_order:
                for y in most_common_individuals:
                    mask = (obs.tct == x) & (obs.external_donor_name == y)
                    if obs.loc[mask, "PSI"].notna().sum() < 3:
                        obs.loc[mask, "PSI"] = np.nan

            g = sns.FacetGrid(
                obs,
                col="tct",
                col_order=tct_order,
                row="external_donor_name",
                row_order=most_common_individuals,
                #hue="Cell type",
                #hue_order=cell_type_order,
                palette="tab10",
                sharex=True,
                sharey=True,
                height=1.5,
                aspect=1.0,
                margin_titles=True,
            )
            g.map_dataframe(
                sns.histplot,
                y="PSI",
                bins=np.linspace(0-eps, 1+eps, 11),
                stat="proportion",
            )
            for ax in g.axes.flat:
                if len(ax.patches) == 0:
                    ax.text(0.5, 0.5, "N/A")
            g.fig.subplots_adjust(wspace=0.25, hspace=0.25)
            g.set_titles(col_template="{col_name}", row_template="Donor:\n{row_name}")
            g.set_xlabels("Freq.")
            g.set_ylabels("PSI")
            g.set(xlim=(0, 1), ylim=(0, 1))
            sns.despine()
            plt.savefig(output[i], bbox_inches='tight')
            plt.close()


rule exp_plot_marrow_b:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
    output:
        expand("output/plots/{gene}.svg", gene=genes_to_plot),
    run:
        adata = anndata.read_h5ad(input[0])
        adata.obs = adata.obs.replace("late pro-B cell", "(1) pro-B")
        adata.obs = adata.obs.replace("precursor B cell", "(2) pre-B")
        adata.obs = adata.obs.replace("immature B cell", "(3) immature B")
        adata.obs = adata.obs.replace("naive B cell", "(4) naive B")
        print(adata.shape)
        adata.obs["Cell type"] = adata.obs.cell_ontology_class
        adata = adata[adata.obs["Cell type"].isin(cell_type_order)]
        print(adata.shape)
        sc.pp.normalize_total(adata, target_sum=1e4)
        sc.pp.log1p(adata)
        for i, gene in enumerate(genes_to_plot):
            obs = adata.obs.copy()
            obs["log norm. expression"] = adata[:, gene].X.toarray().ravel()
            plt.figure(figsize=(3, 2))
            g = sns.violinplot(
                data=obs,
                x="Cell type",
                y="log norm. expression",
                order=cell_type_order,
                #hue="Cell type",
                #hue_order=cell_type_order,
                palette="tab10",
                cut=0.05,
                scale="width",
                width=0.75,
            )
            g.set(xlabel=None)
            sns.despine(bottom=True)
            plt.xticks(rotation=45)
            plt.savefig(output[i], bbox_inches='tight')
            plt.close()
            plt.clf()


rule diff_test_summary:
    input:
        'output/{anything}/expression.tsv',
        'output/{anything}/splicing.clusters.tsv',
    output:
        'output/{anything}/summary.tsv',
    run:
        df_exp = pd.read_csv(input[0], "\t")
        df_spl = pd.read_csv(input[1], "\t")
        n_genes_exp = ((df_exp.p_value_adj < config["fdr"]) & (df_exp.abs_lfc > config["min_abs_lfc"])).sum()
        n_genes_spl = len(df_spl[(df_spl.p_value_adj < config["fdr"]) & (df_spl.max_abs_delta_psi > config["min_abs_delta_psi"])].gene_id.unique())
        ratio = n_genes_spl / n_genes_exp
        intersection = len(list(set(df_exp.gene.unique()[:100]).intersection(set(df_spl.gene_id.unique()[:100]))))
        print(n_genes_exp, n_genes_spl, ratio, intersection)
        res = pd.DataFrame([[n_genes_exp, n_genes_spl, ratio, intersection]], columns=["n_genes_exp", "n_genes_spl", "ratio", "top100_intersection"])
        res.to_csv(output[0], "\t", index=False)


rule get_classification_score:
    input:
        #"output/dimensionality_reduction/marrow/gene-expression/pca_20/latent.txt",
        #"output/dimensionality_reduction/marrow/introns-transitive/pca_20/latent.txt",
        "output/dimensionality_reduction/all/gene-expression/pca_40/latent.txt",
        "output/dimensionality_reduction/all/introns-shared-acceptor/vae_hyperopt/latent.txt",
    output:
        "output/dimensionality_reduction/marrow/classification_score.tsv",
    run:
        tissue = "Marrow"
        cts = [ct for t, ct in tissue_cell_type_pairs if t==tissue]
        cell_idx = np.where((sample_info.tissue=="Marrow") & np.isin(sample_info.cell_type, cts))[0]
        latent_exp = np.loadtxt(input[0])[cell_idx]
        latent_spl = np.loadtxt(input[1])[cell_idx]
        obs = sample_info.iloc[cell_idx]
        groups = obs["mouse.id"].values

        clf = LogisticRegression(random_state=42, max_iter=10000)

        results = []
        for cell_type in cts:
            print(cell_type)
            labels = (obs.cell_type==cell_type)
            for latent, latent_name in zip([latent_exp, latent_spl], ["Expresion", "Splicing"]):
                scores = cross_val_score(clf, latent, labels, groups=groups, scoring="roc_auc", cv=StratifiedGroupKFold(n_splits=3, shuffle=True, random_state=42))
                for split, score in enumerate(scores):
                    results.append([cell_type, latent_name, split, score])
        results = pd.DataFrame(results, columns=['Cell type', 'Latent space', 'split', "ROC_AUC"])
        results.to_csv(output[0], "\t", index=False)


rule plot_classification_score:
    input:
        "output/dimensionality_reduction/marrow/classification_score.tsv",
    output:
        "output/dimensionality_reduction/marrow/classification_score.svg",
    run:
        df = pd.read_csv(input[0], "\t")
        df["Cell type"] = df["Cell type"].str.replace("_", " ")
        df["Cell type"] = df["Cell type"].replace("hematopoietic stem cell", "HSC")
        df["Cell type"] = df["Cell type"].replace("granulocytopoietic cell", "granulocytop.")
        plt.figure(figsize=(4,3))
        g = sns.barplot(x="Cell type", y="ROC_AUC", hue="Latent space", data=df, ci='sd', palette="Accent");
        g.set_xticklabels(g.get_xticklabels(), rotation=45,  horizontalalignment='right')
        g.set(xlabel="")
        sns.despine()
        #plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0., title="latent space");
        plt.tight_layout()
        plt.savefig(output[0], bbox_inches="tight")


rule get_marker_introns:
    input:
        'output/{anything}/splicing.clusters.significant.tsv',
        'output/{anything}/splicing.introns.tsv',
    output:
        'output/{anything}/marker_introns.tsv',
    run:
        groups = pd.read_csv(input[0], "\t")
        introns = pd.read_csv(input[1], "\t")
        introns["id"] = introns.chromosome.astype(str) + ":" + introns.start.astype(str) + "-" + introns.end.astype(str)
        introns = introns.set_index("id")
        groups = groups.sort_values("max_abs_delta_psi", ascending=False).head(30)
        #groups = groups[(groups.p_value_adj < config["fdr"]) & (groups.max_abs_delta_psi >= config["min_abs_delta_psi"])].sort_values("p_value", ascending=True).head(30)
        print(groups.p_value_adj)
        print(groups.max_abs_delta_psi)
        groups["marker_intron"] = groups.cluster.apply(lambda x: introns[introns.cluster==x].delta_psi.idxmax())
        groups.to_csv(output[0], "\t")


rule plot_marker_introns_global:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
        expand('output/differential_splicing/tissue_cell_type/{tct[0]}/{tct[1]}/marker_introns.tsv', tct=tissue_cell_type_pairs),
    output:
        'output/differential_splicing/tissue_cell_type/marker_introns.svg',
        #'output/differential_splicing/tissue_cell_type/marker_introns_genes.svg',
    run:
        adata_spl = anndata.read_h5ad(input[1])
        adata_spl.var["prev_id"] = np.arange(len(adata_spl.var))
        adata_spl.var["id"] = adata_spl.var.chromosome.astype(str) + ":" + adata_spl.var.start.astype(str) + "-" + adata_spl.var.end.astype(str)
        adata_spl.var = adata_spl.var.set_index("id")


        all_marker_introns = []
        i = 0
        for input_path in input[2:]:
            marker_introns = pd.read_csv(input_path, "\t").head(1)
            #print(input_path)
            #print(marker_introns)
            #raise Exception("debug")
            marker_introns["group"] = i
            all_marker_introns.append(marker_introns)
            i += 1
        marker_introns = pd.concat(all_marker_introns, ignore_index=True)
        print(marker_introns.group.value_counts())
        marker_introns = marker_introns.sample(frac=1, random_state=42)
        marker_introns = marker_introns.sort_values("max_abs_delta_psi", kind="mergesort", ascending=False)
        marker_introns = marker_introns.drop_duplicates(["gene_id"])
        marker_introns = marker_introns.sort_values(["group", "max_abs_delta_psi"], ascending=[True, False])
        print(marker_introns.group.value_counts())

        introns = marker_introns.marker_intron.values
        adata_spl.X = group_normalize(adata_spl.X.toarray(), adata_spl.var.cluster.values, smooth=False)
        adata_spl = adata_spl[:, introns]

        print(marker_introns.head(10))

        adata_spl.obs["tissue_cell_type"] = adata_spl.obs.tissue.astype(str) + "/" + adata_spl.obs.cell_type.astype(str)
        tissue_cell_types = [tissue + "/" + cell_type for tissue, cell_type in tissue_cell_type_pairs]
        adata_spl = adata_spl[adata_spl.obs["tissue_cell_type"].isin(tissue_cell_types)]
        for tissue_cell_type in tissue_cell_types:
            print(tissue_cell_type)
            for intron in introns:
                idx_cells = np.where(adata_spl.obs["tissue_cell_type"]==tissue_cell_type)[0]
                n_defined = (~np.isnan(adata_spl[idx_cells, intron].X)).sum()
                if n_defined < 10:
                    adata_spl[idx_cells, intron].X = np.nan

        gene_name = pd.read_csv("output/gene_name.txt", "\t", index_col=0)
        adata_spl.var = adata_spl.var.merge(gene_name, how="left", left_on="gene_id", right_index=True)
        adata_spl.var.loc[:, "id"] = adata_spl.var.gene_name.astype(str) + "_" + adata_spl.var.prev_id.astype(str)
        adata_spl.var = adata_spl.var.set_index("id")
        introns = adata_spl.var.index.values
        #pd.DataFrame(introns).to_csv("output/marker_introns.txt", index=False, header=False)


        df = adata_spl.obs.filter(items=["tissue_cell_type"])
        df[adata_spl.var.index.values] = adata_spl.X
        df = df.groupby("tissue_cell_type").mean().loc[tissue_cell_types]
        df.index.name = None
        width,height = 24, 24
        fig, ax = plt.subplots(figsize=(width,height))
        ax = sns.heatmap(
            df.T,
            cmap="coolwarm",
            center=0.5,
            vmin=0.0, vmax=1.0,
            ax=ax,
            square=True,
            cbar_kws={'label': 'Mean PSI', "location": "top", "shrink": 0.5},
            #cbar=False,
            xticklabels=1,
            yticklabels=1,
        )
        plt.xticks(rotation=90)
        plt.yticks(rotation=0)
        ax.get_figure().savefig(output[0], bbox_inches='tight')
        plt.close("all")

        #adata_exp = anndata.read_h5ad(input[0])
        #sc.pp.normalize_total(adata_exp, target_sum=1e4)
        #sc.pp.log1p(adata_exp)
        #adata_exp.obs["tissue_cell_type"] = adata_exp.obs.subclass_label.str.replace("_", " ").str.replace("slash", "/")
        #adata_exp = adata_exp[adata_exp.obs["tissue_cell_type"].isin(cell_type_order)]
        #gene_ids = marker_introns.gene_id.values
        #gene_names = marker_introns.gene_name.values
        #adata_exp = adata_exp[:, gene_ids]
        #adata_exp.var["gene_name"] = gene_names
        #adata_exp.var = adata_exp.var.set_index("gene_name")
        #adata_exp.X = adata_exp.X.toarray()

        #df = adata_exp.obs.filter(items=["Cell type"])
        #df[adata_exp.var.index.values] = adata_exp.X
        #df = df.groupby("Cell type").mean().loc[cell_type_order]
        #df.index.name = None
        #fig, ax = plt.subplots(figsize=(width,height))
        #ax = sns.heatmap(
        #    df.T,
        #    cmap="viridis",
        #    vmin=0.0, vmax=2.2,
        #    robust=True,
        #    ax=ax,
        #    square=True,
        #    #cbar_kws={'label': 'Mean expression', "location": "top", "shrink": 0.5},
        #    cbar=False,
        #    xticklabels=1,
        #    yticklabels=False,
        #)
        #plt.xticks(rotation=90)
        #plt.yticks(rotation=0)
        #ax.get_figure().savefig(output[1], bbox_inches='tight')
        #plt.close("all")


rule plot_marker_introns_tissue:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
        lambda wildcards: expand('output/differential_splicing/tissue_cell_type/{tct[0]}/{tct[1]}/marker_introns.tsv', tct=[[t, ct] for t, ct in tissue_cell_type_pairs if t == wildcards["tissue"]]),
    output:
        'output/differential_splicing/tissue_cell_type/{tissue}/marker_introns.svg',
        'output/differential_splicing/tissue_cell_type/{tissue}/marker_introns_genes.svg',
        'output/differential_splicing/tissue_cell_type/{tissue}/plot_marker_introns.tsv',
    threads: workflow.cores
    run:
        adata_spl = anndata.read_h5ad(input[1])
        adata_spl.var["prev_id"] = np.arange(len(adata_spl.var))
        adata_spl.var["id"] = adata_spl.var.chromosome.astype(str) + ":" + adata_spl.var.start.astype(str) + "-" + adata_spl.var.end.astype(str)
        adata_spl.var = adata_spl.var.set_index("id")

        all_marker_introns = []
        i = 0
        for input_path in input[2:]:
            if "pancreatic_acinar_cell" in input_path: continue
            marker_introns = pd.read_csv(input_path, "\t").head(12)
            #print(input_path)
            #print(marker_introns)
            #raise Exception("debug")
            marker_introns["group"] = i
            all_marker_introns.append(marker_introns)
            i += 1
        marker_introns = pd.concat(all_marker_introns, ignore_index=True)
        print(marker_introns.group.value_counts())
        marker_introns = marker_introns.sample(frac=1, random_state=42)
        marker_introns = marker_introns.sort_values("max_abs_delta_psi", kind="mergesort", ascending=False)
        marker_introns = marker_introns.drop_duplicates(["gene_id"])
        print(marker_introns.group.value_counts())
        N_MARKERS = 3
        assert((marker_introns.groupby("group").size() >= N_MARKERS).all())
        #marker_introns = marker_introns.drop_duplicates(["cluster"])
        marker_introns = marker_introns.sort_values(["group", "max_abs_delta_psi"], ascending=[True, False])
        idx = marker_introns.groupby("group").max_abs_delta_psi.nlargest(N_MARKERS).to_frame().reset_index().level_1.values
        marker_introns = marker_introns.loc[idx]
        marker_introns = marker_introns.sort_values(["group", "max_abs_delta_psi"], ascending=[True, False])
        print(marker_introns.group.value_counts())
        print(marker_introns.Annotated.value_counts())
        marker_introns.to_csv(output[2], "\t")

        introns = marker_introns.marker_intron.values
        adata_spl.X = group_normalize(adata_spl.X.toarray(), adata_spl.var.cluster.values, smooth=False)
        adata_spl = adata_spl[:, introns]

        cell_types = [cell_type for tissue, cell_type in tissue_cell_type_pairs if tissue==wildcards["tissue"] and cell_type != "pancreatic_acinar_cell"]
        adata_spl = adata_spl[adata_spl.obs.tissue==wildcards["tissue"]]
        adata_spl = adata_spl[adata_spl.obs["cell_type"].isin(cell_types)]
        for cell_type in cell_types:
            for intron in introns:
                idx_cells = np.where(adata_spl.obs["cell_type"]==cell_type)[0]
                n_defined = (~np.isnan(adata_spl[idx_cells, intron].X)).sum()
                if n_defined < 10:
                #if n_defined < 20:
                    adata_spl[idx_cells, intron].X = np.nan

        gene_name = pd.read_csv("output/gene_name.txt", "\t", index_col=0)
        adata_spl.var = adata_spl.var.merge(gene_name, how="left", left_on="gene_id", right_index=True)
        adata_spl.var["Annotated"] = marker_introns.Annotated.astype(bool).values
        adata_spl.var["Annotated_str"] = ""
        adata_spl.var.loc[~(adata_spl.var.Annotated), "Annotated_str"] = "*"
        adata_spl.var.loc[:, "id"] = adata_spl.var.Annotated_str + adata_spl.var.gene_name.astype(str) + "_" + adata_spl.var.prev_id.astype(str)
        adata_spl.var = adata_spl.var.set_index("id")
        introns = adata_spl.var.index.values
        #pd.DataFrame(introns).to_csv("output/marker_introns.txt", index=False, header=False)

        df = adata_spl.obs.filter(items=["cell_type"])
        df[adata_spl.var.index.values] = adata_spl.X
        df = df.groupby("cell_type").mean().loc[cell_types]
        df.index.name = None
        df.index = df.index.str.replace("_", " ")
        width,height = np.array([0.5, 1.5]) * 0.7 * len(cell_types)
        fig, ax = plt.subplots(figsize=(width,height))
        ax = sns.heatmap(
            df.T,
            cmap="coolwarm",
            center=0.5,
            vmin=0.0, vmax=1.0,
            ax=ax,
            square=True,
            #cbar_kws={'label': 'Mean PSI', "location": "top", "shrink": 0.2},
            cbar_kws={'label': 'Mean PSI', "location": "top", "aspect": 10},
            #cbar=False,
            xticklabels=1,
            yticklabels=1,
        )
        plt.xticks(rotation=90)
        plt.yticks(rotation=0)
        ax.get_figure().savefig(output[0], bbox_inches='tight')
        plt.close("all")

        adata_exp = anndata.read_h5ad(input[0])
        sc.pp.normalize_total(adata_exp, target_sum=1e4)
        sc.pp.log1p(adata_exp)
        adata_exp = adata_exp[adata_exp.obs.tissue==wildcards["tissue"]]
        adata_exp = adata_exp[adata_exp.obs["cell_type"].isin(cell_types)]
        gene_ids = marker_introns.gene_id.values
        gene_names = marker_introns.gene_name.values
        adata_exp = adata_exp[:, gene_ids]
        adata_exp.var["gene_name"] = gene_names
        adata_exp.var = adata_exp.var.set_index("gene_name")
        adata_exp.X = adata_exp.X.toarray()

        df = adata_exp.obs.filter(items=["cell_type"])
        df[adata_exp.var.index.values] = adata_exp.X
        df = df.groupby("cell_type").mean().loc[cell_types]
        df.index.name = None
        df.index = df.index.str.replace("_", " ")
        fig, ax = plt.subplots(figsize=(width,height))
        ax = sns.heatmap(
            df.T,
            cmap="viridis",
            robust=True,
            ax=ax,
            square=True,
            #cbar_kws={'label': 'Mean expression', "location": "top", "shrink": 0.2},
            cbar_kws={'label': 'Mean expr.', "location": "top", "aspect": 10},
            #cbar=False,
            xticklabels=1,
            yticklabels=False,
            #yticklabels=1,
        )

        cb = ax.collections[0].colorbar
        tick_locator = ticker.MaxNLocator(nbins=1)
        cb.locator = tick_locator
        cb.update_ticks()

        plt.xticks(rotation=90)
        plt.yticks(rotation=0)
        ax.get_figure().savefig(output[1], bbox_inches='tight')
        plt.close("all")


rule make_bam_paths_tct:
    input:
        expand("output/mapping/filtered_bams/{sample_id}.bam", sample_id=sample_ids)
    output:
        "output/bam_paths/tct/{tissue}/{cell_type}/paths.txt",
    run:
        s_ids = sample_ids[np.where((sample_info.tissue==wildcards["tissue"]) & (sample_info.cell_type==wildcards["cell_type"]))[0]]
        cwd = os.getcwd()
        pd.DataFrame(dict(sample_id=s_ids, bam_path=[f"{cwd}/output/mapping/filtered_bams/{sample_id}.bam" for sample_id in s_ids])).to_csv(output[0], "\t", index=False, header=False)


rule make_coverage_track_tct:
    input:
        "output/bam_paths/tct/{tissue}/{cell_type}/paths.txt",
    output:
        "output/coverage_track/tct/{tissue}/{cell_type}/coverage.bw",
    threads: workflow.cores // 4
    shell:
        "python -m scquint.quantification.run coverage_track/Snakefile --cores {threads} -d output/coverage_track/tct/{wildcards.tissue}/{wildcards.cell_type}/ --config bam_paths=../../../../bam_paths/tct/{wildcards.tissue}/{wildcards.cell_type}/paths.txt chrom_sizes_path={chrom_sizes_path}"


quantifications_all = ["gene-expression", "introns-transitive", "SE", "introns-shared-acceptor"]

rule clustering_patterns_comparison_quantitative:
    input:
        expand("output/dimensionality_reduction/mg/{quantification}/pca_20/latent.txt", quantification=quantifications_mg_individual_expanded),
        #expand("output/dimensionality_reduction/mg_3_38_F/{quantification}/pca_10/latent.txt", quantification=quantifications_mg_individual_expanded),
        #expand("output/dimensionality_reduction/all/{quantification}/pca_20/latent.txt", quantification=quantifications_all),
    output:
        "output/comparison/mg_3_38_F/classification_results.tsv",
    run:
        obs = sample_info[(sample_info.tissue=="Mammary_Gland")]
        #obs = sample_info[(sample_info.tissue=="Mammary_Gland") & (sample_info["mouse.id"]=="3_38_F")]
        obs = obs.sort_index()
        #obs = sample_info.copy()
        idx = np.where((obs.tissue=="Mammary_Gland") & (obs["mouse.id"]=="3_38_F") & (obs.cell_ontology_class!="endothelial cell"))[0]
        obs = obs.iloc[idx]
        obs["Cell type"] = obs.cell_ontology_class
        obs["Plate ID"] = obs.plate_id
        results = []
        #for name, input_path in zip(quantifications_mg_individual_expanded, input):
        for name, input_path in zip(quantifications_mg_individual_expanded + ["random"], input + ["random"]):
            print(name)
            #if name == "gene-expression":
            #    name = "Gene expression \n (featureCounts)"
            #if name == "kallisto":
            #    name = "Isoform proportions \n (kallisto)"
            #if name == "bins-nmf":
            #    name = "100-bp bin cov. proportions \n (ODEGR-NMF)"
            #if name == "introns-shared-acceptor":
            #    name = "Alt. intron proportions \n (scQuint)"
            if name != "random":
                latent = np.loadtxt(input_path)
                latent = latent[idx]

            for cell_type in obs.cell_ontology_class.unique():

                for seed in range(100):
                    if name == "random":
                        latent = np.random.normal(size=latent.shape)
                    latent_a = latent[obs.cell_ontology_class==cell_type]
                    latent_b = latent[obs.cell_ontology_class!=cell_type]
                    dist_cell_type = np.linalg.norm(latent_a.mean(axis=0) - latent_b.mean(axis=0))
                    latent_a = latent[(obs.cell_ontology_class==cell_type) & (obs.plate_id=="B002438")]
                    latent_b = latent[(obs.cell_ontology_class==cell_type) & (obs.plate_id!="B002438")]
                    dist_plate = np.linalg.norm(latent_a.mean(axis=0) - latent_b.mean(axis=0))
                    dist_logratio = np.log(dist_cell_type) - np.log(dist_plate)
                    #print(dist_cell_type, dist_plate, dist_logratio)

                    results.append((name, cell_type, "cell_type", seed, *calculate_classification_metrics(latent, obs.cell_ontology_class==cell_type, slice(None), seed), dist_logratio))
                    results.append((name, cell_type, "plate", seed, *calculate_classification_metrics(latent, obs.plate_id=="B002438", np.where(obs.cell_ontology_class==cell_type)[0], seed), dist_logratio))


        results = pd.DataFrame(results, columns=["Method", 'Cell type', 'Prediction', 'seed', 'accuracy', 'F1 score', "AUROC", "dist_logratio"])
        results.to_csv(output[0], "\t", index=False)


rule clustering_patterns_comparison_quantitative_plot:
    input:
        "output/comparison/mg_3_38_F/classification_results.tsv",
    output:
        directory("output/comparison/mg_3_38_F/classification_results_plots/"),
    run:
        df = pd.read_csv(input[0], "\t")
        df = df.replace("luminal epithelial cell of mammary gland", "luminal\nepithelial")
        df = df.replace("basal cell", "basal")
        df = df.replace("stromal cell", "stromal")
        df.Method.replace({
            "bins-nmf": "ODEGR-NMF",
            "exons": "DEXSeq",
            "introns-gene": "DESJ",
            "introns-transitive": "LeafCutter",
            "introns-shared-acceptor": "scQuint",
        }, inplace=True)

        os.makedirs(output[0])

        for metric in ["AUROC", "accuracy", "dist_logratio"]:
            g = sns.catplot(
                x="Method",
                y=metric,
                data=df,
                order=["random", "gene-expression", "kallisto", "LeafCutter", "SE", "SE-shared-donor", "SE-shared-acceptor", "scQuint"],
                #order=["random", "gene-expression", "LeafCutter", "SE", "scQuint"],
                ci='sd',
                palette="Accent",
                row="Cell type",
                col="Prediction" if metric != "dist_logratio" else None,
                #kind="bar",
                kind="point", join=False,
                margin_titles=True,
                height=3.5,
                aspect=1.0,
               )
            g.set_xticklabels(rotation=45, horizontalalignment="right")
            #g.set(ylim=(0.45, 1), yticks=[0.5, 0.6, 0.7, 0.8, 0.9, 1.0])
            g.set(xlabel="")
            #sns.despine()
            #plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0., title="latent space");
            plt.tight_layout()
            plt.savefig(os.path.join(output[0], f"{metric}.svg"), bbox_inches="tight")
            plt.close()


rule dimensionality_reduction_marrow_expression:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/marrow/gene-expression/pca_20/latent.txt",
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue=="Marrow")]
        sc.pp.filter_genes(adata, min_cells=30)
        sc.pp.normalize_total(adata, target_sum=1e4)
        sc.pp.log1p(adata)
        X = adata.X.toarray()
        latent = PCA(n_components=20).fit_transform(X)
        np.savetxt(output[0], latent)


rule dimensionality_reduction_marrow_pca:
    input:
        "output/quantification/{quantification}/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/marrow/{quantification}/pca_{K}/latent.txt",
    wildcard_constraints: quantification='.+(?<!gene-expression)'
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue=="Marrow")]
        adata = filter_min_cells_per_feature(adata, 30)
        if wildcards["quantification"] == "introns-transitive":
            adata = recluster(adata)
        latent = run_pca(adata, int(wildcards["K"]))
        np.savetxt(output[0], latent)


rule dimensionality_reduction_tissue_pca:
    input:
        "output/quantification/{quantification}/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/tissue/{tissue}/{quantification}/pca_{k}/latent.txt",
    wildcard_constraints: quantification='.+(?<!gene-expression)'
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue==wildcards["tissue"])]
        if wildcards["quantification"] != "bins-nmf":
            adata = filter_min_cells_per_feature(adata, 100)
        print(adata.shape)
        if wildcards["quantification"] == "introns-transitive":
            adata = recluster(adata)
        print(adata.shape)
        latent = run_pca(adata, int(wildcards["k"]))
        np.savetxt(output[0], latent)


rule dimensionality_reduction_tissue_vae:
    input:
        "output/quantification/{quantification}/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/tissue/{tissue}/{quantification}/vae_{transform}_{sample}/latent.txt",
    wildcard_constraints: quantification='.+(?<!gene-expression)'
    resources:
        gpu=1
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue==wildcards["tissue"])]
        if wildcards["quantification"] != "bins-nmf":
            adata = filter_min_cells_per_feature(adata, 100)
        if np.isin(wildcards["quantification"], ["introns-transitive", "introns-transitive-shuffled", "leafcutter", "introns-subset-lc", "introns-subset-lc-nonfiltered"]):
            adata = recluster(adata)
        print(wildcards["quantification"], adata.shape)
        input_transform = wildcards["transform"]
        if input_transform == "frequency-smoothed":
            feature_addition = group_normalize(adata.X.sum(axis=0), adata.var.cluster, smooth=False).A1.ravel()
            print(feature_addition.shape)
        else:
            feature_addition = None
        latent, model = run_vae(
            adata, n_epochs_kl_warmup=30, regularization_gaussian_std=10.0,
            use_cuda=True, input_transform=input_transform,
            feature_addition=feature_addition, sample=wildcards["sample"]=="True" or wildcards["sample"]=="true",
        )
        np.savetxt(output[0], latent)


rule dimensionality_reduction_tissue_expression:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/tissue/{tissue}/gene-expression/pca_{k}/latent.txt",
    run:
        adata = anndata.read_h5ad(input[0])
        adata = adata[(adata.obs.tissue==wildcards["tissue"])]
        sc.pp.filter_genes(adata, min_cells=100)
        sc.pp.normalize_total(adata, target_sum=1e4)
        sc.pp.log1p(adata)
        X = adata.X.toarray()
        latent = PCA(n_components=int(wildcards["k"])).fit_transform(X)
        np.savetxt(output[0], latent)


rule dimensionality_reduction_tissue_joint_pca:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
    output:
        "output/dimensionality_reduction/tissue/{tissue}/joint/pca_{k}/latent.txt",
    run:
        adata_exp = anndata.read_h5ad(input[0])
        adata_exp = adata_exp[(adata_exp.obs.tissue==wildcards["tissue"])]
        sc.pp.filter_genes(adata_exp, min_cells=30)
        sc.pp.normalize_total(adata_exp, target_sum=1e4)
        sc.pp.log1p(adata_exp)
        X_exp = adata_exp.X.toarray()

        adata_spl = anndata.read_h5ad(input[1])
        adata_spl = adata_spl[(adata_spl.obs.tissue==wildcards["tissue"])]
        adata_spl = filter_min_cells_per_feature(adata_spl, 30)
        X_spl = group_normalize(adata_spl.X.toarray(), adata_spl.var.cluster.values, smooth=True)
        intron_clusters = adata_spl.var.cluster.values
        all_intron_clusters = np.unique(intron_clusters)
        first_indices_dict = {}
        for i, c in enumerate(intron_clusters):
            if c not in first_indices_dict:
                first_indices_dict[c] = i
        first_indices = np.array([first_indices_dict[c] for c in all_intron_clusters])
        X_spl = np.delete(X_spl, first_indices, axis=1)

        X = np.hstack([X_exp, X_spl])
        print(X_exp.shape, X_spl.shape, X.shape)

        latent = PCA(n_components=int(wildcards["k"])).fit_transform(X)
        np.savetxt(output[0], latent)


rule leafcutter_extract:
    input:
        "output/mapping/filtered_bams/{sample_id}.bam",
    output:
        "output/quantification/leafcutter/juncs/{sample_id}.junc"
    shell:
        """
        samtools index {input} &&
        regtools junctions extract -s0 -a 8 -m 50 -M 500000 {input} -o {output}
        """


rule leafcutter_prepare_junc_files:
    input:
        expand("output/quantification/leafcutter/juncs/{sample_id}.junc", sample_id=sample_ids_leafcutter),
    output:
        "output/quantification/leafcutter/juncfiles.txt",
    run:
        for path in input:
            shell(f"echo {path} >> {output}")


rule leafcutter_cluster:
    input:
        "output/quantification/leafcutter/juncfiles.txt",
    output:
        "output/quantification/leafcutter/_perind_numers.counts.gz",
    shell:
        """
        ~/miniconda2/bin/python leafcutter/clustering/leafcutter_cluster_regtools.py --checkchrom -j {input} -m 50 -o output/quantification/leafcutter/ -l 500000
        """


rule leafcutter_make_adata:
    input:
        "output/quantification/leafcutter/_perind_numers.counts.gz",
    output:
        "output/quantification/leafcutter/adata.h5ad",
    run:
        df = pd.read_csv(input[0], " ", index_col=0).T
        print(df)
        df = df.loc[sample_ids_leafcutter]
        print(df)
        X = sp_sparse.csr_matrix(df.values)
        print(X.shape)
        obs = pd.DataFrame(index=df.index)
        print(obs)
        var = pd.DataFrame(index=df.columns)
        var["chromosome"] = var.index.str.split(":").str[0]
        var["start"] = var.index.str.split(":").str[1]
        var["end"] = var.index.str.split(":").str[2]
        var["cluster"] = var.index.str.split(":").str[3]
        print(var)
        adata = anndata.AnnData(X=X, obs=obs, var=var)
        print(adata.shape)
        adata = filter_min_cells_per_feature(adata, 30)
        print(adata.shape)
        adata = recluster(adata)
        print(adata.shape)
        adata.var["original_cluster"] = adata.var.cluster
        adata.write(output[0], compression="gzip")
#
#
###methods_compare_latent_all = ["gene-expression/pca_20", "introns-shared-acceptor/pca_100_20", "introns-shared-acceptor/vae_hyperopt"]
methods_compare_latent_all = ["gene-expression/pca_40", "introns-shared-acceptor/vae_hyperopt"]
###methods_compare_latent_all = ["gene-expression/pca_20", "introns-shared-acceptor/vae_frequency-smoothed_False_100", "introns-shared-acceptor/vae_hyperopt"]
###methods_compare_latent_all_names = ["Expression", "Splicing-PCA", "Splicing-VAE"]
methods_compare_latent_all_names = ["Expression", "Splicing"]
##methods_compare_latent_all_names = ["Expression", "Splicing-VAE", "Splicing-PCA", "Splicing-PCA-HV"]
##
rule compare_latent_all:
    input:
        expand(
            "output/dimensionality_reduction/all/{method}/latent.txt",
            method=methods_compare_latent_all,
        )
    output:
        directory("output/comparison/all"),
    run:
        obs = sample_info
        latents = [np.loadtxt(input_path) for input_path in input]
        latent_names = methods_compare_latent_all_names

        labels_list = [obs.cell_ontology_class, obs.sex, obs["mouse.id"], obs.plate_id, obs.tissue, obs.subtissue]
        label_names = ["Cell type", "Sex", "Individual", "Plate", "Tissue", "Subtissue"]
        projector = lambda: UMAP(min_dist=0.5, n_neighbors=15, random_state=42)
        projector_name = "UMAP_15"
        main_path = output[0]

        path = os.path.join(main_path, "all")
        plot_comparison(latents, latent_names, slice(None), labels_list, label_names, path, projector=projector, save=True)

        for tissue in obs.tissue.unique():
            path = os.path.join(main_path, "tissue", tissue, projector_name)
            plot_comparison(latents, latent_names, np.where(obs.tissue==tissue)[0], labels_list, label_names, path, projector=projector, save=True)

        path = os.path.join(main_path, "mg_individual_3_56_F", projector_name)
        plot_comparison(latents, latent_names, np.where((obs.tissue=="Mammary_Gland") & (obs["mouse.id"]=="3_56_F"))[0], labels_list, label_names, path, projector=projector, save=True)
        path = os.path.join(main_path, "mg_individual_3_57_F", projector_name)
        plot_comparison(latents, latent_names, np.where((obs.tissue=="Mammary_Gland") & (obs["mouse.id"]=="3_57_F"))[0], labels_list, label_names, path, projector=projector, save=True)
        path = os.path.join(main_path, "mg_individual_3_38_F", projector_name)
        plot_comparison(latents, latent_names, np.where((obs.tissue=="Mammary_Gland") & (obs["mouse.id"]=="3_38_F"))[0], labels_list, label_names, path, projector=projector, save=True)
        path = os.path.join(main_path, "mg_individual_3_39_F", projector_name)
        plot_comparison(latents, latent_names, np.where((obs.tissue=="Mammary_Gland") & (obs["mouse.id"]=="3_39_F"))[0], labels_list, label_names, path, projector=projector, save=True)

        path = os.path.join(main_path, "Marrow_B", projector_name)
        plot_comparison(latents, latent_names, np.where((obs.tissue=="Marrow") & (obs.cell_ontology_class.str.endswith("B cell")))[0], labels_list, label_names, path, projector=projector, save=True)


rule prepare_adata_for_cellxgene_exp_spl:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
        "output/gene_name.txt",
        "output/dimensionality_reduction/all/gene-expression/pca_40/latent.txt",
        "output/dimensionality_reduction/all/introns-shared-acceptor/vae_hyperopt/latent.txt",
    output:
        "output/adata_cellxgene_exp_spl.h5ad",
    run:
        gene_name = pd.read_csv(input[2], "\t", index_col=0)
        latent_exp = np.loadtxt(input[3])
        latent_spl = np.loadtxt(input[4])
        adata_exp = anndata.read_h5ad(input[0])
        adata_exp.var = adata_exp.var.merge(gene_name, how="left", left_index=True, right_index=True)
        adata_exp.var["gene_id"] = adata_exp.var.index.values
        adata_exp.var = adata_exp.var.set_index("gene_name")
        adata_exp.var_names_make_unique()
        adata_exp.obs.index = adata_exp.obs.index.astype(str)
        adata_spl = anndata.read_h5ad(input[1])
        adata_spl.var = adata_spl.var.merge(gene_name, how="left", left_on="gene_id", right_index=True)
        adata_spl.var["id"] = adata_spl.var.gene_name + "_" + adata_spl.var.chromosome.astype(str) + ":" + adata_spl.var.start.astype(str) + "-" + adata_spl.var.end.astype(str)
        adata_spl.var = adata_spl.var.set_index("id", drop=False)
        adata_spl.obs.index = adata_spl.obs.index.astype(str)
        adata_spl.var.index = adata_spl.var.index.astype(str)
        assert((adata_exp.obs.index == adata_spl.obs.index).all())
        sc.pp.normalize_total(adata_exp, target_sum=1e4)
        sc.pp.log1p(adata_exp)
        adata_exp.X = adata_exp.X.toarray()
        adata_spl.X = group_normalize(adata_spl.X.toarray(), adata_spl.var.cluster.values, smooth=False)

        X = np.hstack((adata_exp.X, adata_spl.X))
        print(adata_exp.shape, adata_spl.shape, X.shape)
        obs = adata_exp.obs
        var = pd.concat([adata_exp.var, adata_spl.var])
        adata = anndata.AnnData(X=X, obs=obs, var=var)
        adata.obsm["X_umap_expression"] = UMAP(min_dist=0.5, n_neighbors=15, random_state=42).fit_transform(latent_exp)
        adata.obsm["X_umap_splicing"] = UMAP(min_dist=0.5, n_neighbors=15, random_state=42).fit_transform(latent_spl)
        adata.write_h5ad(output[0], compression="gzip")


rule compare_latent_pca_vae:
    input:
        "output/dimensionality_reduction/all/gene-expression/pca_20/latent.txt",
        "output/dimensionality_reduction/all/introns-shared-acceptor/pca_100_20/latent.txt",
        "output/dimensionality_reduction/all/introns-shared-acceptor/vae_hyperopt/latent.txt",
    output:
        "output/comparison/pca_vae/{tissue}/cell_type.svg",
    run:
        cell_idx = np.where(sample_info.tissue==wildcards["tissue"])[0]
        dfs = []
        names = ["Expression", "Splicing (PCA)", "Splicing (VAE)"]
        for name, input_path in zip(names, input):
            df = sample_info.iloc[cell_idx].copy()
            latent = np.loadtxt(input_path)[cell_idx]
            proj = UMAP(min_dist=0.5, n_neighbors=15, random_state=42).fit_transform(latent)
            df["UMAP 1"] = proj[:, 0]
            df["UMAP 2"] = proj[:, 1]
            df["Quantification"] = name
            dfs.append(df)
        df = pd.concat(dfs)
        g = sns.relplot(
            data=df,
            x="UMAP 1",
            y="UMAP 2",
            hue="Cell type",
            col="Quantification",
            col_order=names,
            kind="scatter",
            facet_kws={'sharey': False, 'sharex': False},
            height=3,
            palette="tab10",
            edgecolor="none",
            s=4,
        )
        g.set_titles(col_template="{col_name}")
        g.fig.subplots_adjust(wspace=0.1)
        for ax in g.axes.flat:
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_ylabel("UMAP 2")
        sns.despine()
        plt.savefig(output[0], bbox_inches='tight')


rule compare_latent_tissue:
    input:
        #"output/dimensionality_reduction/tissue/{tissue}/gene-expression/pca_20/latent.txt",
        #"output/dimensionality_reduction/tissue/{tissue}/introns-shared-acceptor/pca_20/latent.txt",
        #"output/dimensionality_reduction/tissue/{tissue}/joint/pca_20/latent.txt",
        #"output/dimensionality_reduction/tissue/{tissue}/introns-shared-acceptor/vae_frequency-smoothed_False/latent.txt",
        expand(
            "output/dimensionality_reduction/all/{method}/latent.txt",
            method=methods_compare_latent_all,
        )
    output:
        expand("output/comparison/tissue/{{tissue}}/{label}.svg", label=labels)
    run:
        obs = sample_info.copy()
        idx = np.where((obs.tissue==wildcards["tissue"]))[0]
        obs = obs.iloc[idx]

        dfs = []
        names = ["Expression latent space", "Splicing latent space"]
        for name, input_path in zip(names, input):
            df = obs.copy()
            latent = np.loadtxt(input_path)[idx]

            proj = UMAP(min_dist=0.5, n_neighbors=15, random_state=42).fit_transform(latent)
            df["UMAP 1"] = proj[:, 0]
            df["UMAP 2"] = proj[:, 1]
            df["Quantification"] = name
            dfs.append(df)
        df = pd.concat(dfs)
        for i, label in enumerate(labels):
            g = sns.relplot(
                data=df, x="UMAP 1", y="UMAP 2",
                col="Quantification", hue=label,
                kind="scatter", facet_kws={'sharey': False, 'sharex': False},
                height=3, palette="tab20", edgecolor="none", s=4,
                hue_order=obs["Cell type"].value_counts().index.values if label=="Cell type" else None
               )
            g.set_titles(col_template="{col_name}")
            g.fig.subplots_adjust(wspace=0.1)
            for ax in g.axes.flat:
                ax.set_xticks([])
                ax.set_yticks([])
                ax.set_ylabel("UMAP 2")
            sns.despine()
            plt.savefig(output[i], bbox_inches='tight')


rule differential_test_marrow_b:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
    output:
        'output/differential_splicing/marrow_b/{cell_type}/expression.tsv',
        'output/differential_splicing/marrow_b/{cell_type}/splicing.clusters.tsv',
        'output/differential_splicing/marrow_b/{cell_type}/splicing.introns.tsv',
    threads: workflow.cores // 4
    run:
        adata_exp = anndata.read_h5ad(input[0])
        adata_exp.obs.index = adata_exp.obs.index.astype(str)
        adata_spl = anndata.read_h5ad(input[1])
        adata_spl.obs.index = adata_spl.obs.index.astype(str)
        adata_spl.var.index = adata_spl.var.index.astype(str)
        assert((adata_exp.obs.index == adata_spl.obs.index).all())
        obs = adata_exp.obs
        obs = obs.replace("late pro-B cell", "(1) pro-B")
        obs = obs.replace("precursor B cell", "(2) pre-B")
        obs = obs.replace("immature B cell", "(3) immature B")
        obs = obs.replace("naive B cell", "(4) naive B")
        mask = obs.cell_ontology_class.isin(cell_type_order)

        cell_idx_a = np.where(mask & (obs.cell_ontology_class==wildcards["cell_type"]))[0]
        cell_idx_b = np.where(mask & (obs.cell_ontology_class!=wildcards["cell_type"]))[0]

        MIN_FEATURES = 50
        diff_exp = run_differential_expression(adata_exp, cell_idx_a, cell_idx_b, MIN_FEATURES)
        diff_exp.to_csv(output[0], '\t', index=False)
        diff_spl_clusters, diff_spl_introns = run_differential_splicing(
            adata_spl,
            cell_idx_a,
            cell_idx_b,
            min_cells_per_cluster=MIN_FEATURES,
            min_total_cells_per_intron=MIN_FEATURES,
            n_jobs=threads,
            do_recluster=False,
        )
        diff_spl_clusters.to_csv(output[1], '\t')
        diff_spl_introns.to_csv(output[2], '\t')


rule plot_marker_introns_dist_marrow_b:
    input:
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
        expand('output/differential_splicing/marrow_b/{cell_type}/marker_introns.tsv', cell_type=cell_type_order),
    output:
        "output/differential_splicing/marrow_b/marker_introns_dist.svg",
    run:
        adata_spl = anndata.read_h5ad(input[0])
        adata_spl.var["prev_id"] = adata_spl.var.index.values
        adata_spl.var["id"] = adata_spl.var.chromosome.astype(str) + ":" + adata_spl.var.start.astype(str) + "-" + adata_spl.var.end.astype(str)
        adata_spl.var = adata_spl.var.set_index("id")

        adata_spl.obs = adata_spl.obs.replace("late pro-B cell", "(1) pro-B")
        adata_spl.obs = adata_spl.obs.replace("precursor B cell", "(2) pre-B")
        adata_spl.obs = adata_spl.obs.replace("immature B cell", "(3) immature B")
        adata_spl.obs = adata_spl.obs.replace("naive B cell", "(4) naive B")
        mask = adata_spl.obs.cell_ontology_class.isin(cell_type_order)
        #adata_spl = adata_spl[mask]
        adata_spl = adata_spl[mask].copy()

        all_marker_introns = []
        for input_path in input[1:]:
            marker_introns = pd.read_csv(input_path, "\t", header=None).values.astype(str).ravel().tolist()[:5]
            all_marker_introns += marker_introns
        print(all_marker_introns)
        print(len(all_marker_introns))
        all_marker_introns = pd.unique(all_marker_introns)
        print(len(all_marker_introns))
        adata_spl.X = group_normalize(adata_spl.X.toarray(), adata_spl.var.cluster.values, smooth=False)
        adata_spl = adata_spl[:, all_marker_introns]
        for cell_type in cell_type_order:
            print(cell_type)
            for intron in all_marker_introns:
                idx_cells = np.where(adata_spl.obs.cell_ontology_class==cell_type)[0]
                n_defined = (~np.isnan(adata_spl[idx_cells, intron].X)).sum()
                if n_defined < 10:
                    adata_spl[idx_cells, intron].X = np.nan

        gene_name = pd.read_csv("output/gene_name.txt", "\t", index_col=0)
        adata_spl.var = adata_spl.var.merge(gene_name, how="left", left_on="gene_id", right_index=True)
        #adata_spl.var.loc[:, "id"] = adata_spl.var.gene_name.astype(str) + "_" + adata_spl.var.index
        adata_spl.var.loc[:, "id"] = adata_spl.var.gene_name.astype(str) + "_" + adata_spl.var.prev_id.astype(str)
        adata_spl.var = adata_spl.var.set_index("id")

        marker_introns = adata_spl.var.index.values
        print("len(marker_introns): ", len(marker_introns))

        obs = adata_spl.obs
        obs.loc[:, marker_introns] = adata_spl.X
        obs["sample_id"] = obs.index.values
        obs["Cell type"] = obs.cell_ontology_class
        obs = obs.melt(id_vars=["sample_id", "Cell type"], value_vars=marker_introns, var_name="Intron", value_name="PSI")

        g = sns.FacetGrid(
            obs,
            col="Cell type",
            col_order=cell_type_order,
            hue="Cell type",
            row="Intron",
            palette="tab10",
            sharex=False,
            sharey=True,
            height=1.5,
            aspect=1,
            margin_titles=True,
        )
        eps = 1e-4
        #n_bins = 10
        n_bins = 5
        g.map_dataframe(
            sns.histplot,
            y="PSI",
            bins=np.linspace(0-eps, 1+eps, n_bins+1),
            stat="probability",
        )
        g = g.map(lambda y, **kw: plt.axhline(y.mean(), color=kw["color"], ls="--"), 'PSI')

        g.fig.subplots_adjust(wspace=0.1, hspace=0.12)
        g.set_titles(col_template="{col_name}", row_template="{row_name}")
        g.set_ylabels("PSI")
        g.set_xlabels("")
        g.set(xticks=[])
        g.set(xlim=(0, 1), ylim=(0, 1))
        #sns.despine(bottom=True)
        sns.despine()
        g.savefig(output[0], bbox_inches="tight")


rule plot_marker_introns_trajectory_marrow_b:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
        expand('output/differential_splicing/marrow_b/{cell_type}/marker_introns.tsv', cell_type=cell_type_order),
    output:
        "output/differential_splicing/marrow_b/marker_introns_trajectory_spl.svg",
    run:
        adata_exp = anndata.read_h5ad(input[0])

        adata_spl = anndata.read_h5ad(input[1])
        #adata_spl.var["prev_id"] = adata_spl.var.index.values
        adata_spl.var["prev_id"] = np.arange(len(adata_spl.var))
        adata_spl.var["id"] = adata_spl.var.chromosome.astype(str) + ":" + adata_spl.var.start.astype(str) + "-" + adata_spl.var.end.astype(str)
        adata_spl.var = adata_spl.var.set_index("id")

        adata_spl.obs = adata_spl.obs.replace("late pro-B cell", "(1)")
        adata_spl.obs = adata_spl.obs.replace("precursor B cell", "(2)")
        adata_spl.obs = adata_spl.obs.replace("immature B cell", "(3)")
        adata_spl.obs = adata_spl.obs.replace("naive B cell", "(4)")
        new_cell_type_order = ["(1)", "(2)", "(3)", "(4)"]
        mask = adata_spl.obs.cell_ontology_class.isin(new_cell_type_order)
        #adata_spl = adata_spl[mask]

        adata_exp = adata_exp[mask].copy()
        adata_spl = adata_spl[mask].copy()

        sc.pp.normalize_total(adata_exp, target_sum=1e4)
        sc.pp.log1p(adata_exp)
        #X = adata.X.toarray()

        all_marker_introns = []
        i = 0
        for input_path in input[2:]:
            n_examples = [3, 3, 4, 5]
            marker_introns = pd.read_csv(input_path, "\t").head(n_examples[i])
            marker_introns["group"] = i
            all_marker_introns.append(marker_introns)
            if i == 3:
                marker_introns = pd.read_csv(input_path, "\t")
                marker_introns = marker_introns[marker_introns.marker_intron=="chr6:99260420-99266347"]
                marker_introns["group"] = i
                all_marker_introns.append(marker_introns)
            i += 1
        marker_introns = pd.concat(all_marker_introns, ignore_index=True)
        print(marker_introns.group.value_counts())
        marker_introns = marker_introns.drop_duplicates(["gene_id"])
        marker_introns = marker_introns.iloc[[0, 3, 6, 9, 1, 4, 7, 10, 2, 5, 8, 11]]
        #marker_introns = marker_introns.iloc[[8, 0, 1, 5, 6, 2, 11, 7]]
        marker_introns = marker_introns.iloc[[8, 1, 6, 11, 0, 5, 2, 7]]
        print(marker_introns.group.value_counts())

        gene_ids = marker_introns.gene_id.values
        gene_names = marker_introns.gene_name.values
        adata_exp = adata_exp[:, gene_ids]
        adata_exp.X = adata_exp.X.toarray()

        introns = marker_introns.marker_intron.values
        adata_spl.X = group_normalize(adata_spl.X.toarray(), adata_spl.var.cluster.values, smooth=False)
        adata_spl = adata_spl[:, introns]

        for cell_type in new_cell_type_order:
            print(cell_type)
            for intron in introns:
                idx_cells = np.where(adata_spl.obs.cell_ontology_class==cell_type)[0]
                n_defined = (~np.isnan(adata_spl[idx_cells, intron].X)).sum()
                if n_defined < 3:
                    print("here")
                    raise Exception("debug")
                    adata_spl[idx_cells, intron].X = np.nan

        gene_name = pd.read_csv("output/gene_name.txt", "\t", index_col=0)
        adata_spl.var = adata_spl.var.merge(gene_name, how="left", left_on="gene_id", right_index=True)
        #adata_spl.var.loc[:, "id"] = adata_spl.var.gene_name.astype(str) + "_" + adata_spl.var.index
        adata_spl.var.loc[:, "id"] = adata_spl.var.gene_name.astype(str) + "_" + adata_spl.var.prev_id.astype(str)
        adata_spl.var = adata_spl.var.set_index("id")
        introns = adata_spl.var.index.values
        pd.DataFrame(introns).to_csv("output/marker_introns.txt", index=False, header=False)

        obs = adata_spl.obs
        #obs.loc[:, gene_names] = adata_exp.X
        obs.loc[:, introns] = adata_spl.X
        obs.loc[:, gene_names] = adata_exp.X.toarray()

        obs["sample_id"] = obs.index.values
        obs["Cell type"] = obs.cell_ontology_class
        obs_melt = obs.melt(id_vars=["sample_id", "Cell type"], value_vars=introns, var_name="Intron", value_name="PSI")

        color_exp = "C9"
        color_spl = "C4"

        obs_melt["foo"] = "foo"
        g = sns.catplot(
            data=obs_melt,
            col="Intron",
            col_order=introns,
            col_wrap=4,
            #row="Intron",
            #row_order=introns,
            palette=[color_spl],
            sharex=True,
            sharey=False,
            #height=2.0,
            height=1.7,
            aspect=1.5,
            #margin_titles=True,
            kind="point",
            ci=None,
            #kind="box",
            order=new_cell_type_order,
            hue="foo",
            x="Cell type",
            y="PSI",
            legend=False,
            marker="o",
        )
        g.fig.subplots_adjust(hspace=0.3, wspace=0.7)

        i = 0
        for idx, marker_intron in marker_introns.iterrows():
            ax = g.axes.flat[i]
            line_spl = ax.lines[0]
            i += 1
            ax2 = ax.twinx()
            mean_exp = obs.groupby("Cell type")[marker_intron.gene_name].mean().loc[new_cell_type_order].values
            mean_spl = line_spl.get_ydata()
            line_exp, = ax2.plot(ax.get_xticks(), mean_exp, color=color_exp, marker="o", linestyle="-", linewidth=3, markersize=8)
            line_spl, = ax.plot(ax.get_xticks(), mean_spl, color=color_spl, marker="o", linestyle="-", linewidth=3, markersize=8)
            ax.yaxis.set_tick_params(labelsize=8)
            ax2.yaxis.set_tick_params(labelsize=8)
            ax.set_ylabel("Mean PSI", color=color_spl, labelpad=0.5)
            ax2.set_ylabel("Mean expr.", color=color_exp, labelpad=0.5)

        #plt.legend([line_exp, line_spl], ['Expression', 'Splicing'], ncol=2, loc=9)
        g.fig.legend([line_exp, line_spl], ['Expression', 'Splicing'], ncol=2, loc=9, frameon=False)
        fig_width, fig_height = g.fig.get_size_inches()
        g.fig.set_size_inches(fig_width, fig_height + 1.2)
        g.set_titles(col_template="{col_name}", row_template="{row_name}")
        sns.despine(bottom=True, right=False)
        #g.savefig(output[0], bbox_inches="tight")
        g.savefig(output[0])


rule produce_latent_space_features_dendrogram:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
        "output/dimensionality_reduction/all/{features}/latent.txt",
        #"output/dimensionality_reduction/all/gene-expression/pca_40/latent.txt",
        #"output/dimensionality_reduction/all/introns-shared-acceptor/vae_hyperopt/latent.txt",
    output:
        "output/dendrogram/{features}/features.tsv",
        "output/dendrogram/{features}/colors.tsv",
        "output/dendrogram/{features}/color_reference.tsv",
        "output/dendrogram/{features}/legend.svg",
    run:
        adata = anndata.read_h5ad(input[0])
        adata.obs["functional_class"] = sample_info.functional_class.values
        #adata.obs["group"] = adata.obs.tissue.astype(str) + "/" + adata.obs.cell_type.astype(str)
        adata.obs["group"] = adata.obs.cell_ontology_class
        adata.obs["color_group"] = adata.obs.functional_class
        group_color_group = adata.obs.groupby("group").color_group.agg(pd.Series.mode)
        adata.obs = adata.obs.filter(items=["group"])
        latent = np.loadtxt(input[1])
        print("latent.shape: ", latent.shape)
        latent_names = [f"latent_{i}" for i in range(latent.shape[1])]
        print(latent_names)
        adata.obs[latent_names] = latent

        min_cells_per_group = 100
        group_counts = adata.obs.group.value_counts()
        groups = group_counts[group_counts >= min_cells_per_group].index.values
        print(adata.shape)
        adata = adata[adata.obs["group"].isin(groups)]

        print(adata.shape)
        #groupby = adata.obs.groupby("group").mean()
        groupby = adata.obs.groupby("group").median()
        print(groupby)

        group_color_group = group_color_group.loc[groupby.index.values].to_frame()
        print(group_color_group)
        ordering = ["immune", "epithelial", "endothelial", "stromal", "stem cell/progenitor", "Brain_Non-Myeloid", "Pancreas", "other"]
        group_color_group["color_group_code"] = group_color_group.color_group.apply(lambda x: ordering.index(str(x)))
        group_color_group.sort_values("color_group_code", inplace=True)
        print(group_color_group)

        groupby = groupby.loc[group_color_group.index.values]
        #raise Exception("debug")
        #group_color_group.color_group = group_color_group.color_group.astype("category")
        #print(group_color_group)
        #group_color_group["color_group_code"] = group_color_group.color_group.cat.codes

        import matplotlib
        from matplotlib import cm
        #cmap = cm.get_cmap("Dark2")
        #cmap = cm.get_cmap("tab20")
        initial_cmap = cm.get_cmap("tab10")
        def cmap(i):
            if i == 6:
                #new_i = 9
                new_i = 8
            else:
                new_i = i
            return initial_cmap(new_i)
        group_color_group["color_group_hex"] = group_color_group.color_group_code.apply(lambda i: matplotlib.colors.rgb2hex(cmap(i)))

        print(group_color_group)
        groupby.to_csv(output[0], "\t")
        group_color_group.to_csv(output[1], columns=["color_group_hex"], header=False, index=False)
        color_reference = group_color_group.groupby("color_group_code").color_group.first()
        color_reference.to_csv(output[2], sep="\t")
        print(color_reference)

        fig = plt.figure("Line plot")
        legendFig = plt.figure("Legend plot")
        ax = fig.add_subplot(111)

        lines = []
        for i, label in color_reference.iteritems():
            line, = ax.plot([0, 1], [0, 1], marker="o", linestyle="None", label=label, color=cmap(i))
            lines.append(line)

        legendFig.legend(lines, color_reference.values, loc='center')
        legendFig.savefig(output[3], bbox_inches="tight")


rule get_dendrogram_entanglement:
    input:
        "output/dendrogram/{features1}/features.tsv",
        "output/dendrogram/{features2}/features.tsv",
        "output/dendrogram/{features1}/colors.tsv",
    output:
        "output/dendrogram/{features1}_vs_{features2}/entanglement.txt",
        "output/dendrogram/{features1}_vs_{features2}/tanglegram.svg",
    shell:
        "Rscript compute_entanglement.R {input} {output}"


rule make_intron_ids:
    input:
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
        "output/gene_name.txt",
    output:
        "output/intron_id.tsv",
    run:
        adata_spl = anndata.read_h5ad(input[0])
        adata_spl.var["id"] = np.arange(len(adata_spl.var))
        adata_spl.var["coordinates"] = adata_spl.var.chromosome.astype(str) + ":" + adata_spl.var.start.astype(str) + "-" + adata_spl.var.end.astype(str)
        gene_name = pd.read_csv(input[1], "\t", index_col=0)
        adata_spl.var = adata_spl.var.merge(gene_name, how="left", left_on="gene_id", right_index=True)
        adata_spl.var.to_csv(output[0], "\t", index=False, columns=["gene_name", "id", "coordinates"])


rule make_latex_table_intron_coordinates_marrow_b:
    input:
        "output/intron_id.tsv",
        "output/marker_introns.txt",
    output:
        "output/intron_coordinates.txt",
        "output/intron_coordinates.tsv",
    run:
        id_coordinate_mapping = pd.read_csv(input[0], "\t", index_col=1)
        introns = pd.read_csv(input[1], header=None).values.astype(str).ravel().tolist()
        introns = [intron.replace('*', '') for intron in introns if "_" in intron]
        intron_id = [int(intron.split("_")[1]) for intron in introns]
        table = pd.DataFrame({"Intron id": introns})
        print(table)
        table["Intron coordinate"] = id_coordinate_mapping.loc[intron_id].coordinates.values
        print(table)
        table.sort_values("Intron id").to_latex(output[0], index=False)
        table.sort_values("Intron id").to_csv(output[1], sep="\t", index=False)


rule cluster_cells:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
    output:
        "output/cell_clusters_{n}.tsv",
    run:
        adata_spl.obs["tissue_cell_type"] = adata_spl.obs.tissue.astype(str) + "/" + adata_spl.obs.cell_type.astype(str)
        tissue_cell_types = [tissue + "/" + cell_type for tissue, cell_type in tissue_cell_type_pairs_including_singletons]
        adata = anndata.read_h5ad(input[0])
        adata = adata[adata.obs.tissue_cell_type.isin(tissue_cell_types)]
        sc.pp.normalize_total(adata, target_sum=1e4)
        sc.pp.log1p(adata)
        sc.pp.highly_variable_genes(adata)
        adata = adata[:, adata.var.highly_variable]
        sc.tl.pca(adata, n_comps=40)

        results = []

        for tissue_cell_type in tissue_cell_types:
            print(tissue_cell_type)
            adata_ct = adata[adata.obs.tissue_cell_type==tissue_cell_type]
            latent = adata_ct.obsm["X_pca"]
            print(latent.shape)
            n_clusters = len(latent) // int(wildcards["n"])
            model = SameSizeKMeansMinCostFlow(n_clusters, max_iters=10000)
            model.fit(latent)
            adata_ct.obs["cluster"] = model.labels_
            adata_ct.obs["cluster"] = cell_type + "-" + adata_ct.obs.cluster.astype(str)
            print((adata_ct.obs.cluster.value_counts() < 7).sum())
            results.append(adata_ct.obs.cluster)

        results = pd.concat(results)
        results.to_csv(output[0], "\t")


rule train_regression_clusters:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
        "input/diff_spl_gene_ids.txt",  # "input/target_genes.txt",
        "input/input_genes.txt",
        #"output/annotated_exons.bed",
        "input/AS_events/SE.most.introns.txt",
        "output/cell_clusters_{cluster_size}.tsv",
    output:
        "output/regression/clusters/{cluster_size}_{model}/metrics.tsv",
        #"output/regression/clusters/{cluster_size}_{model}/true.tsv",
        #"output/regression/clusters/{cluster_size}_{model}/predicted.tsv",
        "output/regression/clusters/{cluster_size}_{model}/coefficients.tsv",
    run:
        adata_exp = anndata.read_h5ad(input[0])
        adata_exp.var["feature_type"] = "Expression"
        adata_spl = anndata.read_h5ad(input[1])
        adata_spl.var["feature_type"] = "Splicing"
        adata_spl.var["id"] = adata_spl.var.gene_id.astype(str) + "_" + adata_spl.var.chromosome.astype(str) + ":" + adata_spl.var.start.astype(str) + "-" + adata_spl.var.end.astype(str)
        adata_spl.var["intron"] = adata_spl.var.chromosome.astype(str) + ":" + adata_spl.var.start.astype(str) + "-" + adata_spl.var.end.astype(str)
        adata_spl.var["length"] = adata_spl.var.end - adata_spl.var.start
        adata_spl.var = adata_spl.var.set_index("id")
        assert((adata_exp.obs.index == adata_spl.obs.index).all())

        clusters = pd.read_csv(input[5], "\t", index_col=0)
        adata_exp.obs = adata_exp.obs.merge(clusters, how="left", left_index=True, right_index=True)
        adata_spl.obs = adata_spl.obs.merge(clusters, how="left", left_index=True, right_index=True)

        sc.pp.normalize_total(adata_exp, target_sum=1e4)
        sc.pp.filter_genes(adata_exp, min_cells=100)
        sc.pp.log1p(adata_exp)

        diff_spl_genes = pd.read_csv(input[2], header=None).values.astype(str).ravel()
        input_genes = pd.read_csv(input[3]).values.astype(str).ravel()
        input_genes = [x for x in input_genes if x in adata_exp.var.index.values]
        target_genes = list(set(diff_spl_genes) - set(input_genes))

        adata_exp.var["gene_id"] = adata_exp.var.index.values
        adata_exp = adata_exp[:, input_genes]
        adata_spl_input = adata_spl[:, (adata_spl.var.gene_id.isin(input_genes)) & (adata_spl.var.gene_id.isin(diff_spl_genes))]
        adata_spl_output = adata_spl[:, adata_spl.var.gene_id.isin(target_genes)]


        adata_spl_input = filter_min_cells_per_feature(adata_spl_input, 100)
        adata_spl_input.X = group_normalize(adata_spl_input.X.toarray(), adata_spl_input.var.cluster.values, smooth=True)
        adata_spl_input.var["variance"] = np.nanvar(adata_spl_input.X, axis=0)
        adata_spl_input.var["priority"] = False
        priority_introns = [
            "ENSMUSG00000022332_chr15:68929658-68995052",
            "ENSMUSG00000008658_chr16:5763912-6173605",
        ]
        adata_spl_input.var.loc[priority_introns, "priority"] = True
        chosen_indices = adata_spl_input.var.sort_values(["priority", "variance", "annotated", "length"], ascending=[False, False, True, True]).groupby("cluster").head(1).sort_values("cluster").index.values
        adata_spl_input = adata_spl_input[:, chosen_indices]

        adata_spl_output = filter_min_cells_per_feature(adata_spl_output, 100)
        SE_introns = np.unique(pd.read_csv(input[4], header=None).values.ravel())
        adata_spl_output = adata_spl_output[:, adata_spl_output.var.intron.isin(SE_introns)]
        adata_spl_output = filter_singletons(adata_spl_output)
        cluster_counts = adata_spl_output.var.cluster.value_counts()
        correct_clusters = cluster_counts[cluster_counts==2].index.values
        adata_spl_output = adata_spl_output[:, adata_spl_output.var.cluster.isin(correct_clusters)]
        adata_spl_output = filter_singletons(adata_spl_output)
        adata_output_counts = adata_spl_output.copy()
        adata_spl_output.X = group_normalize(adata_spl_output.X.toarray(), adata_spl_output.var.cluster.values, smooth=False)

        chosen_indices = adata_spl_output.var.groupby("cluster").length.idxmin()
        adata_spl_output = adata_spl_output[:, chosen_indices]
        adata_output = adata_spl_output

        adata_output_counts = adata_output_counts[:, (adata_output_counts.var.sort_values(["cluster", "length"])).index.values]

        X = np.hstack((adata_exp.X.toarray(), adata_spl_input.X))
        obs = adata_exp.obs
        var = pd.concat([adata_exp.var, adata_spl_input.var])
        adata_input = anndata.AnnData(X=X, obs=obs, var=var)
        adata_input.obs = adata_input.obs.filter(items=["cluster"])
        adata_input.obs[adata_input.var.index.values] = adata_input.X
        groupby_input = adata_input.obs.groupby("cluster").mean().sample(frac=1, random_state=42)
        groupby_input = groupby_input.loc[:, groupby_input.std() > 0.05]  # TODO: tune this
        cell_types = groupby_input.index.str.split("-").str[0].values
        X = groupby_input.values  # might need intercept for some models
        X = StandardScaler().fit_transform(X)

        adata_output.obs = adata_output.obs.filter(items=["cluster"])
        adata_output.obs[adata_output.var.index.values] = adata_output.X
        groupby_output = adata_output.obs.groupby("cluster").mean().sample(frac=1, random_state=42)
        #mask_cols = (groupby_output.std() > 0.2) & (groupby_output.isna().mean() < 0.05)
        #top20 = pd.read_csv("top20.txt", header=None).values.ravel()
        #above05 = pd.read_csv("above_0.5.txt", header=None).values.ravel()
        output_chosen = pd.read_csv("top30.txt", header=None).values.ravel()
        mask_cols = np.isin(groupby_output.columns.values, output_chosen)
        groupby_output = groupby_output.loc[:, mask_cols]
        Y = groupby_output.values
        print("Y.shape: ", Y.shape)


        intron_id = pd.read_csv("output/intron_id.tsv", "\t", index_col=2)
        gene_name = pd.read_csv("input/gene_name.txt", "\t", index_col=0)
        replacement_columns = {
            col: gene_name.loc[col.split("_")[0]].gene_name + "_" + str(intron_id.loc[col.split("_")[1]].id) if "_" in col else gene_name.loc[col].gene_name
            for col in groupby_input.columns.values
        }
        groupby_input.rename(columns=replacement_columns, inplace=True)
        cell_types = [x.replace("_", " ").replace("slash", "/") for x in cell_types]
        input_ordered = pd.read_csv("output/regression/clusters/30_dir-multi-l1/coefficients_clustering_cols.txt", header=None).values.astype(str).ravel()
        output_ordered = pd.read_csv("output/regression/clusters/30_dir-multi-l1/coefficients_clustering_rows.txt", header=None).values.astype(str).ravel()

        groupby_input = (groupby_input-groupby_input.mean())/groupby_input.std()
        input_groupby_cell_type = groupby_input.groupby(cell_types).mean().loc[cell_type_order, input_ordered]

        #idx = np.concatenate([np.arange(0, 20), np.arange(83-20, 83)])
        #input_groupby_cell_type = input_groupby_cell_type.iloc[:, idx]

        #width,height = 7,2.5
        width,height = 14,2.5
        fig, ax = plt.subplots(figsize=(width,height))
        ax = sns.heatmap(
            input_groupby_cell_type,
            cmap="PiYG",
            center=0,
            robust=True,
            ax=ax,
            square=True,
            #cbar_kws={'label': 'Mean z-score', "location": "top", "shrink": 0.3},
            #cbar_kws={'label': 'Mean z-score', "location": "left", "shrink": 0.13, "anchor": (1.5, 0.25)},
            cbar_kws={"location": "bottom", "shrink": 0.2, "label": "Mean z-score", "anchor": (-0.3, 0.0), "aspect": 10},
            xticklabels=1,
            yticklabels=1,
        )
        label_fontsize = 13
        #cbar = ax.collections[0].colorbar
        #cbar.ax.tick_params(labelsize=label_fontsize)
        #cbar.ax.set_title('Mean z-score',fontsize=label_fontsize)
        ax.set_xlabel("Regulator", fontsize=label_fontsize)
        ax.set_ylabel("Cell type", fontsize=label_fontsize)
        ax.get_figure().savefig("output/input.svg", bbox_inches='tight')
        plt.close("all")

        replacement_columns = {
            col: gene_name.loc[col.split("_")[0]].gene_name + "_" + str(intron_id.loc[col.split("_")[1]].id) if "_" in col else gene_name.loc[col].gene_name
            for col in groupby_output.columns.values
        }
        groupby_output.rename(columns=replacement_columns, inplace=True)
        output_groupby_cell_type = groupby_output.groupby(cell_types).mean().loc[cell_type_order, output_ordered]
        width,height = 16, 7
        fig, ax = plt.subplots(figsize=(width,height))
        ax = sns.heatmap(
            output_groupby_cell_type.T,
            cmap="PuOr",
            center=0.5,
            #robust=True,
            ax=ax,
            square=True,
            cbar_kws={'label': 'Mean PSI', "location": "top", "shrink": 0.08, "anchor": (0.38, 1.1), "aspect": 10},
            xticklabels=1,
            yticklabels=1,
        )
        ax.xaxis.set_ticks_position('top')
        ax.xaxis.set_label_position('top')
        plt.xticks(rotation=90)
        #plt.xlabel("Predictor")
        #plt.ylabel("Target")
        ax.set_xlabel("Cell type", fontsize=label_fontsize)
        ax.set_ylabel("Target", fontsize=label_fontsize)
        ax.get_figure().savefig("output/output.svg", bbox_inches='tight')
        plt.close("all")
        raise Exception("debug")

        adata_output_counts.obs = adata_output_counts.obs.filter(items=["cluster"])
        adata_output_counts.obs[adata_output_counts.var.index.values] = adata_output_counts.X.toarray()
        groupby_output_counts = adata_output_counts.obs.groupby("cluster").sum().sample(frac=1, random_state=42)
        idx_cols = 2 * np.where(mask_cols)[0]
        idx_cols = sorted(np.concatenate([idx_cols, idx_cols+1]))
        groupby_output_counts = groupby_output_counts.iloc[:, idx_cols]
        Y_counts = groupby_output_counts.values

        assert((groupby_output.index==groupby_input.index).all())
        assert((groupby_output_counts.index==groupby_input.index).all())

        pred_groupby_output = groupby_output.copy()
        results = []
        coefficients = []

        for i, target in enumerate(groupby_output.columns):
            print(target)
            y = Y[:, i]
            y_counts = Y_counts[:, [2*i,2*i+1]]

            mask_nan = np.isnan(y)
            y_final = y[~mask_nan]
            X_final = X[~mask_nan]
            y_counts_final = y_counts[~mask_nan]
            cell_types_final = cell_types[~mask_nan]

            seed = 42

            if wildcards["model"] == "multi-l1":
                model = CVLassoMultinomialGLM(0.0, 400, 20)
                X_final = np.hstack((np.ones((len(X_final), 1), dtype=float), X_final))
                model.fit(X_final, y_counts_final, cell_types_final, device="cpu")
                print("L1 penalty: ", model.l1_penalty)
                y_pred = None
                r2 = 1.0
                coeff = model.model.A.cpu().detach().numpy()[1:].ravel()
            if wildcards["model"] == "dir-multi-l1":
                #model = CVLassoDirichletMultinomialGLM(0.0, 100, 20)
                model = CVLassoDirichletMultinomialGLM(10.0, 70, 30)
                X_final = np.hstack((np.ones((len(X_final), 1), dtype=float), X_final))
                model.fit(X_final, y_counts_final, cell_types_final, device="cpu")
                print("L1 penalty: ", model.l1_penalty)
                y_pred = None
                r2 = 1.0
                coeff = model.model.A.cpu().detach().numpy()[1:].ravel()
                print("log_alpha: ", model.model.log_alpha.cpu().detach().numpy())
            if wildcards["model"] == "ridge":
                alphas = np.logspace(-1, 4, 100)
                model = RidgeCV(alphas=alphas)
                model.fit(X_final, y_final)
                r2 = model.score(X_final, y_final)
                y_pred = model.predict(X)
                coeff = model.coef_
            elif wildcards["model"] == "lasso":
                alphas = np.logspace(-5, 2, 200)
                model = LassoCV(random_state=seed, max_iter=10000, alphas=alphas, cv=10)
                model.fit(X_final, y_final)
                r2 = model.score(X_final, y_final)
                y_pred = model.predict(X)
                coeff = model.coef_
            elif wildcards["model"] == "elasticnet":
                l1_ratio = [.1, .5, .7, .9, .95, .99, 1]
                model = ElasticNetCV(random_state=seed, max_iter=10000, cv=10, l1_ratio=l1_ratio)
                model.fit(X_final, y_final)
                r2 = model.score(X_final, y_final)
                y_pred = model.predict(X)
                coeff = model.coef_
            if wildcards["model"] in ["ridge", "lasso", "elasticnet"]:
                print(model.alpha_)
            #print(r2)
            #pred_groupby_output[target] = y_pred
            results.append([target, r2])
            coefficients.append(coeff)

        results = pd.DataFrame(results, columns=["target", "r2",])
        results.to_csv(output[0], "\t", index=False)
        #groupby_output.to_csv(output[1], "\t")
        #pred_groupby_output.to_csv(output[2], "\t")
        coefficients = pd.DataFrame(coefficients, index=groupby_output.columns, columns=groupby_input.columns)
        #coefficients.to_csv(output[3], "\t")
        coefficients.to_csv(output[1], "\t")


rule plot_regression_clusters_coefficients:
    input:
        "output/regression/clusters/{anything}/coefficients.tsv",
        "output/regression/clusters/{anything}/metrics.tsv",
        "input/gene_name.txt",
        "output/intron_id.tsv",
    output:
        "output/regression/clusters/{anything}/coefficients_alphabetical.svg",
        "output/regression/clusters/{anything}/coefficients_clustering.svg",
        "output/regression/clusters/{anything}/coefficients_clustering_rows.txt",
        "output/regression/clusters/{anything}/coefficients_clustering_cols.txt",
    run:
        coefficients = pd.read_csv(input[0], "\t", index_col=0)
        metrics = pd.read_csv(input[1], "\t", index_col=0)
        assert((coefficients.index==metrics.index).all())
        print(coefficients.shape)

        intron_id = pd.read_csv(input[3], "\t", index_col=2)

        #coefficients = coefficients[metrics.r2 >= 0.5]
        #top_idx = np.argsort(metrics.r2.values)[::-1][:20]
        #print(metrics.iloc[top_idx])
        #coefficients = coefficients.iloc[top_idx]
        chosen = pd.read_csv("top30.txt", header=None).values.ravel()
        coefficients = coefficients.loc[chosen]

        print(coefficients.shape)
        gene_name = pd.read_csv(input[2], "\t", index_col=0)
        replacement_columns = {
            col: gene_name.loc[col.split("_")[0]].gene_name + "_" + str(intron_id.loc[col.split("_")[1]].id) if "_" in col else gene_name.loc[col].gene_name
            for col in coefficients.columns.values
        }
        coefficients.rename(columns=replacement_columns, inplace=True)
        replacement_index = {
            idx: gene_name.loc[idx.split("_")[0]].gene_name + "_" + str(intron_id.loc[idx.split("_")[1]].id)
            for idx in coefficients.index.values
        }
        coefficients.rename(index=replacement_index, inplace=True)
        coefficients = coefficients.reindex(sorted(coefficients.columns), axis=1)
        coefficients.sort_index(inplace=True)
        print(coefficients)

        #sns.set(font_scale=1.2)
        width,height = 16,7
        fig, ax = plt.subplots(figsize=(width,height))
        ax = sns.heatmap(
            coefficients,
            cmap="bwr",
            center=0,
            robust=True,
            ax=ax,
            square=True,
            cbar_kws={'label': 'Regression coefficient', "location": "top", "shrink": 0.3},
            xticklabels=1,
            yticklabels=1,
        )
        ax.set_xlabel("Regulator")
        ax.set_ylabel("Target")
        ax.get_figure().savefig(output[0], bbox_inches='tight')
        plt.close("all")

        #coefficients_abs = coefficients.copy()
        #coefficients_abs.loc[:, coefficients_abs.columns.str.contains("_")] = coefficients_abs.loc[:, coefficients_abs.columns.str.contains("_")].abs()

        #g = sns.clustermap(coefficients)
        #new_row_idx = g.dendrogram_row.reordered_ind

        #g = sns.clustermap(coefficients_abs)
        #g = sns.clustermap(coefficients)
        #new_col_idx = g.dendrogram_col.reordered_ind

        new_row_idx = leaves_list(linkage(coefficients, optimal_ordering=True, method="ward"))
        new_col_idx = leaves_list(linkage(coefficients.T, optimal_ordering=True, method="ward"))
        coefficients_reordered = coefficients.iloc[new_row_idx, new_col_idx]

        idx = np.concatenate([np.arange(0, 20), np.arange(83-20, 83)])

        #fig, ax = plt.subplots(figsize=(12,6))
        fig, ax = plt.subplots(figsize=(24,6))
        ax = sns.heatmap(
            coefficients_reordered,
            #coefficients_reordered.iloc[:, idx],
            cmap="bwr",
            center=0,
            robust=True,
            ax=ax,
            square=True,
            cbar_kws={'label': 'Regression coefficient', "location": "top", "shrink": 0.15, "aspect": 10},
            #xticklabels=1,
            #yticklabels=1,
            xticklabels=False,
            yticklabels=False,
        )
        #ax.set_xlabel("Regulator")
        #ax.set_ylabel("Target")
        ax.get_figure().savefig(output[1], bbox_inches='tight')
        plt.close()
        pd.DataFrame(coefficients_reordered.index.values).to_csv(output[2], index=False, header=False)
        pd.DataFrame(coefficients_reordered.columns.values).to_csv(output[3], index=False, header=False)


rule make_latex_table_intron_coordinates:
    input:
        "output/intron_id.tsv",
        "output/{anything}/coefficients_clustering_cols.txt",
        "output/{anything}/coefficients_clustering_rows.txt",
    output:
        "output/{anything}/intron_coordinates.txt",
        "output/{anything}/intron_coordinates.tsv",
    run:
        id_coordinate_mapping = pd.read_csv(input[0], "\t", index_col=1)
        introns = pd.read_csv(input[1], header=None).values.astype(str).ravel().tolist() + pd.read_csv(input[2], header=None).values.astype(str).ravel().tolist()
        introns = [intron for intron in introns if "_" in intron]
        intron_id = [int(intron.split("_")[1]) for intron in introns]
        table = pd.DataFrame({"Intron id": introns})
        print(table)
        table["Intron coordinate"] = id_coordinate_mapping.loc[intron_id].coordinates.values
        print(table)
        table.sort_values("Intron id").to_latex(output[0], index=False)
        table.sort_values("Intron id").to_csv(output[1], sep="\t", index=False)


rule compare_latent_functional_class:
    input:
        "output/dimensionality_reduction/all/gene-expression/pca_40/latent.txt",
        "output/dimensionality_reduction/all/introns-shared-acceptor/vae_hyperopt/latent.txt",
    output:
        "output/comparison/functional_class/{functional_class}/tissue_cell_type.svg",
    run:
        cell_idx = np.where(sample_info.functional_class.str.contains(wildcards["functional_class"]))[0]
        dfs = []
        names = ["Expression", "Splicing (VAE)"]
        for name, input_path in zip(names, input):
            df = sample_info.iloc[cell_idx].copy()
            df["Tissue-Cell type"] = df.tissue.astype(str) + "." + df.cell_ontology_class.astype(str)
            latent = np.loadtxt(input_path)[cell_idx]
            proj = UMAP(min_dist=0.5, n_neighbors=15, random_state=42).fit_transform(latent)
            df["UMAP 1"] = proj[:, 0]
            df["UMAP 2"] = proj[:, 1]
            df["Quantification"] = name
            dfs.append(df)
        df = pd.concat(dfs)
        g = sns.relplot(
            data=df,
            x="UMAP 1",
            y="UMAP 2",
            hue="Tissue-Cell type",
            col="Quantification",
            col_order=names,
            kind="scatter",
            facet_kws={'sharey': False, 'sharex': False},
            height=3,
            palette="tab10",
            edgecolor="none",
            s=4,
        )
        g.set_titles(col_template="{col_name}")
        g.fig.subplots_adjust(wspace=0.1)
        for ax in g.axes.flat:
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_ylabel("UMAP 2")
        sns.despine()
        plt.savefig(output[0], bbox_inches='tight')


rule differential_test_functional_class:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
        "output/gene_name.txt",
    output:
        'output/differential_splicing/functional_class/{functional_class}/{tct}/expression.tsv',
        'output/differential_splicing/functional_class/{functional_class}/{tct}/splicing.clusters.tsv',
        'output/differential_splicing/functional_class/{functional_class}/{tct}/splicing.introns.tsv',
    threads: workflow.cores // 4
    run:
        adata_exp = anndata.read_h5ad(input[0])
        adata_exp.obs.index = adata_exp.obs.index.astype(str)
        adata_spl = anndata.read_h5ad(input[1])
        adata_spl.obs.index = adata_spl.obs.index.astype(str)
        adata_spl.var.index = adata_spl.var.index.astype(str)
        gene_name = pd.read_csv(input[2], sep="\t", index_col=0)
        print(adata_spl.var.head())
        adata_spl.var = adata_spl.var.merge(gene_name, how="left", left_on="gene_id", right_index=True)
        print(adata_spl.var.head())
        assert((adata_exp.obs.index == adata_spl.obs.index).all())
        assert((adata_exp.obs.index == sample_info.index).all())
        obs = sample_info
        cell_idx_a = np.where((obs.tct==wildcards["tct"]))[0]
        cell_idx_b = np.where((obs.tct!=wildcards["tct"]) &
                              (obs.tct.isin(tcts_functional_class[wildcards["functional_class"]])))[0]
        MIN_FEATURES = 50
        diff_exp = run_differential_expression(adata_exp, cell_idx_a, cell_idx_b, MIN_FEATURES)
        diff_exp.to_csv(output[0], '\t', index=False)
        diff_spl_clusters, diff_spl_introns = run_differential_splicing(
            adata_spl,
            cell_idx_a,
            cell_idx_b,
            min_cells_per_cluster=MIN_FEATURES,
            min_total_cells_per_intron=MIN_FEATURES,
            n_jobs=threads,
            do_recluster=False,
        )
        diff_spl_clusters.to_csv(output[1], '\t')
        diff_spl_introns.to_csv(output[2], '\t')


rule plot_marker_introns_functional_class:
    input:
        "output/quantification/gene-expression/adata_annotated.h5ad",
        "output/quantification/introns-shared-acceptor/adata_annotated.h5ad",
        lambda wildcards: expand('output/differential_splicing/functional_class/{{functional_class}}/{tct}/marker_introns.tsv', tct=tcts_functional_class[wildcards["functional_class"]]),
    output:
        'output/differential_splicing/functional_class/{functional_class}/marker_introns.svg',
        'output/differential_splicing/functional_class/{functional_class}/marker_introns_genes.svg',
        'output/differential_splicing/functional_class/{functional_class}/plot_marker_introns.tsv',
    threads: workflow.cores
    run:
        adata_spl = anndata.read_h5ad(input[1])
        adata_spl.var["prev_id"] = np.arange(len(adata_spl.var))
        adata_spl.var["id"] = adata_spl.var.chromosome.astype(str) + ":" + adata_spl.var.start.astype(str) + "-" + adata_spl.var.end.astype(str)
        adata_spl.var = adata_spl.var.set_index("id")
        assert((adata_spl.obs.index == sample_info.index).all())
        adata_spl.obs["tct"] = sample_info.tct.values

        all_marker_introns = []
        i = 0
        for input_path in input[2:]:
            marker_introns = pd.read_csv(input_path, "\t").head(10)
            #print(input_path)
            #print(marker_introns)
            #raise Exception("debug")
            marker_introns["group"] = i
            all_marker_introns.append(marker_introns)
            i += 1
        marker_introns = pd.concat(all_marker_introns, ignore_index=True)
        print(marker_introns.group.value_counts())
        marker_introns = marker_introns.sample(frac=1, random_state=42)
        marker_introns = marker_introns.sort_values("max_abs_delta_psi", kind="mergesort", ascending=False)
        marker_introns = marker_introns.drop_duplicates(["gene_id"])
        print(marker_introns.group.value_counts())
        N_MARKERS = 3
        assert((marker_introns.groupby("group").size() >= N_MARKERS).all())
        #marker_introns = marker_introns.drop_duplicates(["cluster"])
        marker_introns = marker_introns.sort_values(["group", "max_abs_delta_psi"], ascending=[True, False])
        idx = marker_introns.groupby("group").max_abs_delta_psi.nlargest(N_MARKERS).to_frame().reset_index().level_1.values
        marker_introns = marker_introns.loc[idx]
        marker_introns = marker_introns.sort_values(["group", "max_abs_delta_psi"], ascending=[True, False])
        print(marker_introns.group.value_counts())
        print(marker_introns.Annotated.value_counts())
        marker_introns.to_csv(output[2], "\t")

        introns = marker_introns.marker_intron.values
        adata_spl.X = group_normalize(adata_spl.X.toarray(), adata_spl.var.cluster.values, smooth=False)
        adata_spl = adata_spl[:, introns]

        tcts = tcts_functional_class[wildcards["functional_class"]]
        adata_spl = adata_spl[adata_spl.obs.tct.isin(tcts)]
        for tct in tcts:
            for intron in introns:
                idx_cells = np.where(adata_spl.obs.tct==tct)[0]
                n_defined = (~np.isnan(adata_spl[idx_cells, intron].X)).sum()
                if n_defined < 10:
                    adata_spl[idx_cells, intron].X = np.nan

        gene_name = pd.read_csv("output/gene_name.txt", "\t", index_col=0)
        adata_spl.var = adata_spl.var.merge(gene_name, how="left", left_on="gene_id", right_index=True)
        adata_spl.var["Annotated"] = marker_introns.Annotated.astype(bool).values
        adata_spl.var["Annotated_str"] = ""
        adata_spl.var.loc[~(adata_spl.var.Annotated), "Annotated_str"] = "*"
        adata_spl.var.loc[:, "id"] = adata_spl.var.Annotated_str + adata_spl.var.gene_name.astype(str) + "_" + adata_spl.var.prev_id.astype(str)
        adata_spl.var = adata_spl.var.set_index("id")
        introns = adata_spl.var.index.values
        #pd.DataFrame(introns).to_csv("output/marker_introns.txt", index=False, header=False)

        df = adata_spl.obs.filter(items=["tct"])
        df[adata_spl.var.index.values] = adata_spl.X
        df = df.groupby("tct").mean().loc[tcts]
        df.index.name = None
        df.index = df.index.str.replace("_", " ")
        width,height = np.array([0.5, 1.5]) * 0.7 * len(tcts)
        fig, ax = plt.subplots(figsize=(width,height))
        ax = sns.heatmap(
            df.T,
            cmap="coolwarm",
            center=0.5,
            vmin=0.0, vmax=1.0,
            ax=ax,
            square=True,
            #cbar_kws={'label': 'Mean PSI', "location": "top", "shrink": 0.2},
            cbar_kws={'label': 'Mean PSI', "location": "top", "aspect": 10},
            #cbar=False,
            xticklabels=1,
            yticklabels=1,
        )
        plt.xticks(rotation=90)
        plt.yticks(rotation=0)
        ax.get_figure().savefig(output[0], bbox_inches='tight')
        plt.close("all")

        adata_exp = anndata.read_h5ad(input[0])
        assert((adata_exp.obs.index == sample_info.index).all())
        adata_exp.obs["tct"] = sample_info.tct.values
        sc.pp.normalize_total(adata_exp, target_sum=1e4)
        sc.pp.log1p(adata_exp)
        adata_exp = adata_exp[adata_exp.obs.tct.isin(tcts)]
        gene_ids = marker_introns.gene_id.values
        gene_names = marker_introns.gene_name.values
        adata_exp = adata_exp[:, gene_ids]
        adata_exp.var["gene_name"] = gene_names
        adata_exp.var = adata_exp.var.set_index("gene_name")
        adata_exp.X = adata_exp.X.toarray()

        df = adata_exp.obs.filter(items=["tct"])
        df[adata_exp.var.index.values] = adata_exp.X
        df = df.groupby("tct").mean().loc[tcts]
        df.index.name = None
        df.index = df.index.str.replace("_", " ")
        fig, ax = plt.subplots(figsize=(width,height))
        ax = sns.heatmap(
            df.T,
            cmap="viridis",
            robust=True,
            ax=ax,
            square=True,
            #cbar_kws={'label': 'Mean expression', "location": "top", "shrink": 0.2},
            cbar_kws={'label': 'Mean expr.', "location": "top", "aspect": 10},
            #cbar=False,
            xticklabels=1,
            yticklabels=False,
            #yticklabels=1,
        )

        cb = ax.collections[0].colorbar
        tick_locator = ticker.MaxNLocator(nbins=1)
        cb.locator = tick_locator
        cb.update_ticks()

        plt.xticks(rotation=90)
        plt.yticks(rotation=0)
        ax.get_figure().savefig(output[1], bbox_inches='tight')
        plt.close("all")


rule make_latex_table_intron_coordinates_anything:
    input:
        "output/intron_id.tsv",
        "output/{anything}/plot_marker_introns.tsv",
    output:
        "output/{anything}/plot_marker_intron_coordinates.txt",
        "output/{anything}/plot_marker_intron_coordinates.tsv",
    run:
        coordinate_id_mapping = pd.read_csv(input[0], "\t", index_col=2)
        introns = pd.read_csv(input[1], sep="\t")
        coordinate_id_mapping = coordinate_id_mapping.loc[introns.marker_intron.values]
        print(coordinate_id_mapping)
        coordinate_id_mapping["Intron id"] = coordinate_id_mapping.gene_name.astype(str) + "_" + coordinate_id_mapping["id"].astype(str)
        coordinate_id_mapping["Intron coordinate"] = coordinate_id_mapping.index.values
        table = coordinate_id_mapping.filter(items=["Intron id", "Intron coordinate"])
        table.sort_values("Intron id").to_latex(output[0], index=False)
        table.sort_values("Intron id").to_csv(output[1], sep="\t", index=False)


rule plot_tissues_per_donor:
    output:
        "output/plots/tissues_per_donor.pdf",
    run:
        tissues = sorted(sample_info.tissue.unique())
        print(tissues)
        #res = sample_info[~sample_info["mouse.id"].str.contains("/")].groupby("mouse.id").tissue.agg(lambda ts: (ts.value_counts().reindex(tissues).fillna(0) > 0).values.tolist())
        #df = sample_info[~sample_info["mouse.id"].str.contains("/")].groupby("mouse.id").tissue.agg(lambda ts: ts.value_counts().reindex(tissues).fillna(0).astype(int).values.tolist()).to_frame()
        df = pd.pivot_table(sample_info[~sample_info["mouse.id"].str.contains("/")], index="tissue", columns="mouse.id", values="cell_ontology_class", aggfunc=len).fillna(0).astype(int)
        print(df)
        sns.heatmap(df, cbar_kws={'label': 'Number of cells'}, square=True)
        plt.savefig(output[0], bbox_inches='tight')


rule plot_marker_introns_trajectory_marrow_b:
    input:
        expand('output/differential_splicing/marrow_b/{cell_type}/splicing.clusters.significant.tsv', cell_type=cell_type_order),
    output:
        "output/plots/marrow_b_n_events.svg",
    run:
        df = pd.DataFrame({
            "Cell type": cell_type_order,
            "Diff. spl. events": [pd.read_csv(path, "\t").shape[0] for path in input],
        })
        plt.figure(figsize=(2, 2))
        sns.barplot(data=df, x="Cell type", y="Diff. spl. events")
        plt.xticks(rotation=45, ha="right")
        plt.xlabel("")
        sns.despine()
        plt.savefig(output[0], bbox_inches='tight')
