"""
Übung 2 — k-Prototypes: Stock Data
CAS PML, Unsupervised Learning (Exercise 2 aus dem Skript)

Aufgabe (Skript):
  1. Datensatz (tab-separiert) laden
  2. Daten clustern

Warum k-Prototypes? Die Stock-Daten sind gemischt (numerisch + kategorial).
  - k-Means  -> nur numerisch
  - k-Modes  -> nur kategorial (Modus statt Mittel, Mismatch-Distanz)
  - k-Prototypes = k-Means + k-Modes -> gemischte Attribute
"""

import matplotlib

from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from kmodes.kprototypes import KPrototypes
from sklearn.preprocessing import RobustScaler

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

matplotlib.use("Agg")  # headless: Plots nur in Dateien
BASE = Path(__file__).resolve().parent.parent  # .../aufgaben
DATA = BASE / "data" / "stocks2.csv"
PLOTS = BASE / "plots"
PLOTS.mkdir(exist_ok=True)
RANDOM_SEED = 0


# =========================================================
# 1. DATEN LADEN & VERSTEHEN
# =========================================================

# Daten aus dem Datensatz laden
df = pd.read_csv(DATA)

# Überblich verschaffen, vor dem Clustering
print(df.head())
print(df.shape)
print(df.dtypes)
print(df.nunique())
print(df.describe(include="all"))

# Frage: Welche Spalten sind numerisch, welche kategorial?
# - Numerisch: Dividend, PE, CAP
# - Kategorisch: Sector
# Frage: Gibt es eine ID-/Name-Spalte, die NICHT ins Clustering gehört?
# - ID: Symbol, Stock -> Reine Identifikation
# Entscheidung:
# - Wir verwenden also Dividend, PE, CAP und Sector fürs Clustering, der Rest fliegt raus!

# =========================================================
# 2. FEATURES VORBEREITEN
# =========================================================
# k-Prototypes braucht:
#   - ein gemischtes Array X (numerische Spalten als Zahl, kategoriale als String)
#   - die INDIZES (0-basiert, bezogen auf X) der kategorialen Spalten

# features auswählen, welche wir oben selektiert haben
feature_cols = ["Dividend", "PE", "Cap", "Sector"]
# und nach numpy laden
X = df[feature_cols].to_numpy()
# Sector ist die 4. Spalte in feature_cols -> Index 3 (0-basiert)
categorical_idx = [3]
# Ausgeben ob wir NaN im Datensatz haben, die wir erst bereinigen müssten
print(df[feature_cols].isna().sum())
# Überblick über die Daten verschaffen
print(df[["Dividend", "PE", "Cap"]].describe())

# Das Describe oben liefert zutage: CAP hat Werte von 13.9 bis 592771
# Und ist damit viel grösser als DIV und PE. Das bedeutet wir müssen CAP entsprechend
# skalieren, damit wir nicht nur auf die Marktkapitalisierung clustern.
# - Entscheidung: Deshalb: numerische Spalten skalieren, damit jede ungefähr gleich viel beiträgt.

num_cols = ["Dividend", "PE", "Cap"]
cat_cols = ["Sector"]

# RobustScaler -> gegen die PE/Cap-Ausreisser, spreizt die Hauptmasse besser
scaler = RobustScaler()
num_scaled = scaler.fit_transform(df[num_cols])  # -> numerisch, standardisiert

# X zusammenbauen: skalierte numerische Spalten + kategoriale Spalte
X = np.column_stack([num_scaled, df[cat_cols].to_numpy()])

# Reihenfolge in X: [Dividend, PE, Cap, Sector] -> Sector ist Index 3
categorical_idx = [3]

# Ausgeben was der Scaler genau macht
print("=== VOR Skalierung ===")
print(df[num_cols].describe().loc[["mean", "std", "min", "max"]])
print("\n=== NACH Skalierung ===")
num_scaled_df = pd.DataFrame(num_scaled, columns=num_cols)
print(num_scaled_df.describe().loc[["mean", "std", "min", "max"]])

