"""Tests for `openomics` package."""

from openomics import MessengerRNA, MicroRNA, LncRNA, Protein, SomaticMutation

from openomics import MultiOmics

from .test_clinical import *

cohort_folder_path = "tests/data/TCGA_LUAD/"

@pytest.fixture

def generate_TCGA_LUAD_MessengerRNA():

    data = MessengerRNA(

        data=os.path.join(cohort_folder_path, "LUAD__geneExp.txt"),

        transpose=True,

        usecols="GeneSymbol|TCGA",

        gene_index="GeneSymbol",

        gene_level="gene_name",

    )

    data.drop_genes(data.expressions.columns[50:])

    data.drop_samples(data.expressions.index[:100])

    return data

@pytest.fixture

def generate_TCGA_LUAD_MessengerRNA_dask():

    data = MessengerRNA(

        data=os.path.join(cohort_folder_path, "LUAD__geneExp.txt"),

        transpose=True,

        usecols="GeneSymbol|TCGA",

        gene_index="GeneSymbol",

        gene_level="gene_name",

        npartitions=4,

    )

    data.drop_genes(data.expressions.columns[50:])

    data.init_annotations()

    return data

@pytest.fixture

def generate_TCGA_LUAD_MicroRNA():

    data = MicroRNA(

        data=os.path.join(cohort_folder_path, "LUAD__miRNAExp__RPM.txt"),

        transpose=True,

        usecols="GeneSymbol|TCGA",

        gene_index="GeneSymbol",

        gene_level="gene_name",

    )

    data.drop_genes(data.expressions.columns[50:])

    data.drop_samples(data.expressions.index[:100])

    return data

@pytest.fixture

def generate_TCGA_LUAD_LncRNA():

    data = LncRNA(

        data=os.path.join(cohort_folder_path, "TCGA-rnaexpr.tsv"),

        transpose=True,

        usecols="Gene_ID|TCGA",

        gene_index="Gene_ID",

        gene_level="gene_id",

    )

    data.drop_genes(data.expressions.columns[50:])

    data.drop_samples(data.expressions.index[:100])

    return data

@pytest.fixture

def generate_TCGA_LUAD_SomaticMutation():

    data = SomaticMutation(

        data=os.path.join(cohort_folder_path,

                          "LUAD__somaticMutation_geneLevel.txt"),

        transpose=True,

        usecols="GeneSymbol|TCGA",

        gene_index="gene_name",

    )

    data.drop_genes(data.expressions.columns[50:])

    data.drop_samples(data.expressions.index[:100])

    return data

@pytest.fixture

def generate_TCGA_LUAD_Protein():

    data = Protein(

        data=os.path.join(cohort_folder_path, "protein_RPPA.txt"),

        transpose=True,

        usecols="GeneSymbol|TCGA",

        gene_index="GeneSymbol",

        gene_level="protein_name",

    )

    data.drop_genes(data.expressions.columns[50:])

    data.drop_samples(data.expressions.index[:100])

    return data

def test_import_MessengerRNA_Dask(generate_TCGA_LUAD_MessengerRNA_dask):

    """

    Args:

        generate_TCGA_LUAD_MessengerRNA_dask:

    """

    assert generate_TCGA_LUAD_MessengerRNA_dask.expressions is not None

def test_import_expression_table_size(generate_TCGA_LUAD_MessengerRNA, generate_TCGA_clinical):

    """

    Args:

        generate_TCGA_LUAD_MessengerRNA:

    """

    cohort_name = "LUAD"

    luad_data = MultiOmics(cohort_name)

    luad_data.add_clinical_data(generate_TCGA_clinical)

    luad_data.add_omic(generate_TCGA_LUAD_MessengerRNA)

    luad_data.build_samples()

    print(luad_data.data.keys())

    assert (luad_data.data[MessengerRNA.name()].shape ==

            generate_TCGA_LUAD_MessengerRNA.expressions.shape)

@pytest.fixture

def generate_TCGA_LUAD(

    generate_TCGA_clinical,

    generate_TCGA_LUAD_MessengerRNA,

    generate_TCGA_LUAD_MicroRNA,

    generate_TCGA_LUAD_LncRNA,

    generate_TCGA_LUAD_Protein,

):

    """

    Args:

        generate_TCGA_LUAD_MessengerRNA:

        generate_TCGA_LUAD_MicroRNA:

        generate_TCGA_LUAD_LncRNA:

        generate_TCGA_LUAD_Protein:

    """

    cohort_name = "LUAD"

    luad_data = MultiOmics(cohort_name)

    luad_data.add_clinical_data(generate_TCGA_clinical)

    luad_data.add_omic(generate_TCGA_LUAD_MessengerRNA)

    luad_data.add_omic(generate_TCGA_LUAD_MicroRNA)

    luad_data.add_omic(generate_TCGA_LUAD_LncRNA)

    luad_data.add_omic(generate_TCGA_LUAD_Protein)

    return luad_data

def test_TCGA_LUAD_multiomics_transcriptomics(generate_TCGA_LUAD):

    """

    Args:

        generate_TCGA_LUAD:

    """

    assert all(elem in generate_TCGA_LUAD.get_omics_list() for elem in [

        MessengerRNA.name(),

        MicroRNA.name(),

        LncRNA.name(),

        Protein.name(),

    ])