import numpy as np

from ai_toolbox import (
    Callback,
    CNNClassifier,
    CNNRegressor,
    MLPClassifier,
    MLPRegressor,
    RNNClassifier,
    RNNRegressor,
    pad_sequences,
    sequence_collate_fn,
)


def test_mlp_smoke():
    X = np.random.randn(64, 10).astype("float32")
    y = (X[:, 0] > 0).astype("int64")
    model = MLPClassifier(hidden_layers=(16,), epochs=1, batch_size=16)
    model.fit(X, y, verbose=False)
    out = model.predict(X[:5])
    assert out.shape == (5,)


def test_cnn_smoke():
    X = np.random.randn(32, 1, 16, 16).astype("float32")
    y = np.random.randint(0, 3, size=32)
    model = CNNClassifier(conv_channels=(8, 16), classifier_hidden=16, epochs=1, batch_size=8)
    model.fit(X, y, verbose=False)
    probs = model.predict_proba(X[:4])
    assert probs.shape == (4, 3)


def test_rnn_smoke():
    X = np.random.randn(40, 12, 5).astype("float32")
    y = np.random.randint(0, 4, size=40)
    model = RNNClassifier(hidden_size=8, rnn_type="gru", epochs=1, batch_size=10)
    model.fit(X, y, verbose=False)
    metrics = model.evaluate(X, y)
    assert "accuracy" in metrics and "loss" in metrics


def test_mlp_save_load_roundtrip(tmp_path):
    X = np.random.randn(48, 6).astype("float32")
    y = np.random.randint(0, 3, size=48)

    model = MLPClassifier(hidden_layers=(12,), epochs=1, batch_size=8, random_state=7)
    model.fit(X, y, verbose=False)

    path = tmp_path / "mlp_roundtrip.pt"
    model.save(path)

    loaded = MLPClassifier.load(path)

    probs_a = model.predict_proba(X[:10])
    probs_b = loaded.predict_proba(X[:10])
    pred_a = model.predict(X[:10])
    pred_b = loaded.predict(X[:10])

    np.testing.assert_allclose(probs_a, probs_b, rtol=1e-6, atol=1e-6)
    np.testing.assert_array_equal(pred_a, pred_b)


def test_mlp_early_stopping_checkpoint(tmp_path):
    rng = np.random.default_rng(123)
    X = rng.normal(size=(120, 8)).astype("float32")
    y = ((X[:, 0] + 0.3 * X[:, 1]) > 0).astype("int64")

    X_train, y_train = X[:84], y[:84]
    X_val, y_val = X[84:], y[84:]

    ckpt_path = tmp_path / "best_mlp_checkpoint.pt"
    model = MLPClassifier(hidden_layers=(16,), epochs=12, batch_size=16, random_state=123)
    model.fit(
        X_train,
        y_train,
        val_data=(X_val, y_val),
        early_stopping=True,
        patience=3,
        checkpoint_path=ckpt_path,
        verbose=False,
    )

    assert ckpt_path.exists()
    assert model.monitor_ == "val_loss"
    assert model.best_epoch_ is not None
    assert model.best_score_ is not None
    assert "val_loss" in model.history_
    assert len(model.history_["loss"]) <= 12

    loaded = MLPClassifier.load(ckpt_path)
    pred = loaded.predict(X_val[:5])
    assert pred.shape == (5,)


def test_mlp_resume_from_checkpoint(tmp_path):
    rng = np.random.default_rng(222)
    X = rng.normal(size=(96, 8)).astype("float32")
    y = ((X[:, 0] - 0.2 * X[:, 1]) > 0).astype("int64")

    ckpt_path = tmp_path / "resume_mlp.pt"

    model_a = MLPClassifier(hidden_layers=(16,), epochs=1, batch_size=16, random_state=222)
    model_a.fit(X, y, checkpoint_path=ckpt_path, verbose=False)
    assert ckpt_path.exists()

    model_b = MLPClassifier(hidden_layers=(16,), epochs=1, batch_size=16, random_state=222)
    model_b.fit(X, y, resume_from_checkpoint=ckpt_path, checkpoint_path=ckpt_path, verbose=False)

    assert len(model_b.history_["loss"]) >= 2
    assert "f1" in model_b.evaluate(X, y)