# Plotten was der Scaler genau macht
fig, axes = plt.subplots(1, 2, figsize=(12, 5))
df[num_cols].boxplot(ax=axes[0])
axes[0].set_title("Vor Skalierung")
axes[0].set_yscale("log")  # sonst erdrückt Cap alles
num_scaled_df.boxplot(ax=axes[1])
axes[1].set_title("Nach StandardScaler")
fig.tight_layout()
fig.savefig(PLOTS / "u2_skalierung_vergleich.png", dpi=150, bbox_inches="tight")
plt.close(fig)

# =========================================================
# 3. PASSENDES k FINDEN (Elbow über cost_)
# =========================================================
# KPrototypes hat nach dem Fit das Attribut .cost_ (analog zur Inertia bei k-Means).

# Elbow über cost_
k_values = range(2, 10)
costs = []

GAMMA = 0.5  # fixes Gewicht kategorial vs. numerisch -> cost_ über k vergleichbar

#for k in k_values:
#    kp = KPrototypes(
#        n_clusters=k, init="Cao", random_state=RANDOM_SEED, n_init=5, gamma=GAMMA
#    )
#    kp.fit_predict(X, categorical=categorical_idx)
#    costs.append(kp.cost_)
#    print(f"k={k}: cost={kp.cost_:.1f}")

#fig, ax = plt.subplots()
#ax.plot(list(k_values), costs, marker="o")
#ax.set_xlabel("Cluster Count k")
#ax.set_ylabel("cost_ (Unähnlichkeit)")
#ax.set_title("Elbow – k-Prototypes")
#ax.grid(True)
#fig.savefig(PLOTS / "u2_elbow_kprototypes.png", dpi=150, bbox_inches="tight")
#plt.close(fig)
#print("Plot gespeichert: plots/u2_elbow_kprototypes.png")

# Ergebnis: Elbow: stärkster Steigungsabfall nach k=4

# =========================================================
# 4. FINALES CLUSTERING
# =========================================================

best_k = 4  # Elbow: stärkster Steigungsabfall nach k=4

kproto = KPrototypes(
    n_clusters=best_k, init="Cao", random_state=RANDOM_SEED, n_init=5, gamma=GAMMA
)
clusters = kproto.fit_predict(X, categorical=categorical_idx)

print("Cluster-Grössen:", np.bincount(clusters))
print("\nZentren (Dividend, PE, Cap skaliert | Sector):")
print(kproto.cluster_centroids_)

# =========================================================
# 5. INTERPRETATION
# =========================================================

df["cluster"] = clusters

# 5a) Numerische Zentren zurück in Originaleinheiten
centroids = kproto.cluster_centroids_
num_centers_scaled = centroids[:, :len(num_cols)].astype(float)   # erste 3 Spalten
num_centers = scaler.inverse_transform(num_centers_scaled)        # zurückrechnen
sector_modes = centroids[:, len(num_cols)]                        # 4. Spalte = Sector-Modus

profile = pd.DataFrame(num_centers, columns=num_cols)
profile["Sector (Modus)"] = sector_modes
profile["Grösse"] = np.bincount(clusters)
print("=== Cluster-Profile (Originaleinheiten) ===")
print(profile.round(2))

# 5b) Plausibilitätscheck: direkt aus den Daten je Cluster mitteln
print("\n=== Mittelwerte je Cluster (direkt aus df) ===")
print(df.groupby("cluster")[num_cols].mean().round(2))

# 5c) Cluster visualisieren (Cap vs PE, log-log) logarithmisch, da viele Ausreisser
fig, ax = plt.subplots(figsize=(8, 6))
ax.scatter(df["Cap"], df["PE"], c=df["cluster"], cmap="tab10", s=12, alpha=0.6)
ax.scatter(num_centers[:, num_cols.index("Cap")],
           num_centers[:, num_cols.index("PE")],
           c="k", marker="X", s=200, label="Zentren")
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xlabel("Cap [log]")
ax.set_ylabel("PE [log]")
ax.set_title(f"k-Prototypes Cluster (k={best_k})")
ax.legend()
fig.savefig(PLOTS / "u2_cluster_scatter.png", dpi=150, bbox_inches="tight")
plt.close(fig)
print("Plot gespeichert: plots/u2_cluster_scatter.png")
