"""Authored synthetic tabular exercise; not an official AICE question or score.

Requires Python 3.10+, numpy, scikit-learn, tensorflow and matplotlib.
Creates a new aice-practice-output directory in the current working directory.
An existing output directory is never overwritten.
"""
from pathlib import Path
import csv
import json
import platform

import numpy as np
import sklearn
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import confusion_matrix, precision_recall_fscore_support
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import tensorflow as tf
import matplotlib
matplotlib.use("Agg")
from matplotlib import pyplot as plt


def main():
    output = Path("aice-practice-output")
    if output.exists():
        raise FileExistsError("aice-practice-output already exists; choose a fresh working folder")
    tf.keras.utils.set_random_seed(17)
    rng = np.random.default_rng(17)
    x = rng.normal(size=(300, 4)).astype("float32")
    y = (x[:, 0] + 0.5 * x[:, 1] > 0).astype("int32")
    ids = np.arange(len(y))
    train_val, test = train_test_split(ids, test_size=0.2, random_state=17, stratify=y)
    train, validation = train_test_split(
        train_val, test_size=0.25, random_state=17, stratify=y[train_val]
    )
    assert (len(train), len(validation), len(test)) == (180, 60, 60)
    assert not (set(train) & set(validation) or set(train_val) & set(test))

    scaler = StandardScaler().fit(x[train])
    x_train = scaler.transform(x[train]).astype("float32")
    x_val = scaler.transform(x[validation]).astype("float32")
    baseline = LogisticRegression(random_state=17).fit(x_train, y[train])
    model = tf.keras.Sequential([
        tf.keras.Input(shape=(x_train.shape[1],)),
        tf.keras.layers.Dense(8, activation="relu"),
        tf.keras.layers.Dense(1, activation="sigmoid"),
    ])
    model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
    stop = tf.keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=4, restore_best_weights=True
    )
    history = model.fit(
        x_train, y[train], validation_data=(x_val, y[validation]),
        epochs=30, batch_size=16, callbacks=[stop], verbose=0,
    )
    assert np.isfinite(history.history["val_loss"]).all()
    # Test data is evaluated only after fitting and stopping are finished.
    x_test = scaler.transform(x[test]).astype("float32")
    probabilities = model.predict(x_test, verbose=0)
    assert probabilities.shape == (60, 1)
    assert np.isfinite(probabilities).all()
    assert ((probabilities >= 0) & (probabilities <= 1)).all()
    predictions = (probabilities[:, 0] >= 0.5).astype("int32")
    baseline_predictions = baseline.predict(x_test)
    scores = {}
    for name, values in [("baseline", baseline_predictions), ("neural", predictions)]:
        p, r, f1, _ = precision_recall_fscore_support(
            y[test], values, average="binary", zero_division=0
        )
        matrix = confusion_matrix(y[test], values, labels=[0, 1])
        assert matrix.sum() == len(test)
        scores[name] = {"precision": float(p), "recall": float(r), "f1": float(f1)}

    # mkdir(exist_ok=False) also protects against a concurrent creator.
    output.mkdir(exist_ok=False)
    with (output / "predictions.csv").open("w", newline="", encoding="utf-8") as stream:
        writer = csv.writer(stream)
        writer.writerow(["id", "probability", "prediction"])
        writer.writerows(zip(test.tolist(), probabilities[:, 0].tolist(), predictions.tolist()))
    with (output / "predictions.csv").open(newline="", encoding="utf-8") as stream:
        rows = list(csv.DictReader(stream))
    assert [int(row["id"]) for row in rows] == test.tolist()
    assert len(rows) == len(test) and len({row["id"] for row in rows}) == len(test)

    fig, ax = plt.subplots()
    ax.plot(history.history["loss"], label="train")
    ax.plot(history.history["val_loss"], label="validation")
    ax.set(xlabel="Epoch (zero-based index)", ylabel="Binary crossentropy")
    ax.legend()
    fig.savefig(output / "learning-curve.png", bbox_inches="tight")
    plt.close(fig)
    model.save(output / "neural.keras")
    # Preserve the fitted preprocessing contract alongside the neural model.
    np.savez(output / "scaler.npz", mean=scaler.mean_, scale=scaler.scale_)
    versions = {"python": platform.python_version(), "numpy": np.__version__,
                "sklearn": sklearn.__version__, "tensorflow": tf.__version__,
                "matplotlib": matplotlib.__version__}
    report = {"versions": versions, "scores": scores,
              "splits": {"train": train.tolist(), "validation": validation.tolist(),
                         "test": test.tolist()}, "history": history.history,
              "note": "synthetic data; not a real-world or official exam score"}
    (output / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    print("PASS: split 180/60/60, probability shape (60, 1), finite outputs, CSV IDs")
    print(json.dumps({"versions": versions, "scores": scores}, indent=2))
    print("Synthetic data only; no production or certification readiness claim.")


if __name__ == "__main__":
    main()