import unittest
import warnings
import numpy as np
from sklearn.preprocessing import LabelBinarizer
import tensorflow as tf
import tensorflow.keras as keras

from useckit.v2 import model_descriptions as md


# Copied intelligent noise generator (do not import from other modules)
def make_some_intelligent_noise(labels: int = 4, shape: tuple = (8, 100, 100, 4), noisiness: float = 0.5):
    biases = np.array([i / labels for i in range(labels)]) + (1 / (2 * labels))
    x = []
    y = []
    shape = list(shape)
    shape[0] = int(shape[0] / labels)

    for i, b in enumerate(biases):
        x.extend(np.random.normal(b, noisiness, shape))
        y.extend(np.ones((shape[0],)) * i)

    x = np.array(x)
    y = [f'User {int(y_elem)}' for y_elem in y]
    y = np.array(y)
    x = np.clip(x, -1, 1)
    return x, y


def _gen_data_single_input(n=48, t=24, v=2, num_classes=3, seed=42, noisiness=0.2):
    """Generate synthetic single-input TSC data using intelligent noise.

    Parameters
    - n: number of samples (e.g., 48)
    - t: number of time steps per sample (sequence length, e.g., 24)
    - v: number of variables/channels per time step (features, e.g., 2)
    - num_classes: number of distinct class labels
    - seed: RNG seed for reproducibility
    - noisiness: standard deviation of the Gaussian noise around class-specific biases
    """
    np.random.seed(seed)
    x, y = make_some_intelligent_noise(labels=num_classes, shape=(n, t, v), noisiness=noisiness)
    lb = LabelBinarizer()
    y_ohc = lb.fit_transform(y)
    return x.astype(np.float32), y, y_ohc.astype(np.float32)


def _gen_data_multi_input(n=48, t=24, v=2, num_classes=3, seed=42, noisiness=0.2):
    """Generate synthetic multi-input TSC data (list of V inputs).

    Parameters
    - n: number of samples (e.g., 48)
    - t: number of time steps per sample (sequence length, e.g., 24)
    - v: number of variables/channels per time step (features, e.g., 2)
    - num_classes: number of distinct class labels
    - seed: RNG seed for reproducibility
    - noisiness: standard deviation of the Gaussian noise around class-specific biases
    """
    x, y, y_ohc = _gen_data_single_input(n=n, t=t, v=v, num_classes=num_classes, seed=seed, noisiness=noisiness)
    x_list = [x[:, :, i:i+1] for i in range(v)]
    return x_list, y, y_ohc


def _assert_acc_ok(testcase: unittest.TestCase, model, x, y_ohc):
    results = model.evaluate(x, y_ohc, verbose=0)
    names = model.metrics_names
    # find first accuracy-like metric
    acc_idx = None
    for i, n in enumerate(names):
        if 'acc' in n:
            acc_idx = i
            break
    testcase.assertIsNotNone(acc_idx, f"No accuracy metric found in {names}")
    acc = float(results[acc_idx])
    testcase.assertTrue(np.isfinite(acc))
    testcase.assertGreaterEqual(acc, 0.35)
    testcase.assertLessEqual(acc, 1.0)


class TestModelDescriptionsV2(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        warnings.simplefilter("ignore")
        if tf is not None:
            tf.keras.backend.clear_session()
            tf.random.set_seed(123)
        np.random.seed(123)

        cls.n = 48
        cls.t = 24
        cls.v = 2
        cls.num_classes = 3

    def _train_single_input(self, builder_fn):
        x, y, y_ohc = _gen_data_single_input(self.n, self.t, self.v, self.num_classes, seed=1)
        input_layer = keras.layers.Input((self.t, self.v))
        model, _ = builder_fn(input_layer, self.num_classes)
        hist = model.fit(x, y_ohc, epochs=10, batch_size=8, verbose=0, validation_split=0.2)
        self.assertTrue(len(hist.history) > 0)
        _assert_acc_ok(self, model, x, y_ohc)

    def _train_multi_input(self, builder_fn):
        v = self.v
        x_list, y, y_ohc = _gen_data_multi_input(self.n, self.t, v, self.num_classes, seed=2)
        # The description functions derive V from the provided input shape
        input_layer = keras.layers.Input((self.t, v))
        model, _ = builder_fn(input_layer, self.num_classes)
        hist = model.fit(x_list, y_ohc, epochs=10, batch_size=8, verbose=0, validation_split=0.2)
        self.assertTrue(len(hist.history) > 0)
        _assert_acc_ok(self, model, x_list, y_ohc)

    def test_cnn_valid(self):
        self._train_single_input(md.dl4tsc_cnn_padding_valid)

    def test_cnn_same(self):
        self._train_single_input(md.dl4tsc_cnn_padding_same)

    def test_fcn(self):
        self._train_single_input(md.dl4tsc_fcn)

    def test_mlp(self):
        self._train_single_input(md.dl4tsc_mlp)

    def test_resnet(self):
        self._train_single_input(md.dl4tsc_resnet)

    def test_encoder(self):
        self._train_single_input(md.dl4tsc_encoder)

    def test_inception(self):
        self._train_single_input(md.dl4tsc_inception)

    def test_tlenet(self):
        self._train_single_input(md.dl4tsc_tlenet)

    def test_mcdcnn(self):
        self._train_multi_input(md.dl4tsc_mcdcnn)

    def test_mcnn(self):
        self._train_multi_input(md.dl4tsc_mcnn)

    @unittest.skip("TWIESN is non-Keras and not epoch-based; requires legacy Dataset pipeline")
    def test_twiesn(self):  # documentation-only placeholder
        pass


if __name__ == '__main__':
    unittest.main()