def test_mlp_classifier_callbacks_scheduler_and_weighting_smoke():
    class StopAfterOneEpoch(Callback):
        def __init__(self) -> None:
            super().__init__()
            self.epochs = []

        def on_epoch_end(self, epoch, logs=None):
            self.epochs.append((epoch, dict(logs or {})))
            self.stop_training = True

    rng = np.random.default_rng(223)
    X = rng.normal(size=(80, 6)).astype("float32")
    y = (X[:, 0] > 0).astype("int64")
    sample_weight = np.where(y == 1, 2.0, 1.0).astype("float32")
    cb = StopAfterOneEpoch()

    model = MLPClassifier(hidden_layers=(12,), epochs=3, batch_size=16, random_state=223)
    model.fit(
        X,
        y,
        class_weight={0: 1.0, 1: 1.5},
        sample_weight=sample_weight,
        scheduler="step",
        scheduler_kwargs={"step_size": 1, "gamma": 0.8},
        callbacks=[cb],
        verbose=False,
    )

    assert model.early_stopped_ is True
    assert model.stopped_epoch_ == 1
    assert len(model.history_["loss"]) == 1
    assert len(cb.epochs) == 1
    assert "lr" in cb.epochs[0][1]


def test_mlp_regressor_smoke():
    rng = np.random.default_rng(11)
    X = rng.normal(size=(64, 10)).astype("float32")
    y = (1.2 * X[:, 0] - 0.7 * X[:, 1] + 0.1 * rng.normal(size=64)).astype("float32")

    model = MLPRegressor(hidden_layers=(16,), epochs=1, batch_size=16)
    model.fit(X, y, verbose=False)
    pred = model.predict(X[:5])
    metrics = model.evaluate(X, y)

    assert pred.shape == (5,)
    assert "mse" in metrics and "mae" in metrics and "rmse" in metrics and "r2" in metrics and "loss" in metrics


def test_mlp_regressor_scheduler_and_weighting_smoke():
    rng = np.random.default_rng(224)
    X = rng.normal(size=(96, 5)).astype("float32")
    y = (0.7 * X[:, 0] - 0.3 * X[:, 1] + 0.1 * rng.normal(size=96)).astype("float32")
    sample_weight = np.linspace(0.5, 1.5, num=96, dtype=np.float32)

    model = MLPRegressor(hidden_layers=(10,), epochs=2, batch_size=16, random_state=224)
    model.fit(
        X[:72],
        y[:72],
        sample_weight=sample_weight[:72],
        val_data=(X[72:], y[72:]),
        scheduler="plateau",
        scheduler_kwargs={"factor": 0.5, "patience": 1},
        scheduler_monitor="val_loss",
        verbose=False,
    )
    metrics = model.evaluate(X, y)
    assert "rmse" in metrics and "r2" in metrics


def test_cnn_regressor_smoke():
    rng = np.random.default_rng(12)
    X = rng.normal(size=(24, 1, 16, 16)).astype("float32")
    y = X.mean(axis=(1, 2, 3)).astype("float32")

    model = CNNRegressor(conv_channels=(8, 16), classifier_hidden=16, epochs=1, batch_size=8)
    model.fit(X, y, verbose=False)
    pred = model.predict(X[:4])
    metrics = model.evaluate(X, y)

    assert pred.shape == (4,)
    assert "mse" in metrics and "mae" in metrics and "rmse" in metrics and "r2" in metrics


def test_rnn_regressor_smoke():
    rng = np.random.default_rng(13)
    X = rng.normal(size=(40, 12, 5)).astype("float32")
    y = (0.6 * X[:, -1, 0] - 0.2 * X[:, :, 1].mean(axis=1)).astype("float32")

    model = RNNRegressor(hidden_size=8, rnn_type="lstm", epochs=1, batch_size=10)
    model.fit(X, y, verbose=False)
    pred = model.predict(X[:6])
    metrics = model.evaluate(X, y)

    assert pred.shape == (6,)
    assert "mse" in metrics and "mae" in metrics and "rmse" in metrics and "r2" in metrics and "loss" in metrics


