"""Test sequence model
"""
from bpnet.seqmodel import SeqModel
from bpnet.heads import ScalarHead, BinaryClassificationHead, ProfileHead
import numpy as np
import keras.layers as kl


class TopDense:
    """Class to be used as functional model interpretation
    """

    def __init__(self, pool_size=2):
        self.pool_size = pool_size

    def __call__(self, inp):
        x = kl.GlobalAvgPool1D()(inp)
        return kl.Dense(1)(x)


class TopConv:
    """Class to be used as functional model interpretation
    """

    def __init__(self, n_output=2):
        self.n_output = n_output

    def __call__(self, inp):
        return kl.Conv1D(self.n_output, 1)(inp)


class BaseNet:
    """Class to be used as functional model interpretation
    """

    def __init__(self, activation='relu'):
        self.activation = activation

    def __call__(self, inp):
        x = kl.Conv1D(16, kernel_size=3, activation=self.activation, padding='same')(inp)
        return x


def test_interpret_wo_bias():
    from bpnet.metrics import RegressionMetrics, ClassificationMetrics, PeakPredictionProfileMetric
    from concise.preprocessing import encodeDNA
    # test the model
    seqs = encodeDNA(['ACAGA'] * 100)

    inputs = {"seq": seqs,
              "bias/a/profile": np.random.randn(100, 5, 2)}

    # Let's use regression
    targets = {"a/class": np.random.randint(low=0, high=2, size=(100, 1)).astype(float),
               "a/counts": 1 + np.ceil(np.abs(np.random.randn(100))),
               "a/profile": 1 + np.ceil(np.abs(np.random.randn(100, 5, 2))),
               }

    import keras.backend as K
    # K.clear_session()
    # use bias
    m = SeqModel(
        body=BaseNet('relu'),
        heads=[BinaryClassificationHead('{task}/class',
                                        net=TopDense(pool_size=2),
                                        use_bias=False),
               ScalarHead('{task}/counts',
                          loss='mse',
                          metric=RegressionMetrics(),
                          net=TopDense(pool_size=2),
                          use_bias=False),
               ProfileHead('{task}/profile',
                           loss='mse',
                           metric=PeakPredictionProfileMetric(neg_max_threshold=0.05,
                                                              required_min_pos_counts=0),
                           net=TopConv(n_output=2),
                           use_bias=True,
                           bias_shape=(5, 2)),  # NOTE: the shape currently has to be hard-coded to the sequence length
               ],
        tasks=['a']
    )
    m.model.fit(inputs, targets)

    o = m.contrib_score_all(seqs)
    assert 'a/profile/wn' in o
    assert o['a/profile/wn'].shape == seqs.shape
    assert 'a/profile/wn' in o
    assert o['a/profile/wn'].shape == seqs.shape

    # evaluate the dataset -> setup an array dataset (NumpyDataset) -> convert to
    from bpnet.data import NumpyDataset
    ds = NumpyDataset({"inputs": inputs, "targets": targets})
    o = m.evaluate(ds)
    assert 'avg/counts/mad' in o