"""
Übung 1 — k-Means: Delivery Fleet
CAS PML, Unsupervised Learning

Aufgabe (Skript):
  1. Datensatz aus tab-separiertem File laden
  2. Daten visualisieren, sinnvolle Clusterzahl k finden (visuell + systematisch)
  3. Daten clustern

Umbau aus dem Lösungs-Notebook (lektion_1.ipynb):
  - .as_matrix() -> .to_numpy()  (in aktuellem Pandas entfernt)
  - plt.show()   -> savefig nach plots/  (headless, kein Display in der devenv-Shell)
  - n_init=10, random_state=0 ergänzt (Reproduzierbarkeit, keine FutureWarnings)
"""

import matplotlib

matplotlib.use("Agg")  # headless: nur in Dateien rendern

from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from scipy.spatial.distance import cdist, pdist

# =========================================================
# EINSTELLUNGEN / PFADE
# =========================================================

BASE = Path(__file__).resolve().parent.parent  # .../aufgaben
DATA = BASE / "data" / "data_delivery_fleet.tsv"
PLOTS = BASE / "plots"
PLOTS.mkdir(exist_ok=True)

RANDOM_SEED = 0
RUN_HARTIGAN = False  # Legacy-Mass (gilt als veraltet), siehe unten


def _scatter_clusters(ax, data, model):
    """Punkte nach Cluster-Label einfärben, Zentren als schwarzes X."""
    ax.scatter(data[:, 0], data[:, 1], c=model.labels_, cmap="tab10", s=10)
    centers = model.cluster_centers_
    ax.scatter(centers[:, 0], centers[:, 1], c="k", marker="X", s=150)
    ax.set_xlabel("distance [mi]")
    ax.set_ylabel("speeding [%]")


def save_kmeans(data, model, outfile, title=None):
    fig, ax = plt.subplots()
    _scatter_clusters(ax, data, model)
    ax.set_title(title or f"k-Means (k={len(model.cluster_centers_)})")
    fig.savefig(PLOTS / outfile, dpi=150, bbox_inches="tight")
    plt.close(fig)


# =========================================================
# 1. DATEN LADEN & VORBEREITEN
# =========================================================

if not DATA.exists():
    raise SystemExit(
        f"Daten nicht gefunden: {DATA}\n"
        "Lege 'data_delivery_fleet.tsv' in den data/-Ordner."
    )

df = pd.read_csv(DATA, sep="\t", header=0)
print(df.head())

# Spalten:
#   Driver_ID         - ID der Fahrer:in
#   Distance_Feature  - mittlere gefahrene Meilen pro Tag
#   Speeding_Feature  - % der Zeit > Tempolimit (>5 mph)
data = df[["Distance_Feature", "Speeding_Feature"]].to_numpy()
print("shape:", data.shape)

# Rohdaten
fig, ax = plt.subplots()
ax.scatter(df.Distance_Feature, df.Speeding_Feature, marker=".")
ax.set_xlabel("distance [mi]")
ax.set_ylabel("speeding fraction [%]")
ax.grid(True)
fig.savefig(PLOTS / "01_rohdaten.png", dpi=150, bbox_inches="tight")
plt.close(fig)

# =========================================================
# 2. ERSTES CLUSTERING (frei gewähltes k)
# =========================================================

k_means = KMeans(n_clusters=4, n_init=10, random_state=RANDOM_SEED).fit(data)
save_kmeans(data, k_means, "02_kmeans_k4.png")

# Cluster-Zuordnung für k = 1 ... 8 als Übersichts-Grid (statt 8 Einzelplots)
fig, axes = plt.subplots(2, 4, figsize=(20, 8))
for k, ax in zip(range(1, 9), axes.ravel()):
    model = KMeans(n_clusters=k, n_init=10, random_state=RANDOM_SEED).fit(data)
    _scatter_clusters(ax, data, model)
    ax.set_title(f"k={k}")
fig.tight_layout()
fig.savefig(PLOTS / "03_kmeans_sweep_k1-8.png", dpi=150, bbox_inches="tight")
plt.close(fig)

# =========================================================
# 3. OPTIMALE CLUSTERZAHL k
# =========================================================