def test_rnn_lstm_variable_length_classifier_smoke():
    rng = np.random.default_rng(31)
    n_samples = 72
    max_len = 14
    n_features = 4

    lengths = rng.integers(4, max_len + 1, size=n_samples)
    X = np.zeros((n_samples, max_len, n_features), dtype="float32")
    y = rng.integers(0, 2, size=n_samples)

    for i, seq_len in enumerate(lengths):
        seq = rng.normal(0.0, 0.2, size=(seq_len, n_features)).astype("float32")
        seq[-1, 0] += 1.5 if y[i] == 1 else -1.5
        X[i, :seq_len] = seq
        # Misleading padding values to verify lengths are handled.
        X[i, seq_len:, 0] = -4.0 if y[i] == 1 else 4.0

    split = 54
    X_train, y_train, lengths_train = X[:split], y[:split], lengths[:split]
    X_val, y_val, lengths_val = X[split:], y[split:], lengths[split:]

    model = RNNClassifier(
        hidden_size=12,
        rnn_type="lstm",
        aggregate="mean",
        epochs=1,
        batch_size=12,
        random_state=31,
    )
    model.fit(
        X_train,
        y_train,
        lengths=lengths_train,
        val_data=(X_val, y_val, lengths_val),
        verbose=False,
    )

    probs = model.predict_proba(X_val[:6], lengths=lengths_val[:6])
    pred = model.predict(X_val[:6], lengths=lengths_val[:6])
    metrics = model.evaluate(X_val, y_val, lengths=lengths_val)

    assert probs.shape == (6, 2)
    assert pred.shape == (6,)
    assert "accuracy" in metrics and "precision" in metrics and "recall" in metrics and "f1" in metrics and "loss" in metrics


def test_rnn_lstm_variable_length_regressor_smoke():
    rng = np.random.default_rng(32)
    n_samples = 64
    max_len = 12
    n_features = 3

    lengths = rng.integers(3, max_len + 1, size=n_samples)
    X = np.zeros((n_samples, max_len, n_features), dtype="float32")
    y = np.zeros(n_samples, dtype="float32")

    for i, seq_len in enumerate(lengths):
        seq = rng.normal(0.0, 0.3, size=(seq_len, n_features)).astype("float32")
        X[i, :seq_len] = seq
        y[i] = float(seq[:, 0].mean() + 0.5 * seq[-1, 1])
        X[i, seq_len:, 0] = 10.0
        X[i, seq_len:, 1] = -10.0

    model = RNNRegressor(
        hidden_size=10,
        rnn_type="lstm",
        aggregate="last",
        epochs=1,
        batch_size=8,
        random_state=32,
    )
    model.fit(X, y, lengths=lengths, verbose=False)
    pred = model.predict(X[:7], lengths=lengths[:7])
    metrics = model.evaluate(X, y, lengths=lengths)

    assert pred.shape == (7,)
    assert "mse" in metrics and "mae" in metrics and "rmse" in metrics and "r2" in metrics and "loss" in metrics


def test_pad_sequences_smoke():
    seqs = [
        np.array([[1.0], [2.0]], dtype=np.float32),
        np.array([[3.0]], dtype=np.float32),
        np.array([[4.0], [5.0], [6.0]], dtype=np.float32),
    ]
    padded, lengths = pad_sequences(seqs, value=-1.0)
    assert padded.shape == (3, 3, 1)
    assert lengths.tolist() == [2, 1, 3]
    assert float(padded[1, 1, 0]) == -1.0


def test_sequence_collate_fn_smoke():
    batch = [
        (np.ones((3, 2), dtype=np.float32), 1),
        (np.ones((1, 2), dtype=np.float32) * 2, 0),
    ]
    x, y, lengths = sequence_collate_fn(batch, pad_value=0.0)
    assert tuple(x.shape) == (2, 3, 2)
    assert tuple(y.shape) == (2,)
    assert tuple(lengths.shape) == (2,)
    assert lengths.tolist() == [3, 1]


def test_mlp_regressor_save_load_roundtrip(tmp_path):
    rng = np.random.default_rng(14)
    X = rng.normal(size=(48, 6)).astype("float32")
    y = (X[:, 0] + 0.5 * X[:, 1]).astype("float32")

    model = MLPRegressor(hidden_layers=(12,), epochs=1, batch_size=8, random_state=14)
    model.fit(X, y, verbose=False)
    path = tmp_path / "mlp_reg_roundtrip.pt"
    model.save(path)

    loaded = MLPRegressor.load(path)
    pred_a = model.predict(X[:10])
    pred_b = loaded.predict(X[:10])
    np.testing.assert_allclose(pred_a, pred_b, rtol=1e-6, atol=1e-6)
