"""Test bpnet train

"""

import os

import pytest

from bpnet.cli.train import bpnet_train

from pathlib import Path

from bpnet.seqmodel import SeqModel

from concise.preprocessing import encodeDNA

import gin

import keras.backend as K

def test_output_files(trained_model):

    K.clear_session()

    output_files = os.listdir(str(trained_model))

    expected_files = [

        'config.gin',

        'config.gin.json',

        'bpnet-train.kwargs.json',

        'dataspec.yml',

        'evaluate.ipynb',

        'evaluate.html',

        'evaluation.valid.json',

        'history.csv',

        'model.h5',

        'seq_model.pkl',

        'note_params.json',

    ]

    for f in expected_files:

        assert f in output_files

    m = SeqModel.load(trained_model / 'seq_model.pkl')

    m.predict(encodeDNA(["A" * 200]))

def test_output_files_model_w_bias(trained_model_w_bias):

    K.clear_session()

    output_files = os.listdir(str(trained_model_w_bias))

    expected_files = [

        'config.gin',

        'config.gin.json',

        'bpnet-train.kwargs.json',

        'dataspec.yml',

        'evaluate.ipynb',

        'evaluate.html',

        'evaluation.valid.json',

        'history.csv',

        'model.h5',

        'seq_model.pkl',

        'note_params.json',

    ]

    for f in expected_files:

        assert f in output_files

    m = SeqModel.load(trained_model_w_bias / 'seq_model.pkl')

    m.predict(encodeDNA(["A" * 200]))

def test_trained_model_bed6(tmp_path, data_dir, config_gin, dataspec_bed6):

    K.clear_session()

    gin.clear_config()

    bpnet_train(dataspec=str(dataspec_bed6),

                output_dir=str(tmp_path),

                premade='bpnet9',

                config=str(config_gin),

                override='seq_width=100;train.batch_size=8',

                num_workers=2

                )

def test_trained_model_override_in_memory(tmp_path, data_dir, config_gin, dataspec_bias):

    K.clear_session()

    gin.clear_config()

    bpnet_train(dataspec=str(dataspec_bias),

                output_dir=str(tmp_path),

                premade='bpnet9',

                config=str(config_gin),

                in_memory=True,

                override='seq_width=190;train.batch_size=8',

                num_workers=2

                )

def test_train_regions(tmp_path, data_dir, config_gin, dataspec_bias, regions):

    K.clear_session()

    gin.clear_config()

    bpnet_train(dataspec=str(dataspec_bias),

                output_dir=str(tmp_path),

                premade='bpnet9',

                config=str(config_gin),

                override=f'bpnet_data.intervals_file="{regions}"',

                num_workers=2

                )

def test_trained_model_premade_pyspec(tmp_path, data_dir, config_gin, dataspec_bias):

    K.clear_session()

    gin.clear_config()

    bpnet_train(dataspec=str(dataspec_bias),

                output_dir=str(tmp_path),

                premade='bpnet9-pyspec',

                config=str(config_gin),

                num_workers=2

                )

def test_trained_model_vmtouch(tmp_path, data_dir, config_gin, dataspec_bias):

    K.clear_session()

    gin.clear_config()

    bpnet_train(dataspec=str(dataspec_bias),

                output_dir=str(tmp_path),

                premade='bpnet9',

                config=str(config_gin),

                vmtouch=True,

                num_workers=1

                )