import os

import matplotlib
import pkg_resources
import pytest

from janggu.data import Bioseq
from janggu.data import split_train_test
from janggu.data import subset
from janggu.data import view
from janggu.data.data import _data_props

matplotlib.use('AGG')


def test_dna_props_extraction(tmpdir):
    os.environ['JANGGU_OUTPUT'] = tmpdir.strpath
    data_path = pkg_resources.resource_filename('janggu', 'resources/')
    bed_file = os.path.join(data_path, 'sample.bed')

    refgenome = os.path.join(data_path, 'sample_genome.fa')

    dna = Bioseq.create_from_refgenome('dna', refgenome=refgenome,
                                    storage='ndarray',
                                    roi=bed_file,
                                    binsize=200, stepsize=200,
                                    order=1)

    props = _data_props(dna)
    assert 'dna' in props
    assert props['dna']['shape'] == (200, 1, 4)

    with pytest.raises(Exception):
        _data_props((0,))


def test_split_train_test():
    data_path = pkg_resources.resource_filename('janggu', 'resources/')
    bed_file = os.path.join(data_path, 'sample.bed')

    refgenome = os.path.join(data_path, 'sample_genome.fa')

    dna = Bioseq.create_from_refgenome('dna', refgenome=refgenome,
                                       storage='ndarray',
                                       roi=bed_file,
                                       binsize=200, stepsize=200,
                                       order=1, store_whole_genome=True)

    traindna, testdna = split_train_test(dna, holdout_chroms='chr2')

    assert len(traindna) == 50
    assert len(testdna) == 50
    assert len(dna) == len(traindna) + len(testdna)

    traindna, testdna = split_train_test([dna, dna], holdout_chroms='chr2')

    assert len(traindna[0]) == 50
    assert len(testdna[0]) == 50
    assert len(dna) == len(traindna[0]) + len(testdna[0])


def test_subset_include_chrname_test():
    data_path = pkg_resources.resource_filename('janggu', 'resources/')
    bed_file = os.path.join(data_path, 'sample.bed')

    refgenome = os.path.join(data_path, 'sample_genome.fa')

    dna = Bioseq.create_from_refgenome('dna', refgenome=refgenome,
                                       storage='ndarray',
                                       roi=bed_file,
                                       binsize=200, stepsize=200,
                                       order=1, store_whole_genome=True)

    subdna = subset(dna, include_regions='chr2')

    assert len(subdna) == 50


def test_subset_exclude_chrname_test():
    data_path = pkg_resources.resource_filename('janggu', 'resources/')
    bed_file = os.path.join(data_path, 'sample.bed')

    refgenome = os.path.join(data_path, 'sample_genome.fa')

    dna = Bioseq.create_from_refgenome('dna', refgenome=refgenome,
                                       storage='ndarray',
                                       roi=bed_file,
                                       binsize=200, stepsize=200,
                                       order=1, store_whole_genome=True)

    subdna = subset(dna, exclude_regions='chr2')

    assert len(subdna) == 50


def test_view_bed_test():
    data_path = pkg_resources.resource_filename('janggu', 'resources/')
    bed_file = os.path.join(data_path, 'sample.bed')
    bedsub_file = os.path.join(data_path, 'scored_sample.bed')

    refgenome = os.path.join(data_path, 'sample_genome.fa')

    dna = Bioseq.create_from_refgenome('dna', refgenome=refgenome,
                                       storage='ndarray',
                                       roi=bed_file,
                                       binsize=200, stepsize=200,
                                       order=1, store_whole_genome=True)

    subdna = view(dna, use_regions=bedsub_file)

    assert len(subdna) == 4