import csv
import itertools
import more_itertools
import matplotlib
matplotlib.use('pdf')
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
#import scanpy as sc
import scipy.sparse as sp_sparse
from scipy.stats import chi2, entropy, kde, mannwhitneyu, ttest_ind
import seaborn as sns
sns.set_theme(style="whitegrid")
from skbio.stats.composition import clr, ilr, alr
from sklearn.decomposition import NMF, PCA
from sklearn.metrics import roc_auc_score, roc_curve
from statsmodels.stats.multitest import multipletests
import re
from tqdm import tqdm
#import GPUtil

import anndata
from scquint.differential_splicing import run_differential_splicing
#from spliceVI.differential_splicing import run_differential_splicing
# from spliceVI.dimensionality_reduction import Dataset, VAE, UnsupervisedTrainer, Posterior
#from spliceVI.utils import relabel, group_normalize, filter_min_cells_per_feature


original_gtf_path = "/global/scratch/projects/fc_songlab/gbenegas/genomes/mus_musculus/mm10/mm10_filt.gtf"
chrom_sizes_path = "/global/scratch/projects/fc_songlab/gbenegas/genomes/mus_musculus/mm10/mm10.chrom.sizes"
sjdb_path = "/global/scratch/projects/fc_songlab/gbenegas/genomes/mus_musculus/mm10/STAR_index/sjdbList.out.tab"
genome_fasta_path = "/global/scratch/projects/fc_songlab/gbenegas/genomes/mus_musculus/mm10/mm10.fa"

n_total_samples_per_group = 300
n_test_samples_per_group = [100, 200, 300]

groups = ["a", "b"]
repeats = [str(x).zfill(2) if x < 10 else str(x) for x in np.arange(1, n_total_samples_per_group+1)]
sample_ids = expand("{group}/sample_{repeat}", group=groups, repeat=repeats)
sample_ids = np.array(sample_ids)
sample_ids2 = np.array([s.replace("/", ".") for s in sample_ids])  # unfortunately needed for leafcutter
subset_genes = 200
subset_transcripts = 2
subset_genes_seed = 42

#groupings = ["nontransitive", "transitive", "gene"]
groupings = ["introns-shared-acceptor"]
methods = [
    "scQuint",
    "LeafCutter",
    "DTUrtle",
    "BRIE2",
    #"original_kallisto_permutation-Euclidean",
    #"assembly-unguided-merged_exons_permutation-Euclidean",
    #"original_bins-nmf_permutation-Euclidean",
    #"original_brie",
    #"original_splicevi-gene_permutation-Euclidean",
    #"original_splicevi-transitive_permutation-Euclidean",
    #"original_splicevi-nontransitive_permutation-Euclidean",
    #"original_introns-shared-acceptor_permutation-Euclidean",
    #"original_exons_permutation-Euclidean",
    #"assembly-unguided-only_exons_permutation-Euclidean",
    #"assembly-guided-only_exons_permutation-Euclidean",
    #"assembly-guided-merged_exons_permutation-Euclidean",
    #"original_bins_permutation-Euclidean",
    #"reconstructed_exons_permutation-Euclidean",
    #"reconstructed_kallisto_permutation-Euclidean",
    #"reconstructed_bins_permutation-Euclidean",
]

fold_changes = [1.5]
#coverages = [0.16, 0.33, 1.0]
#coverages = [0.33, 1.0]
coverages = [1.0, 2.0]
#coverages = [1.0, 2.0]
#size_fractions = [0.05, 0.11, 0.33]
#size_fractions = [0.11, 0.33]
size_fractions = [0.143, 0.33]
seeds = [100]
biases = ["none", "cdnaf"]
#biases = ["none"]
#biases = ["cdnaf"]

datasets = expand("{fold_change}_{coverage}_{size_fraction}_{seed}_{bias}", fold_change=fold_changes, coverage=coverages, size_fraction=size_fractions, seed=seeds, bias=biases)
references = ["unannotated_0", "unannotated_1"]
#references = ["unannotated_0"]

dataset_references = np.concatenate([
    expand("{fold_change}_{coverage}_{size_fraction}_{seed}_{bias}/unannotated_0", fold_change=fold_changes, coverage=coverages, size_fraction=size_fractions, seed=seeds, bias=biases),
    expand("{fold_change}_{coverage}_{size_fraction}_{seed}_none/unannotated_1", fold_change=fold_changes, coverage=coverages, size_fraction=size_fractions, seed=seeds)
])
print(dataset_references)


def flip_coin():
    return np.random.randint(2) == 0


