import os
import numpy as np
import pytest
from pybedtools import Interval
from kipoiseq.transforms.functional import one_hot_dna
from kipoiseq.datasets.sequence import SeqStringDataset, SeqDataset, parse_dtype, BedDataset


@pytest.fixture
def fasta_file():
    return "tests/data/sample.fasta"


@pytest.fixture
def intervals_file():
    return "tests/data/sample_intervals.bed"


# fasta_file = fasta_file()
# intervals_file = intervals_file()


def test_min_props():
    # minimal set of properties that need to be specified on the object
    min_set_props = ["output_schema", "type", "defined_as", "info", "args", "dependencies", "postprocessing",
                     "_yaml_path", "source", "source_dir"]

    for Dl in [SeqStringDataset, SeqDataset]:
        props = dir(Dl)
        assert all([el in props for el in min_set_props])


def test_parse_dtype():
    dtypes = {'int': int, 'string': str, 'float': float, 'bool': bool}
    assert all([parse_dtype(dt) == dtypes[dt] for dt in dtypes.keys()])
    assert all([parse_dtype(dt) == dt for dt in dtypes.values()])
    with pytest.raises(Exception):
        parse_dtype("int8")
    assert parse_dtype(None) is None


def test_fasta_based_dataset(intervals_file, fasta_file):
    # just test the functionality
    dl = SeqStringDataset(intervals_file, fasta_file)
    ret_val = dl[0]
    assert isinstance(ret_val["inputs"], np.ndarray)
    assert ret_val["inputs"].shape == ()
    # # test with set wrong seqlen:
    # dl = SeqStringDataset(intervals_file, fasta_file, required_seq_len=3)
    # with pytest.raises(Exception):
    #     dl[0]

    dl = SeqStringDataset(intervals_file, fasta_file, label_dtype="string")
    ret_val = dl[0]
    assert isinstance(ret_val['targets'][0], np.str_)
    dl = SeqStringDataset(intervals_file, fasta_file, label_dtype="int")
    ret_val = dl[0]
    assert isinstance(ret_val['targets'][0], np.int_)
    dl = SeqStringDataset(intervals_file, fasta_file, label_dtype="bool")
    ret_val = dl[0]
    assert isinstance(ret_val['targets'][0], np.bool_)
    vals = dl.load_all()
    assert vals['inputs'][0] == 'GT'


def test_seq_dataset(intervals_file, fasta_file):
    dl = SeqDataset(intervals_file, fasta_file)
    ret_val = dl[0]

    assert np.all(ret_val['inputs'] == one_hot_dna("GT"))
    assert isinstance(ret_val["inputs"], np.ndarray)
    assert ret_val["inputs"].shape == (2, 4)


# download example files
def test_examples_exist():
    ex = SeqStringDataset.init_example()
    out = ex.load_all()
    assert len(out['inputs']) == len(ex)
    assert len(ex) == 4

    ex = SeqDataset.init_example()
    out = ex.load_all()
    assert len(out['inputs']) == len(ex)
    assert len(ex) == 4

@@ -2409,16 +2409,80 @@
 e files%0A
+@pytest.mark.parametrize(%22cls%22, %5BSeqStringDataset, SeqDataset%5D)%0A
 def test
@@ -2497,16 +2497,19 @@
 s_exist(
+cls
 ):%0A    e
@@ -2512,32 +2512,19 @@
    ex = 
-SeqStringDataset
+cls
 .init_ex
@@ -2539,182 +2539,249 @@
     
-out = ex.load_all()%0A    assert len(out%5B'inputs'%5D) == len(ex)%0A    assert len(ex) == 4%0A%0A    ex = SeqDataset.init_example()%0A    out = ex.load_all()%0A    assert len(out%5B'inputs'%5D)
+example_kwargs = cls.example_kwargs%0A    bed_fh = open(example_kwargs%5B'intervals_file'%5D, %22r%22)%0A    bed_entries = len(%5Bel for el in bed_fh%5D)%0A    bed_fh.close()%0A%0A    dl_entries = 0%0A    for out in ex:%0A        dl_entries += 1%0A    assert dl_entries
  == 
@@ -2810,10 +2810,20 @@
 (ex) == 
-4
+bed_entries
 %0A

fix tests