import json
from pathlib import Path

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

FEATURES = [
    "latency_ms",
    "error_rate",
    "cpu_utilization",
    "memory_utilization",
    "queue_depth",
]
OUTPUT_DIR = Path(__file__).resolve().parent.parent / "models"


def export_model(version: str, seed: int, threshold: float) -> None:
    rng = np.random.default_rng(seed)
    healthy = rng.normal([170, 0.015, 0.38, 0.45, 5], [45, 0.01, 0.12, 0.12, 4], size=(160, 5))
    at_risk = rng.normal([420, 0.12, 0.82, 0.78, 38], [130, 0.08, 0.12, 0.13, 18], size=(160, 5))
    x = np.vstack([healthy, at_risk])
    y = np.array([0] * len(healthy) + [1] * len(at_risk))

    pipeline = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000, random_state=seed))
    pipeline.fit(x, y)

    scaler = pipeline.named_steps["standardscaler"]
    classifier = pipeline.named_steps["logisticregression"]
    artifact = {
        "version": version,
        "trained_with": "scikit-learn LogisticRegression",
        "threshold": threshold,
        "bias": float(classifier.intercept_[0]),
        "features": FEATURES,
        "weights": [float(value) for value in classifier.coef_[0]],
        "baseline": {
            feature: {"mean": float(mean), "stddev": float(stddev)}
            for feature, mean, stddev in zip(FEATURES, scaler.mean_, scaler.scale_, strict=True)
        },
    }

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    with (OUTPUT_DIR / f"model_{version}.json").open("w", encoding="utf-8") as handle:
        json.dump(artifact, handle, indent=2)
        handle.write("\n")


if __name__ == "__main__":
    export_model("v1", seed=42, threshold=0.55)
    export_model("v2", seed=84, threshold=0.5)