rule all:
    input:
        "output/results/relplot.pdf",
        #expand("output/{dataset}/{reference}/differential_splicing/{n_samples}/LeafCutter/genes.tsv", dataset=datasets, reference=references, n_samples=n_test_samples_per_group),
        #expand("output/{dataset}/{reference}/quantification/LeafCutter/_perind_numers.counts.gz", dataset=datasets, reference=references),
        #"references/unannotated_0_truncated.fa",
        #"output/1.5_1.0_0.33_100/unannotated_0/differential_splicing/100/DTUrtle/genes.tsv",
        #"output/1.5_1.0_0.33_100/unannotated_0/quantification/DTUrtle",
        #"output/1.5_1.0_0.33_100/unannotated_0/quantification/BRIE2/brie_count.h5ad",
        #"output/1.5_1.0_0.33_100/unannotated_0/quantification/SCATS_300/tmp/das_script",
        #"output/1.5_1.0_0.33_100/unannotated_0/differential_splicing/200/BRIE2/clusters.tsv",
        #"output/1.5_1.0_0.33_100/unannotated_0/results/300/summary.txt",
        #"output/1.5_1.0_0.33_100/unannotated_1/results/300/summary.txt",
        #"references/gene_transcript_info.txt"
        #expand("output/{dataset}/{reference}/differential_splicing/{n_samples}/scQuint/clusters.tsv", dataset=datasets, reference=references, n_samples=n_test_samples_per_group),
        #expand("output/{dataset}/{reference}/quantification/scQuint/output/introns-shared-acceptor/adata.h5ad", dataset=datasets, reference=references),
        #expand("output/{dataset}/{reference}/bam_paths.txt", dataset=datasets, reference=references),
        #expand("output/{dataset}/{reference}/reference.gtf", dataset=datasets, reference=references),
        #expand("output/{dataset}/fastq/{group}", dataset=datasets, group=groups),
        #expand("output/{dataset}/fastq", dataset=datasets),
        #"output/1.5_0.33_0.11_100/unannotated_0/results/null/p_value_dist.pdf",
        #expand("output/{dataset}/{reference}/quantification/original/introns-shared-acceptor/adata.h5ad", dataset=datasets, reference=references),
        #expand("output/{dataset}/{reference}/quantification/original/{quantification}/adata.h5ad", dataset=datasets, reference=references, quantification=["kallisto", "bins-nmf", "splicevi-nontransitive"]),
        #expand("output/{dataset}/{reference}/quantification/assembly-unguided-merged/exons/adata.h5ad", dataset=datasets, reference=references),
        #expand("output/{dataset}/{reference}/differential_splicing/{n_samples}/original_brie/genes.tsv", dataset=datasets, reference=references, n_samples=n_test_samples_per_group)
        # expand("output/{dataset}/{reference}/assembly_annotated.gtf", dataset=datasets, reference=references)
        #expand("output/{dataset}/{reference}/assembly.gtf", dataset=datasets, reference=references)
        # expand("output/{dataset}/{reference}/quantification/{mode}/{quantification}/adata.h5ad", dataset=datasets, reference=references, quantification=["kallisto", "exons", "bins"], mode=["original", "reconstructed"])
        #expand("output/{dataset}/{reference}/quantification/splicevi/output/splicing_gene/adata.h5ad", dataset=datasets, reference=references)
        #"output/results/summary.txt",


rule make_empty_blacklist:
    output:
        "references/empty_blacklist.bed"
    shell:
        "touch {output}"


rule filter_gtf:
    input:
        original_gtf_path,
        "/global/scratch/projects/fc_songlab/gbenegas/projects/scquint-analysis/tabula_muris/brie_experiments/AS_events/SE.gtf",
    output:
        "references/unannotated_0.gtf",
        "references/gene_transcript_info.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)
        #chromosomes = ["chr1", "chr2", "chr3"]
        #print("subsetting chromosomes")
        #df = df[df.chromosome.isin(chromosomes)]
        #print(df.shape)
        df['gene_id'] = df.attribute.str.extract(r'gene_id "([^;]*)";')
        df['transcript_id'] = df.attribute.str.extract(r'transcript_id "([^;]*)";')



        SE_gtf = pd.read_csv(
            input[1], '\t', header=None, comment="#",
            names=['chromosome', 'source', 'feature', 'start', 'end', 'score', 'strand', 'frame', 'attribute'],
        )
        #print("subsetting chromosomes again")
        #SE_gtf = SE_gtf[SE_gtf.chromosome.isin(chromosomes)]
        SE_gtf['gene_id'] = SE_gtf.attribute.str.extract(r'gene_id "([^;]*)";')
        SE_gtf['transcript_id'] = SE_gtf.attribute.str.extract(r'transcript_id "([^;]*)";')

        SE_gtf["coordinates"] = SE_gtf.chromosome + ":" + SE_gtf.start.astype(str) + "-" + SE_gtf.end.astype(str)
        SE_exons = SE_gtf[SE_gtf.feature=="exon"].sort_values(["chromosome", "start", "end"]).groupby("transcript_id").coordinates.apply(",".join)

        df["coordinates"] = df.chromosome + ":" + df.start.astype(str) + "-" + df.end.astype(str)
        tx_exons = df[df.feature=="exon"].sort_values(["chromosome", "start", "end"]).groupby("transcript_id").coordinates.apply(",".join)

        def find_transcripts(se_rows):
            se_txs = se_rows.transcript_id.unique()
            if len(se_txs) != 2: return []
            se_t1, se_t2 = se_txs
            se_e1 = SE_exons.loc[se_t1]
            se_e2 = SE_exons.loc[se_t2]

            t1s = tx_exons[tx_exons.str.contains(se_e1)].index.values
            t2s = tx_exons[tx_exons.str.contains(se_e2)].index.values

            for t1 in t1s:
                substring1 = tx_exons.loc[t1].replace(se_e1, "")
                for t2 in t2s:
                    substring2 = tx_exons.loc[t2].replace(se_e2, "")
                    if substring1 == substring2: return [t1, t2]
            return []

        tx_pairs = SE_gtf.groupby("gene_id").apply(find_transcripts).to_frame().rename(columns={0: "tx_pair"})
        print(tx_pairs)
        tx_pairs = tx_pairs[tx_pairs.tx_pair.apply(len) > 0]

        tx2gene = df[["transcript_id", "gene_id"]].drop_duplicates().set_index("transcript_id")
        tx_pairs["gene_id"] = tx_pairs.tx_pair.apply(lambda txs: tx2gene.loc[txs].gene_id.values[0])
        print(tx_pairs)
        tx_pairs = tx_pairs.drop_duplicates("gene_id")
        print(tx_pairs)

        target_transcripts = np.concatenate(tx_pairs.tx_pair.values)
        print(target_transcripts)
        print(len(target_transcripts))
        target_genes = df[df.transcript_id.isin(target_transcripts)].gene_id.unique()
        df = df[df.gene_id.isin(target_genes)]
        df = df[(df.transcript_id.isin(target_transcripts)) | (df.feature=="gene")]
        df = df.drop(columns=["coordinates"])

        #target_transcripts = ["Pex1-201", "Pex1-202", "Adam22-210", "Adam22-213", "Dmtf1-201", "Dmtf1-202", "Magi2-205", "Magi2-201", "Chpf-201", "Chpf-202", "Klhl7-201", "Klhl7-202"]
        #target_genes = ["Pex1", "Adam22", "Dmtf1", "Magi2", "Chpf", "Klhl7"]
        #df["gene_name"] = df.attribute.str.extract(r'gene_name "([^;]*)";')
        #df = df[df.gene_name.isin(target_genes)]
        #df["transcript_name"] = df.attribute.str.extract(r'transcript_name "([^;]*)";')
        #df = df[(df.transcript_name.isin(target_transcripts)) | (df.feature=="gene")]
        #print(df.shape)
        #df = df.drop(columns=["gene_name", "transcript_name"])









        #gene_biotype = df.attribute.str.extract(r'gene_biotype "([^;]*)";').values.ravel()
        #df_exon = df[df.feature=="exon"]
        #df_exon = df_exon[df_exon.attribute.str.contains('tag "basic";', regex=False)]

        #df = df.iloc[(gene_biotype == "protein_coding")]
        #print(df.shape)
        #n_transcripts_per_gene = df_exon.groupby("gene_id").transcript_id.nunique()
        #df = df.iloc[(n_transcripts_per_gene[df.gene_id.values] >= subset_transcripts).values]
        #print(df.shape)

        #gene_ids = df.gene_id.unique()
        #np.random.seed(subset_genes_seed)
        #genes_to_include = np.random.choice(gene_ids, size=subset_genes,replace=False)
        #df = df[df.gene_id.isin(genes_to_include)]
        #print(df.shape)

        #df_exon = df_exon[df_exon.gene_id.isin(genes_to_include)]
        #df = df[(df.transcript_id.isin(df_exon.transcript_id.unique())) | (df.feature=="gene")]

        ## here should filter to transcript feature, to have no dups
        #transcripts_to_include = df[df.feature=="transcript"].groupby("gene_id").transcript_id.sample(n=subset_transcripts, replace=False, random_state=subset_genes_seed).values
        #df = df[(df.transcript_id.isin(transcripts_to_include)) | (df.feature=="gene")]
        #print(df.shape)

        ## TODO: should now filter non-overlap 3'
        #tqdm.pandas()
        #df_exon = df[df.feature=="exon"]

        #def get_transcript_introns(df_transcript):
        #    df_transcript = df_transcript.sort_values("start")
        #    exon_pairs = more_itertools.pairwise(df_transcript.loc[:, ["start", "end"]].values)
        #    introns = [[e1[1], e2[0]] for e1, e2 in exon_pairs]
        #    return introns

        #def check_overlapping_introns(introns1, introns2, strand):
        #    for s1, e1 in introns1:
        #        for s2, e2 in introns2:
        #            if strand == "+":
        #                if s1 != s2 and e1 == e2:
        #                    return True
        #            elif strand == "-":
        #                if s1 == s2 and e1 != e2:
        #                    return True
        #            else:
        #                print(f"strange strand: {strand}")
        #    return False

        #def check_overlapping_introns_gene(df_gene):
        #    transcript_introns = df_gene.groupby("transcript_id").apply(get_transcript_introns).values
        #    comparisons = []
        #    for i, introns1 in enumerate(transcript_introns):
        #        for j, introns2 in enumerate(transcript_introns):
        #            if i < j:
        #                comparisons.append(check_overlapping_introns(introns1, introns2, df_gene.strand.values[0]))
        #    return np.all(comparisons)

        #gene_id_overlapping = df_exon.groupby("gene_id").progress_apply(check_overlapping_introns_gene).reset_index()
        #print(gene_id_overlapping)
        #genes_to_keep = gene_id_overlapping[gene_id_overlapping.loc[:, 0]].gene_id.values
        #df = df[df.gene_id.isin(genes_to_keep)]
        #print(df.shape)

        gene_transcript_info = df[["gene_id", "transcript_id"]].drop_duplicates().dropna()
        print(gene_transcript_info.shape)
        gene_transcript_info.to_csv(output[1], "\t", index=False)

        df = df.drop(columns=["gene_id", "transcript_id"])
        df.to_csv(output[0], "\t", header=False, index=False, quoting=csv.QUOTE_NONE)