# ---- Within-Cluster Sum of Squares (Inertia) ----
Ks = range(1, 10)
inertia = [
    KMeans(n_clusters=k, n_init=10, random_state=RANDOM_SEED).fit(data).inertia_
    for k in Ks
]

fig, ax = plt.subplots()
ax.bar(Ks, inertia)
ax.set_title("Within-Cluster Sum of Squares")
ax.set_xlabel("Cluster Count k")
ax.set_ylabel("Error")
ax.grid(True)
fig.savefig(PLOTS / "04_inertia.png", dpi=150, bbox_inches="tight")
plt.close(fig)

# ---- Percentage of Variance Explained ----
# TSS = WCSS + BSS  ->  Variance Explained = BSS / TSS
ks = range(1, 10)
models = [
    KMeans(n_clusters=k, n_init=10, random_state=RANDOM_SEED).fit(data) for k in ks
]
centroids = [m.cluster_centers_ for m in models]

# Abstand jedes Punkts zum nächsten Zentrum -> WCSS pro k
dist = [np.min(cdist(data, cent, "euclidean"), axis=1) for cent in centroids]
wcss = np.array([np.sum(d**2) for d in dist])
tss = np.sum(pdist(data) ** 2) / data.shape[0]
bss = tss - wcss

fig, ax = plt.subplots()
ax.plot(list(ks), bss / tss * 100, "b*-")
ax.set_ylim((0, 100))
ax.set_xlabel("Cluster Count k")
ax.set_ylabel("Percentage of variance explained")
ax.set_title("Variance Explained vs. k")
ax.grid(True)
fig.savefig(PLOTS / "05_variance_explained.png", dpi=150, bbox_inches="tight")
plt.close(fig)

# ---- Silhouette Score (erst ab k >= 2 definiert) ----
ks_sil = range(2, 10)
silhouette = [
    silhouette_score(
        data,
        KMeans(n_clusters=k, n_init=10, random_state=RANDOM_SEED).fit(data).labels_,
        metric="euclidean",
    )
    for k in ks_sil
]

fig, ax = plt.subplots()
ax.plot(list(ks_sil), silhouette, marker="o")
ax.set_xlabel("Cluster Count k")
ax.set_ylabel("Silhouette Score")
ax.set_title("Silhouette for k-Means")
ax.grid(True)
fig.savefig(PLOTS / "06_silhouette.png", dpi=150, bbox_inches="tight")
plt.close(fig)

best_k = list(ks_sil)[int(np.argmax(silhouette))]
print(f"Bestes k nach Silhouette: {best_k}")

# =========================================================
# 4. FINALES MODELL
# =========================================================

k_means_final = KMeans(n_clusters=best_k, n_init=10, random_state=RANDOM_SEED).fit(data)
save_kmeans(
    data, k_means_final, "07_final.png", title=f"Finales Clustering (k={best_k})"
)

print(f"Plots gespeichert in: {PLOTS}")


# =========================================================
# 5. HARTIGAN-REGEL (optional, Legacy)
# =========================================================
# Im Lösungs-Notebook enthalten, gilt aber als veraltet (vgl. Chiang & Mirkin).
# Silhouette/Elbow werden bevorzugt -> standardmässig deaktiviert (RUN_HARTIGAN).


def hartigan_k(data, threshold=12):
    """'Korrekte' Clusterzahl nach Hartigan: erstes k mit H <= threshold."""
    inertia_list = np.zeros(len(data) + 1)
    num = 0
    h_rule = threshold + 1
    while num < len(data) and h_rule > threshold:
        kmn = KMeans(n_clusters=num + 1, n_init=10, random_state=RANDOM_SEED).fit(data)
        inertia_list[num + 1] = kmn.inertia_
        if num > 0:
            h_rule = ((float(inertia_list[num]) / inertia_list[num + 1]) - 1) * (
                len(data) - num - 1
            )
            print(f"k={num}, H={h_rule:.2f}")
        num += 1
    if h_rule > threshold:
        num += 1
    return num - 1


if RUN_HARTIGAN:
    k_h = hartigan_k(data, threshold=12)
    print(f"Hartigan k: {k_h}")
    model_h = KMeans(n_clusters=k_h, n_init=10, random_state=RANDOM_SEED).fit(data)
    save_kmeans(data, model_h, "08_hartigan.png", title=f"Hartigan (k={k_h})")
