import json
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score, confusion_matrix, f1_score,
    precision_score, recall_score, roc_auc_score
)
from sklearn.model_selection import StratifiedKFold, cross_val_score, train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

SEED = 4172
TASK = "split"
THRESHOLD = 0.50
COMPARE_LEAKAGE = False
CLASS_WEIGHT = None

X, y = make_classification(
    n_samples=480,
    n_features=6,
    n_informative=4,
    n_redundant=1,
    weights=[0.88, 0.12],
    class_sep=1.15,
    flip_y=0.025,
    random_state=SEED,
)

X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.25,
    stratify=y,
    random_state=SEED,
)

pipeline = Pipeline([
    ("scale", StandardScaler()),
    ("model", LogisticRegression(
        class_weight=CLASS_WEIGHT,
        max_iter=1000,
        random_state=SEED,
    )),
])

safe_cv = None
leaky_cv = None
if COMPARE_LEAKAGE:
    folds = StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED)
    safe_cv = float(cross_val_score(
        pipeline, X_train, y_train, cv=folds, scoring="roc_auc"
    ).mean())

    # Bevisst anti-eksempel: scaleren ser alle treningsfoldene før CV.
    globally_scaled = StandardScaler().fit_transform(X_train)
    leaky_cv = float(cross_val_score(
        LogisticRegression(max_iter=1000, random_state=SEED),
        globally_scaled,
        y_train,
        cv=folds,
        scoring="roc_auc",
    ).mean())

pipeline.fit(X_train, y_train)
probability = pipeline.predict_proba(X_test)[:, 1]
prediction = (probability >= THRESHOLD).astype(int)
tn, fp, fn, tp = confusion_matrix(y_test, prediction).ravel()

summary = {
    "task": TASK,
    "trainRows": int(len(y_train)),
    "testRows": int(len(y_test)),
    "trainPositiveRate": float(np.mean(y_train)),
    "testPositiveRate": float(np.mean(y_test)),
    "threshold": float(THRESHOLD),
    "accuracy": float(accuracy_score(y_test, prediction)),
    "precision": float(precision_score(y_test, prediction, zero_division=0)),
    "recall": float(recall_score(y_test, prediction, zero_division=0)),
    "f1": float(f1_score(y_test, prediction, zero_division=0)),
    "rocAuc": float(roc_auc_score(y_test, probability)),
    "falsePositive": int(fp),
    "falseNegative": int(fn),
}
if safe_cv is not None:
    summary["safeCv"] = safe_cv
    summary["leakyCv"] = leaky_cv

print("ML_RESULT=" + json.dumps(summary, sort_keys=True))

matrix = np.array([[tn, fp], [fn, tp]])
fig, ax = plt.subplots(figsize=(5.4, 3.8))
image = ax.imshow(matrix, cmap="viridis")
for row in range(2):
    for col in range(2):
        ax.text(col, row, str(matrix[row, col]), ha="center", va="center",
                color="white" if matrix[row, col] > matrix.max() / 2 else "black",
                fontsize=13, fontweight="bold")
ax.set_xticks([0, 1], ["Predikert 0", "Predikert 1"])
ax.set_yticks([0, 1], ["Faktisk 0", "Faktisk 1"])
ax.set_title(f"Hold-out ved terskel {THRESHOLD:.2f}")
fig.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
fig.tight_layout()