rule make_transcript_fasta:
    input:
        "{anything}.gtf",
        genome_fasta_path,
    output:
        "{anything}.fa",
    shell:
        "gffread -w {output} -g {input[1]} {input[0]}"


rule copy_reference:
    input:
        "references/unannotated_0.gtf",
    output:
        "output/{dataset}/unannotated_0/reference.gtf",
    shell:
        "cp {input} {output}"


rule prepare_simulation_params:
    input:
        "references/gene_transcript_info.txt"
    output:
        "output/{fold_change}_{coverage}_{size_fraction}_{seed}_{bias}/simulation_params.txt",
    run:
        diff_fold_change = float(wildcards["fold_change"])
        diff_fold_change_rec = 1.0 / diff_fold_change
        # np.random.seed(int(wildcards["seed"]))
        np.random.seed(subset_genes_seed)

        gene_transcript_info = pd.read_csv(input[0], "\t")
        gene_ids = gene_transcript_info.gene_id.unique()
        gene_transcript_info["coverage"] = float(wildcards["coverage"])
        gene_transcript_info["size_fraction"] = float(wildcards["size_fraction"])
        gene_transcript_info["fold_change_a"] = 1.0
        gene_transcript_info["fold_change_b"] = 1.0
        gene_transcript_info["differential"] = False
        gene_fold_change = 1.5
        print(gene_transcript_info)
        gene_transcript_info = gene_transcript_info.set_index(["gene_id", "transcript_id"])
        print(gene_transcript_info)
        differential_genes = np.random.choice(gene_ids, size=len(gene_ids)//2, replace=False)
        print("differential_genes.shape: ", differential_genes.shape)
        print("np.unique(differential_genes).shape: ", np.unique(differential_genes).shape)

        for gene_id in gene_ids:
            df_gene = gene_transcript_info.loc[gene_id]
            # Always overexpress the whole gene in one condition
            if np.random.randint(2) == 0:
                col = "fold_change_a"
            else:
                col = "fold_change_b"
            df_gene.loc[:, col] = gene_fold_change

            # For genes chosen to be differential, choose one isoform and overexpress
            if gene_id in differential_genes:
                diff_isoform = np.random.randint(len(df_gene))
                if np.random.randint(2) == 0:
                    col = "fold_change_a"
                else:
                    col = "fold_change_b"
                if np.random.randint(2) == 0:
                    final_fold_change = diff_fold_change
                else:
                    final_fold_change = diff_fold_change_rec
                fc = df_gene[col]
                fc.iloc[diff_isoform] *= final_fold_change
                differential = df_gene["differential"]
                differential.iloc[diff_isoform] = True
                df_gene.loc[:, col] = fc
                df_gene.loc[:, "differential"] = differential
            gene_transcript_info.loc[gene_id] = df_gene.values
        gene_transcript_info.reset_index().to_csv(output[0], "\t", index=False)


rule mask_gtf:
    input:
        "output/{dataset}/simulation_params.txt",
        "output/{dataset}/unannotated_0/reference.gtf",
    output:
        "output/{dataset}/unannotated_1/reference.gtf",
    run:
        sim_params = pd.read_csv(input[0], "\t")
        #transcripts_to_keep = sim_params[~sim_params.differential].transcript_id.values
        transcripts_to_keep = sim_params.sort_values("gene_id").transcript_id.values[::2]  # only keep one per gene
        df = pd.read_csv(
            input[1], '\t', header=None, comment="#",
            names=['chromosome', 'source', 'feature', 'start', 'end', 'score', 'strand', 'frame', 'attribute'],
        )
        df['transcript_id'] = df.attribute.str.extract(r'transcript_id "([^;]*)";')
        df = df[(df.transcript_id.isin(transcripts_to_keep)) | (df.feature=="gene")]
        df = df.drop(columns=["transcript_id"])
        df.to_csv(output[0], "\t", header=False, index=False, quoting=csv.QUOTE_NONE)


rule truncate_fasta:
    input:
        "references/unannotated_0.fa",
    output:
        "references/unannotated_0_truncated.fa",
    run:
        from Bio import SeqIO
        records = list(SeqIO.parse(input[0], "fasta"))
        print([len(record) for record in records])
        new_records = []
        for i in range(len(records)//2):
            r1 = records[2*i]
            r2 = records[2*i+1]
            l = min(len(r1), len(r2))
            new_records.append(r1[-l:])
            new_records.append(r2[-l:])
        print([len(record) for record in new_records])
        SeqIO.write(new_records, output[0], "fasta")


rule simulate_reads:
    input:
        "references/unannotated_0.fa",
        "output/{fold_change}_{coverage}_{size_fraction}_{seed}_{bias}/simulation_params.txt",
        "references/unannotated_0_truncated.fa",
    output:
        directory("output/{fold_change}_{coverage}_{size_fraction}_{seed}_{bias}/fastq/{group}")
    shell:
        "Rscript run_polyester.R {input[0]} {input[1]} {n_total_samples_per_group} {wildcards.seed} {wildcards.group} {wildcards.bias} {output} {input[2]}"


rule prepare_fastq_paths:
    input:
        expand("output/{{dataset}}/fastq/{group}", group=groups),
    output:
        "output/{dataset}/fastq_paths.txt"
    run:
        df = pd.DataFrame(sample_ids, columns=["sample_id"])
        cwd = os.getcwd()
        print("cwd: ", cwd)
        df["fastq_1"] = cwd + f"/output/{wildcards.dataset}/fastq/" + df.sample_id + "_1.fasta.gz"
        df["fastq_2"] = cwd + f"/output/{wildcards.dataset}/fastq/" + df.sample_id + "_2.fasta.gz"
        df.to_csv(output[0], "\t", index=False, header=False)


rule prepare_chromosomes_file:
    input:
        "output/{dataset}/{reference}/reference.gtf"
    output:
        "output/{dataset}/{reference}/chromosomes.txt"
    shell:
        "cut -f1 {input} | sort | uniq > {output}"

## ln -s /global/scratch/gbenegas/spliceVI-project/spliceVI-analysis/simulation/output/first_pass/unannotated_0/runs/splicevi/first_STAR_index first_STAR_index
## ln -s /global/scratch/gbenegas/spliceVI-project/spliceVI-analysis/simulation/output/first_pass/unannotated_0/runs/splicevi/first_STAR_index second_STAR_index


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


rule index_bam:
    input:
        "output/{dataset}/{reference}/mapping/rmdup/{sample_id}/Aligned.rmdup.out.bam",
    output:
        "output/{dataset}/{reference}/mapping/rmdup/{sample_id}/Aligned.rmdup.out.bam.bai",
    shell:
        "samtools index {input}"


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


rule scQuint_quantification:
    input:
        "output/{dataset}/{reference}/bam_paths.txt",
        "output/{dataset}/{reference}/chromosomes.txt",
        "output/{dataset}/{reference}/reference.gtf",
    threads: workflow.cores
    output:
        expand("output/{{dataset}}/{{reference}}/quantification/scQuint/output/{grouping}/adata.h5ad", grouping=groupings),
    shell:
        "python -m scquint.quantification.run introns/Snakefile --cores {threads} -d output/{wildcards.dataset}/{wildcards.reference}/quantification/scQuint/ --config min_cells_per_intron=10 bam_paths=../../bam_paths.txt chromosomes_path=../../chromosomes.txt fasta_path={genome_fasta_path} gtf_path=../../reference.gtf chrom_sizes_path={chrom_sizes_path} sjdb_path=../../mapping/first_STAR_index/sjdbList.out.tab"


rule LeafCutter_extract:
    input:
        lambda wildcards: f"output/{wildcards.dataset}/{wildcards.reference}/mapping/rmdup/{wildcards.sample_id.replace('.', '/')}/Aligned.rmdup.out.bam",
        lambda wildcards: f"output/{wildcards.dataset}/{wildcards.reference}/mapping/rmdup/{wildcards.sample_id.replace('.', '/')}/Aligned.rmdup.out.bam.bai",
    output:
        "output/{dataset}/{reference}/quantification/LeafCutter/juncs/{sample_id}.junc",
    shell:
        """
        regtools junctions extract -s 0 -a 8 -m 50 -M 500000 {input[0]} -o {output}
        """


rule LeafCutter_prepare_junc_files:
    input:
        expand("output/{{dataset}}/{{reference}}/quantification/LeafCutter/juncs/{sample_id}.junc", sample_id=sample_ids2),
    output:
        "output/{dataset}/{reference}/quantification/LeafCutter/juncfiles.txt",
    run:
        for path in input:
            shell(f"echo {path} >> {output}")


rule LeafCutter_cluster:
    input:
        "output/{dataset}/{reference}/quantification/LeafCutter/juncfiles.txt",
    output:
        "output/{dataset}/{reference}/quantification/LeafCutter/_perind_numers.counts.gz",
    params:
        "output/{dataset}/{reference}/quantification/LeafCutter/"
    threads:
        workflow.cores  # important to avoid overwriting temp files
    shell:
        """
        module load python/2.7 &&
        python ./leafcutter/clustering/leafcutter_cluster_regtools.py --checkchrom -j {input} -o {params} -l 500000 &&
        module unload python/2.7
        """

rule LeafCutter_prepare_groups_file:
    output:
        "output/{dataset}/{reference}/differential_splicing/{n_samples}/LeafCutter/input.tsv",
    run:
        cell_idx_a = np.arange(int(wildcards["n_samples"]))
        cell_idx_b = np.arange(int(wildcards["n_samples"])) + n_total_samples_per_group
        cell_idx = np.concatenate([cell_idx_a, cell_idx_b])
        samID = sample_ids2[cell_idx]
        condition = np.ones_like(samID)
        condition[:len(cell_idx_a)] = 0
        pd.DataFrame(dict(samID=samID, condition=condition)).to_csv(output[0], "\t", index=False, header=False)


rule leafcutter_diff_spl:
    input:
        "output/{dataset}/{reference}/quantification/LeafCutter/_perind_numers.counts.gz",
        "output/{dataset}/{reference}/differential_splicing/{n_samples}/LeafCutter/input.tsv",
        "leafcutter/scripts/mm10.txt.gz",
    output:
        "output/{dataset}/{reference}/differential_splicing/{n_samples}/LeafCutter/_cluster_significance.txt",
    params:
        "output/{dataset}/{reference}/differential_splicing/{n_samples}/LeafCutter/",
    threads:
        workflow.cores
    shell:
        "./leafcutter/scripts/leafcutter_ds.R --exon_file {input[2]} --num_threads {threads} --min_samples_per_intron 30 --min_samples_per_group 30 --min_coverage 1 --max_cluster_size=10 -o {params} {input[0]} {input[1]}"


rule clean_leafcutter_output:
    input:
        "{anything}/_cluster_significance.txt",
        "input/gene_name.txt",
        "references/gene_transcript_info.txt",
    output:
        "{anything}/genes.tsv",
    run:
        clusters = pd.read_csv(input[0], "\t")
        clusters = clusters[clusters.status=="Success"]
        clusters = clusters.sort_values("p")
        #clusters.rename(columns={"p.adjust": "p_value_gene_adj", "genes": "gene_name"}, inplace=True)
        clusters.rename(columns={"p": "p_value_gene_adj", "genes": "gene_name"}, inplace=True)

        gene_name_to_id = pd.read_csv(input[1], "\t", index_col=1)
        target_genes = pd.read_csv(input[2], "\t").gene_id.unique()
        gene_name_to_id = gene_name_to_id[gene_name_to_id.gene_id.isin(target_genes)]
        print(gene_name_to_id)

        clusters["gene_id"] = clusters.gene_name.apply(lambda genes: [gene_name_to_id.loc[gene_name].gene_id for gene_name in genes.split(",") if gene_name in gene_name_to_id.index.values][0])
        print(clusters)
        clusters.to_csv(output[0], "\t", index=False)


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


#rule kallisto_quantification:
#    input:
#        "output/{dataset}/fastq_paths.txt",
#        "output/{dataset}/fastq",
#        "output/{dataset}/{reference}/quantification/{mode}/reference.gtf",
#    output:
#        "output/{dataset}/{reference}/quantification/{mode}/kallisto/adata.h5ad"
#    threads: workflow.cores
#    priority: 100
#    shell:
#        "python -m spliceVI.quantification.run kallisto/Snakefile --cores {threads} -d output/{wildcards.dataset}/{wildcards.reference}/quantification/{wildcards.mode}/kallisto/ --config genome_fasta_path={genome_fasta_path} gtf_path=../reference.gtf fastq_paths=../../../../fastq_paths.txt min_cells_per_isoform=10"
#
#
#rule quantify_exons:
#    input:
#        "output/{dataset}/{reference}/bam_paths.txt",
#        "output/{dataset}/{reference}/quantification/{mode}/reference.gtf",
#    threads: workflow.cores
#    priority: 100
#    output:
#        "output/{dataset}/{reference}/quantification/{mode}/exons/adata.h5ad"
#    shell:
#        "python -m spliceVI.quantification.run exons/Snakefile --cores all -d output/{wildcards.dataset}/{wildcards.reference}/quantification/{wildcards.mode}/exons/ --config min_cells_per_exon=10 gtf_path=../reference.gtf bam_paths=../../../bam_paths.txt"
#
#
#rule quantify_bins:
#    input:
#        "output/{dataset}/{reference}/bam_paths.txt",
#        "output/{dataset}/{reference}/quantification/{mode}/reference.gtf",
#    threads: workflow.cores
#    priority: 100
#    output:
#        "output/{dataset}/{reference}/quantification/{mode}/bins/adata.h5ad"
#    shell:
#        "python -m spliceVI.quantification.run bins/Snakefile --cores all -d output/{wildcards.dataset}/{wildcards.reference}/quantification/{wildcards.mode}/bins/ --config min_cells_per_bin=10 gtf_path=../reference.gtf bam_paths=../../../bam_paths.txt"
#
#
#rule transform_adata_with_NMF:
#    input:
#        "output/{dataset}/{reference}/quantification/{mode}/bins/adata.h5ad"
#    output:
#        "output/{dataset}/{reference}/quantification/{mode}/bins-nmf/adata.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 quantify_exons_and_introns:
##    input:
##        "output/{dataset}/{reference}/quantification/exons/adata.h5ad",
##        "output/{dataset}/{reference}/quantification/splicevi-gene/adata.h5ad",
##    output:
##        "output/{dataset}/{reference}/quantification/exons-and-introns/adata.h5ad",
##    run:
##        adata_exons = anndata.read_h5ad(input[0])
##        adata_introns = anndata.read_h5ad(input[1])
##        adata_introns = adata_introns[:, ~adata_introns.var.annotated]
##        print(adata_exons.shape, adata_introns.shape)
##        adata = anndata.AnnData(
##            X=sp_sparse.hstack([adata_exons.X, adata_introns.X], format="csr"),
##            var=pd.concat([adata_exons.var, adata_introns.var], ignore_index=True)
##        )
##        print(adata.shape)
##        adata.var.loc[:, "cluster"] = relabel(adata.var.gene_id)
##        adata.var["original_cluster"] = adata.var.cluster
##        adata.write(output[0], compression="gzip")
#
#

rule BRIE2_quantification:
    input:
        "output/{dataset}/{reference}/reference.gtf",
        "output/{dataset}/{reference}/bam_paths.txt",
    output:
        "output/{dataset}/{reference}/quantification/BRIE2/brie_count.h5ad",
    #threads: workflow.cores
    threads: 4
    shell:
        """
        awk '{{print $2 "\t" $1}}' {input[1]} > {input[1]}.swap.txt &&
        mkdir -p output/{wildcards.dataset}/{wildcards.reference}/quantification/BRIE2 &&
        cd output/{wildcards.dataset}/{wildcards.reference}/quantification/BRIE2 &&
        /global/scratch/projects/fc_songlab/gbenegas/software/miniconda2/bin/briekit-event -a ../../reference.gtf -o AS_events &&
        /global/scratch/projects/fc_songlab/gbenegas/software/miniconda2/bin/briekit-event-filter -a AS_events/SE.gff3.gz --anno_ref=../../reference.gtf -r {genome_fasta_path} --keep_overlap --no_splice_site --as_exon_min=1 --as_exon_max=10000 --as_exon_tss=1 --as_exon_tts=1 &&
        brie-count -a AS_events/SE.filtered.gff3.gz -S ../../bam_paths.txt.swap.txt -o . -p {threads}
        """


rule BRIE2_diff_spl_unannotated_0:
    input:
        "output/{dataset}/unannotated_0/quantification/BRIE2/brie_count.h5ad",
    output:
        "output/{dataset}/unannotated_0/differential_splicing/{n_samples}/BRIE2/cell_anno.tsv",
        "output/{dataset}/unannotated_0/differential_splicing/{n_samples}/BRIE2/brie_quant.h5ad",
        "output/{dataset}/unannotated_0/differential_splicing/{n_samples}/BRIE2/clusters.tsv",
    run:
        cell_idx_a = np.arange(int(wildcards["n_samples"]))
        cell_idx_b = np.arange(int(wildcards["n_samples"])) + n_total_samples_per_group
        cell_idx = np.concatenate([cell_idx_a, cell_idx_b])
        samID = sample_ids[cell_idx]
        condition = np.ones_like(samID)
        condition[:len(cell_idx_a)] = 0
        pd.DataFrame(dict(samID=samID, condition=condition)).to_csv(output[0], "\t", index=False)
        shell(f"brie-quant -i {input[0]} -c {output[0]} -o {output[1]} --interceptMode gene --LRTindex 0 --minCount 20 --minCell 20")
        df = pd.read_csv(output[1].replace(".h5ad", ".brie_ident.tsv"), "\t")
        df["gene_id"] = df.GeneID
        df["p_value"] = df.condition_pval
        df.sort_values("p_value").to_csv(output[2], "\t", index=False)


rule BRIE2_diff_spl_unannotated_1:
    output:
        "output/{dataset}/unannotated_1/differential_splicing/{n_samples}/BRIE2/clusters.tsv",
    run:
        pd.DataFrame(columns=["gene_id", "p_value"]).to_csv(output[0], "\t")


rule DTUrtle_quantification:
    input:
        "output/{dataset}/{reference}/reference.fa",
        "output/{dataset}/fastq_paths.txt",
    output:
        directory("output/{dataset}/{reference}/quantification/DTUrtle"),
    threads: 1
    run:
        shell(f"mkdir -p {output[0]}")
        shell(f"salmon index -t {input[0]} -p {threads} -i {output[0]}/index")
        fastq_paths = pd.read_csv(input[1], "\t", header=None).values
        print(fastq_paths)
        for sample_id, path1, path2 in fastq_paths:
            print(sample_id)
            shell(f"salmon quant -l A -i {output[0]}/index -1 {path1} -2 {path2} --seqBias --gcBias --validateMappings --rangeFactorizationBins 4 --threads 1 -o {output[0]}/{sample_id}")


rule DTUrtle_diff_spl_unannotated_0:
    input:
        "output/{dataset}/unannotated_0/reference.gtf",
        "output/{dataset}/unannotated_0/quantification/DTUrtle/",
    output:
        "output/{dataset}/unannotated_0/differential_splicing/{n_samples}/DTUrtle/sample_metadata.tsv",
        "output/{dataset}/unannotated_0/differential_splicing/{n_samples}/DTUrtle/genes.tsv",
    run:
        cell_idx_a = np.arange(int(wildcards["n_samples"]))
        cell_idx_b = np.arange(int(wildcards["n_samples"])) + n_total_samples_per_group
        cell_idx = np.concatenate([cell_idx_a, cell_idx_b])

        sids = sample_ids[cell_idx]
        condition = ["a"] * len(cell_idx_a) + ["b"] * len(cell_idx_b)

        metadata = pd.DataFrame(dict(
            sample_id=sids,
            condition=condition,
            quant_path = [f"{input[1]}/{sample_id}/quant.sf" for sample_id in sids],
        ))
        metadata.to_csv(output[0], "\t", index=False)

        shell(f"Rscript run_DTUrtle.R {input[0]} {output[0]} {output[1]}")


rule DTUrtle_diff_spl_unannotated_1:
    output:
        "output/{dataset}/unannotated_1/differential_splicing/{n_samples}/DTUrtle/clusters.tsv",
    run:
        pd.DataFrame(columns=["gene_id", "p_value"]).to_csv(output[0], "\t")


#rule SCATS_make_metafile:
#    input:
#        "output/{dataset}/{reference}/bam_paths.txt",
#    output:
#        "output/{dataset}/{reference}/quantification/SCATS_{n_samples}/metafile",
#    run:
#        df = pd.read_csv(input[0], "\t", header=None, names=["sample_id", "bam_path"])
#        if wildcards["n_samples"] == "all":
#            n_samples = n_total_samples_per_group
#        else:
#            n_samples = int(wildcards["n_samples"])
#        cell_idx_a = np.arange(n_samples)
#        cell_idx_b = np.arange(n_samples) + n_total_samples_per_group
#        df.loc[cell_idx_a, "condition"] = "A"
#        df.loc[cell_idx_b, "condition"] = "B"
#        df["UB"] = "UB"
#        df["CB"] = "CB"
#        print(df)
#        df.to_csv(output[0], "\t", header=False, index=False, columns=["sample_id", "condition", "bam_path", "UB", "CB"])
#
#
#rule SCATS_process_gtf:
#    input:
#        "output/{dataset}/{reference}/reference.gtf",
#        "output/{dataset}/{reference}/quantification/SCATS_{n_samples}/metafile",
#    output:
#        "output/{dataset}/{reference}/quantification/SCATS_{n_samples}/reference.genePred",
#        "output/{dataset}/{reference}/quantification/SCATS_{n_samples}/reference.refFile",
#        "output/{dataset}/{reference}/quantification/SCATS_{n_samples}/reference.refgene",
#        "output/{dataset}/{reference}/quantification/SCATS_{n_samples}/reference.gpinfo",
#        expand("output/{{dataset}}/{{reference}}/quantification/SCATS_{{n_samples}}/tmp/count_script/count_{sample_id}.sh", sample_id=sample_ids),
#    shell:
#        """
#        cd output/{wildcards.dataset}/{wildcards.reference}/quantification/SCATS_{wildcards.n_samples}/ &&
#        gtfToGenePred -genePredExt ../../reference.gtf reference.genePred &&
#        awk '{{print "1 \t" $0}}' reference.genePred > reference.refFile &&
#        /global/scratch/projects/fc_songlab/gbenegas/software/miniconda2/bin/python ../../../../../SCATS/SCATS.py -task refgene -ref reference.refFile -out reference.refgene &&
#        /global/scratch/projects/fc_songlab/gbenegas/software/miniconda2/bin/python ../../../../../SCATS/SCATS.py -task group -refgene reference.refgene -out reference.gpinfo &&
#        /global/scratch/projects/fc_songlab/gbenegas/software/miniconda2/bin/python ../../../../../SCATS/SCATS.py -task count -meta metafile -refgene reference.refgene -gpinfo reference.gpinfo -umi no -onebam no
#        """
#
#
#rule SCATS_quantify:
#    input:
#        "output/{dataset}/{reference}/quantification/SCATS_{n_samples}/tmp/count_script/count_{sample_id}.sh",
#    output:
#        "output/{dataset}/{reference}/quantification/SCATS_{n_samples}/tmp/count_script/count_{sample_id}.out",
#    shell:
#        "bash {input}"
#
#
#rule SCATS_diff_spl_part_1:
#    input:
#        "output/{dataset}/{reference}/quantification/SCATS_{n_samples}/metafile",
#        "output/{dataset}/{reference}/quantification/SCATS_{n_samples}/reference.gpinfo",
#        expand("output/{{dataset}}/{{reference}}/quantification/SCATS_{{n_samples}}/tmp/count_script/count_{sample_id}.out", sample_id=sample_ids),
#    output:
#        directory("output/{dataset}/{reference}/quantification/SCATS_{n_samples}/tmp/das_script"),
#    threads: workflow.cores
#    shell:
#        """
#        mkdir -p {output} &&
#        cd output/{wildcards.dataset}/{wildcards.reference}/quantification/SCATS_{wildcards.n_samples}/ &&
#        /global/scratch/projects/fc_songlab/gbenegas/software/miniconda2/bin/python ../../../../../SCATS/SCATS.py -task das -ncore 1 -meta metafile -gpinfo reference.gpinfo &&
#        PATH=/global/scratch/projects/fc_songlab/gbenegas/software/miniconda2/bin:$PATH bash tmp/das_script/das_A_B.sh
#        """

rule scQuint_diff_spl:
    input:
        "output/{dataset}/{reference}/quantification/scQuint/output/introns-shared-acceptor/adata.h5ad",
    output:
        "output/{dataset}/{reference}/differential_splicing/{n_samples}/scQuint/clusters.tsv",
        "output/{dataset}/{reference}/differential_splicing/{n_samples}/scQuint/introns.tsv",
    threads: 1
    run:
        adata = anndata.read_h5ad(input[0])
        print(adata)
        adata.var["original_cluster"] = adata.var.cluster
        if wildcards["n_samples"] == "null":
            cell_idx_a = np.arange(n_total_samples_per_group//2)
            cell_idx_b = np.arange(n_total_samples_per_group//2, n_total_samples_per_group)
        else:
            cell_idx_a = np.arange(int(wildcards["n_samples"]))
            cell_idx_b = np.arange(int(wildcards["n_samples"])) + n_total_samples_per_group

        diff_spl_clusters, diff_spl_introns = run_differential_splicing(
            adata, cell_idx_a, cell_idx_b, min_cells_per_cluster=30, min_total_cells_per_intron=30,
            n_jobs=1, do_recluster=False,
        )
        diff_spl_clusters.to_csv(output[0], '\t')
        diff_spl_introns.to_csv(output[1], '\t')


rule aggregate_diff_spl:
    input:
        "output/{dataset}/{reference}/differential_splicing/{n_samples}/{method}/clusters.tsv",
    output:
        "output/{dataset}/{reference}/differential_splicing/{n_samples}/{method}/genes.tsv",
    run:
        df = pd.read_csv(input[0], "\t")

        if len(df) > 0:
            def multiple_testing_correction(p_values):
                #reject, pvals_corrected, _, _ = multipletests(
                #    p_values.values, 0.05, "fdr_bh"
                #)
                #res = np.min(pvals_corrected)
                res = np.min(p_values)
                return res

            res = df.groupby("gene_id").p_value.agg(multiple_testing_correction).to_frame().reset_index().rename(columns={"p_value": "p_value_gene"})
            #reject, pvals_corrected, _, _ = multipletests(
            #    res.p_value_gene.values, 0.05, "fdr_bh"
            #   )
            res["p_value_gene_adj"] = res.p_value_gene
        else:
            res = pd.DataFrame(columns=["gene_id", "p_value_gene_adj"])
        res.to_csv(output[0], "\t", index=False)


rule dataset_results:
    input:
        "output/{dataset}/simulation_params.txt",
        expand("output/{{dataset}}/{{reference}}/differential_splicing/{{n_samples}}/{method}/genes.tsv", method=methods)
    output:
        "output/{dataset}/{reference}/results/{n_samples}/all.txt",
        "output/{dataset}/{reference}/results/{n_samples}/summary.txt",
    run:
        sim_params = pd.read_csv(input[0], "\t")
        print(sim_params.shape)
        df = sim_params.groupby("gene_id").agg({"fold_change_a": "max", "fold_change_b": "max", "differential": "any"})
        print("df.differential.mean(): ", df.differential.mean())
        #raise Exception("debug")
        #_, _, _, seed = wildcards["dataset"].split("_")
        #seed = int(seed)

        #baseline_method = None
        #baseline_method = "original_splicevi-transitive_permutation-Euclidean"

        summary = []
        for method in methods:
            res = pd.read_csv(f"output/{wildcards['dataset']}/{wildcards['reference']}/differential_splicing/{wildcards['n_samples']}/{method}/genes.tsv", "\t").set_index("gene_id")
            print(method, res)
            df[method] = res.reindex(df.index.values).p_value_gene_adj
            print(method, wildcards["reference"])
            print(df[method])

            # option 1: replace untested with 1.0
            #df[method] = df[method].fillna(1.0)

            # option 2: fill untested at random as significant/non-significant
            #np.random.seed(seed)
            #df[method] = df[method].apply(fill_at_random)
            #print(df[method])


        min_value = 0.0 - 1.0
        max_value = 1.0 + 1.0
        print(min_value, max_value)

        def fill_at_random(x):
            if pd.isnull(x):
                if flip_coin():
                    return min_value
                else:
                    return max_value
            else:
                return x

        #print(df.shape)
        #df = df.dropna()
        #print(df.shape)

        for method in methods:
            np.random.seed(42)
            #auc = np.mean([roc_auc_score(~df.differential.astype(int), df[method].apply(fill_at_random)) for _ in range(100)])
            y_true = ~df.differential.astype(int)
            y_score = df[method]
            print(y_score)
            #mask = ~np.isnan(y_score)
            mask = ~y_score.isna()
            print("sum(mask): ", sum(mask))
            y_true = y_true[mask]
            y_score = y_score[mask]
            if len(y_true) > 0:
                auc = roc_auc_score(y_true, y_score)
            else:
                auc = np.nan
            summary.append([method, auc])
        summary = pd.DataFrame(summary, columns=["Method", "AUC"])
        print(summary)
        df.to_csv(output[0], "\t")
        summary.to_csv(output[1], "\t", index=False)


rule overall_results:
    input:
        expand("output/{dataset_reference}/results/{n_samples}/summary.txt", dataset_reference=dataset_references, n_samples=n_test_samples_per_group)
    output:
        "output/results/summary.txt"
    run:
        overall_summary = []
        #for fold_change, coverage, size_fraction, seed, bias, reference, n_samples in itertools.product(fold_changes, coverages, size_fractions, seeds, biases, references, n_test_samples_per_group):
        for dataset_reference, n_samples in itertools.product(dataset_references, n_test_samples_per_group):
            dataset, reference = dataset_reference.split("/")
            fold_change, coverage, size_fraction, seed, bias = dataset.split("_")
            path = f"output/{fold_change}_{coverage}_{size_fraction}_{seed}_{bias}/{reference}/results/{n_samples}/summary.txt"
            summary = pd.read_csv(path, "\t")
            summary["Fold change"] = fold_change
            summary["Coverage"] = coverage
            summary["size_fraction"] = size_fraction
            summary["seed"] = seed
            summary["bias"] = bias
            summary["Reference"] = reference
            summary["N cells"] = n_samples
            overall_summary.append(summary)
        overall_summary = pd.concat(overall_summary, ignore_index=True)
        overall_summary.to_csv(output[0], "\t", index=False)


rule plot_overall_results:
    input:
        "output/results/summary.txt"
    output:
        "output/results/relplot.pdf",
        "output/results/relplot2.pdf",
    run:
        df = pd.read_csv(input[0], "\t")
        df["Experiment"] = "Even coverage"
        df.loc[df.Reference=="unannotated_1", "Experiment"] = "Unannotated"
        df.loc[df.bias=="cdnaf", "Experiment"] = "Coverage decay"

        df["var_coeff"] = (1 + 1 / df.size_fraction).round(0).astype(int)
        print(df.var_coeff.unique())
        g = sns.relplot(
            data=df,
            x="N cells",
            y="AUC",
            hue="Method",
            style="Experiment",
            col="var_coeff",
            col_order=sorted(df.var_coeff.unique(), reverse=True),
            row="Coverage",
            facet_kws=dict(margin_titles=True),
            kind="line",
            err_style="bars",
            height=3,
            aspect=1.0,
        )
        g.set_titles(col_template="Variance = {col_name} * mean", row_template="Base coverage = {row_name}x")
        g.set(xticks=n_test_samples_per_group)
        plt.savefig(output[0], bbox_inches='tight')

        #df = df[(df.Coverage==1.0) & (df.var_coeff==8)]
        df = df[(df.var_coeff==8)]
        print(df.Coverage.value_counts())
        print(df.var_coeff.value_counts())
        g = sns.relplot(
            data=df,
            y="AUC",
            hue="Method",
            hue_order=methods,
            x="N cells",
            col="Experiment",
            col_order=["Even coverage", "Unannotated", "Coverage decay"],
            row="Coverage",
            kind="line",
            height=3,
            aspect=1.0,
            facet_kws=dict(margin_titles=True),
        )
        g.set_titles(row_template="Base coverage = {row_name}x")
        g.set(xticks=n_test_samples_per_group)
        plt.savefig(output[1], bbox_inches='tight')
