import dask.dataframe as dd

from openomics.database.sequence import GENCODE, MirBase

from .test_multiomics import *

@pytest.fixture

def generate_GENCODE():

    gencode = GENCODE(path="ftp://ftp.ebi.ac.uk/pub/databases/gencode/Gencode_human/release_32/",

                      file_resources={"long_noncoding_RNAs.gtf": "gencode.v32.long_noncoding_RNAs.gtf.gz",

                                      "lncRNA_transcripts.fa": "gencode.v32.lncRNA_transcripts.fa.gz"}, )

    gencode.data = gencode.data.sample(frac=0.01)

    return gencode

@pytest.fixture

def generate_GENCODE_dask():

    gencode = GENCODE(path="ftp://ftp.ebi.ac.uk/pub/databases/gencode/Gencode_human/release_32/",

                      file_resources={"long_noncoding_RNAs.gtf": "gencode.v32.long_noncoding_RNAs.gtf.gz",

                                      "lncRNA_transcripts.fa": "gencode.v32.lncRNA_transcripts.fa.gz"},

                      blocksize='10MB')

    gencode.data = gencode.data.sample(frac=0.01)

    return gencode

@pytest.fixture

def generate_MirBase_ftp():

    mirbase = MirBase(path="ftp://mirbase.org/pub/mirbase/CURRENT/")

    mirbase.data = mirbase.data.sample(frac=0.01)

    return mirbase

def test_import_mirbase_db(generate_MirBase_ftp):

    """

    Args:

        generate_MirBase_ftp:

    """

    assert generate_MirBase_ftp.data_path == "ftp://mirbase.org/pub/mirbase/CURRENT/"

def test_import_GENCODE(generate_GENCODE):

    """

    Args:

        generate_GENCODE:

    """

    assert generate_GENCODE.data_path == 'ftp://ftp.ebi.ac.uk/pub/databases/gencode/Gencode_human/release_32/'

def test_import_dask_GENCODE(generate_GENCODE_dask):

    assert isinstance(generate_GENCODE_dask.data, dd.DataFrame)

def test_annotate_GENCODE(generate_TCGA_LUAD, generate_GENCODE):

    """

    Args:

        generate_TCGA_LUAD:

        generate_GENCODE:

    """

    generate_TCGA_LUAD.LncRNA.annotate_attributes(generate_GENCODE, on="gene_id", columns=['gene_name'])

    # Test join on off-index

    generate_TCGA_LUAD.LncRNA.annotate_attributes(generate_GENCODE,

                                                  on="gene_name",

                                                  columns=['transcript_id'],

                                                  agg="concat")

    assert {'gene_name', 'transcript_id'}.issubset(

        generate_TCGA_LUAD.LncRNA.annotations.columns)

def test_annotate_dask_GENCODE(generate_TCGA_LUAD, generate_GENCODE_dask):

    """

    Args:

        generate_TCGA_LUAD:

        generate_GENCODE_dask:

    """

    generate_GENCODE_dask.data = generate_GENCODE_dask.data[generate_GENCODE_dask.data["gene_id"].notnull()]

    # Test join on index column

    generate_TCGA_LUAD.LncRNA.annotate_attributes(generate_GENCODE_dask,

                                                  on="gene_id",

                                                  columns=['gene_name'],

                                                  agg="concat")

    # Test join on off-index

    generate_TCGA_LUAD.LncRNA.annotate_attributes(generate_GENCODE_dask,

                                                  on="gene_name",

                                                  columns=['transcript_id'],

                                                  agg="concat")

    assert {'gene_name', 'transcript_id'}.issubset(

        generate_TCGA_LUAD.LncRNA.annotations.columns)

def test_annotate_sequence_GENCODE(generate_TCGA_LUAD, generate_GENCODE):

    generate_TCGA_LUAD.LncRNA.annotate_sequences(generate_GENCODE,

                                                 on="gene_id",

                                                 omic="LncRNA",

                                                 agg="longest")

    assert not generate_TCGA_LUAD.LncRNA.annotations["sequence"].empty