# Linear Discriminant Analysis (LDA)
# Reproduces every STEP6 block: Model Parameters, Metrics, Per-Class Metrics,
# Feature Importance, Cross-Validation, LDA Details (explained variance, class
# priors) and the full LDA Statistical Output (Wilks' Lambda, Box's M,
# Eigenvalues/canonical correlations, Structure matrix, ANOVA F, Group Centroids).

# ===== Data & Variables =====
import numpy as np, pandas as pd
from sklearn.datasets import load_iris
# data = pd.read_csv("data.csv"); target = "y"; features = ["x1","x2"]
d = load_iris(as_frame=True); data = d.frame
target = "target"; feature_cols = list(d.feature_names)
X_raw = data[feature_cols].copy()
y = data[target].astype(str).copy()

# Model parameters (match handler)
method = "lda"
solver = "svd"            # 'svd' | 'lsqr' | 'eigen'
shrinkage = None          # None | 'auto' | float (forces solver='lsqr')
n_components = None        # None -> min(n_features, n_classes-1)
test_size = 0.2
cv_folds = 5
random_state = 42

# ===== Analysis =====
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.pipeline import Pipeline
from sklearn.metrics import (accuracy_score, precision_score, recall_score,
                             f1_score, confusion_matrix, classification_report,
                             roc_auc_score)
from scipy import stats

le = LabelEncoder(); y_enc = le.fit_transform(y)
class_labels = [str(c) for c in le.classes_]
n_classes = len(le.classes_); n_features = X_raw.shape[1]

# handler resolves n_components and forces lsqr when shrinkage is set
max_components = min(n_features, n_classes - 1)
n_components_used = n_components if (n_components and n_components <= max_components) else (max_components or None)
solver_used = "lsqr" if (shrinkage is not None and solver == "svd") else solver

Xtr_raw, Xte_raw, ytr, yte = train_test_split(
    X_raw, y_enc, test_size=test_size, random_state=random_state, stratify=y_enc)
scaler = StandardScaler()
Xtr = scaler.fit_transform(Xtr_raw); Xte = scaler.transform(Xte_raw)

model = LinearDiscriminantAnalysis(solver=solver_used, shrinkage=shrinkage,
                                   n_components=n_components_used)
model.fit(Xtr, ytr)
pred = model.predict(Xte); proba = model.predict_proba(Xte)

# ===== Results: Model Parameters =====
print("=== Model Parameters ===")
print("method:", method.upper(), " solver:", solver_used,
      " shrinkage:", shrinkage, " n_components:", n_components_used,
      " test_size:", test_size, " cv_folds:", cv_folds)

# ===== Results: Metrics =====
print("\n=== Metrics ===")
print("accuracy:", accuracy_score(yte, pred),
      " train_accuracy:", accuracy_score(ytr, model.predict(Xtr)))
print("precision_macro:", precision_score(yte, pred, average="macro", zero_division=0),
      " recall_macro:", recall_score(yte, pred, average="macro", zero_division=0),
      " f1_macro:", f1_score(yte, pred, average="macro", zero_division=0))
try:
    if n_classes == 2:
        print("auc:", roc_auc_score(yte, proba[:, 1]))
    else:
        print("auc (ovr macro):", roc_auc_score(yte, proba, multi_class="ovr", average="macro"))
except Exception as e:
    print("auc:", e)

# ===== Results: Per-Class Metrics =====
rep = classification_report(yte, pred, target_names=class_labels, output_dict=True)
print("\n=== Per-Class Metrics ===")
for cl in class_labels:
    r = rep[cl]
    print(f"  {cl}: precision={r['precision']:.4f} recall={r['recall']:.4f} "
          f"f1={r['f1-score']:.4f} support={int(r['support'])}")
print("Confusion matrix:\n", confusion_matrix(yte, pred))

# ===== Results: LDA Details =====
print("\n=== LDA Details ===")
print("explained_variance_ratio:", [float(v) for v in model.explained_variance_ratio_])
print("class_priors:", [float(v) for v in model.priors_])

# ===== Results: Feature Importance =====
# |coef| * feature-std, normalized (matches handler's importance scoring)
std = Xtr.std(axis=0); std[std == 0] = 1e-10
imp = np.abs(model.coef_).mean(axis=0) * std
imp = imp / imp.sum() if imp.sum() > 0 else np.ones(n_features) / n_features
print("\n=== Feature Importance ===")
for name, v in sorted(zip(feature_cols, imp), key=lambda kv: -kv[1]):
    print(f"  {name}: importance={v:.4f} ({v*100:.1f}%)")

# ===== Results: Cross-Validation (pipeline avoids leakage) =====
min_class = int(np.min(np.bincount(y_enc)))
cvf = max(2, min(cv_folds, min_class))
pipe = Pipeline([("scaler", StandardScaler()),
                 ("model", LinearDiscriminantAnalysis(solver=solver_used,
                          shrinkage=shrinkage, n_components=n_components_used))])
cv = cross_val_score(pipe, X_raw.values, y_enc, cv=cvf, scoring="accuracy")
print("\n=== Cross-Validation Results ===")
print("cv_folds:", cvf, " cv_mean:", cv.mean(), " cv_std:", cv.std(),
      " scores:", [round(float(s), 4) for s in cv])

# ===== Results: Prediction Examples =====
# Sample rows from the test set (matches handler's generate_prediction_examples):
# actual label, predicted label, correct, confidence (max predicted probability).
print("\n=== Prediction Examples ===")
_rng = np.random.RandomState(random_state)
_n_ex = min(15, len(yte))
_ex_idx = _rng.choice(len(yte), size=_n_ex, replace=False)
for i in _ex_idx:
    actual_label = class_labels[yte[i]]
    predicted_label = class_labels[pred[i]]
    confidence = proba[i].max()
    print(f"  actual={actual_label} predicted={predicted_label} "
          f"correct={bool(yte[i] == pred[i])} confidence={confidence:.4f}")

# ===== Results: LDA Statistical Output (SPSS-style) =====
# Between/within scatter on standardized training data
overall_mean = Xtr.mean(axis=0)
S_B = np.zeros((n_features, n_features)); S_W = np.zeros((n_features, n_features))
group_means = {}
for i in range(n_classes):
    mask = ytr == i; nk = mask.sum(); mk = Xtr[mask].mean(axis=0)
    group_means[i] = mk
    diff = (mk - overall_mean).reshape(-1, 1)
    S_B += nk * (diff @ diff.T)
    centered = Xtr[mask] - mk
    S_W += centered.T @ centered

# Eigenvalues, canonical correlations, variance explained
S_W_inv = np.linalg.pinv(S_W)
eig_raw, _ = np.linalg.eig(S_W_inv @ S_B)
eig_raw = np.real(eig_raw)
eigenvalues = np.sort(eig_raw)[::-1]
eigenvalues = eigenvalues[eigenvalues > 1e-10][:n_classes - 1]
canon_corr = np.sqrt(eigenvalues / (1 + eigenvalues))
total_eig = eigenvalues.sum()
var_exp = (eigenvalues / total_eig * 100) if total_eig > 0 else eigenvalues * 0
print("\n=== LDA Statistical Output ===")
print("-- Eigenvalues --")
cum = 0.0
for i, (ev, ve, cc) in enumerate(zip(eigenvalues, var_exp, canon_corr)):
    cum += ve
    print(f"  LD{i+1}: eigenvalue={ev:.4f} variance_explained={ve:.1f}% "
          f"cumulative={cum:.1f}% canonical_correlation={cc:.4f}")

# Wilks' Lambda + chi-square test
wilks = float(np.prod([1 / (1 + ev) for ev in eigenvalues])) if len(eigenvalues) else 1.0
n_samples = Xtr.shape[0]; p = n_features; g = n_classes
chi2_stat = -(n_samples - 1 - (p + g) / 2) * np.log(wilks + 1e-10)
df = p * (g - 1)
p_val = float(1 - stats.chi2.cdf(chi2_stat, df))
print("-- Wilks' Lambda --")
print(f"  lambda={wilks:.4f} chi2={chi2_stat:.4f} df={df} p_value={p_val:.4g} "
      f"significant={p_val < 0.05}")

# Box's M test (equality of covariance matrices)
def box_m(X, labels):
    groups = [X[labels == i] for i in range(n_classes)]
    ni = np.array([len(gp) for gp in groups]); N = ni.sum()
    Si = [np.cov(gp, rowvar=False) for gp in groups]
    Sp = sum((ni[i] - 1) * Si[i] for i in range(n_classes)) / (N - n_classes)
    sign, logdet_Sp = np.linalg.slogdet(Sp)
    M = (N - n_classes) * logdet_Sp - sum((ni[i] - 1) * np.linalg.slogdet(Si[i])[1]
                                          for i in range(n_classes))
    c1 = ((2 * p**2 + 3 * p - 1) / (6 * (p + 1) * (n_classes - 1))) * \
         (sum(1 / (ni[i] - 1) for i in range(n_classes)) - 1 / (N - n_classes))
    chi2_box = M * (1 - c1)
    df_box = p * (p + 1) * (n_classes - 1) / 2
    return M, chi2_box, df_box, float(1 - stats.chi2.cdf(chi2_box, df_box))
try:
    M, chi2_box, df_box, p_box = box_m(Xtr, ytr)
    print("-- Box's M --")
    print(f"  M={M:.4f} chi2={chi2_box:.4f} df={df_box:.0f} p_value={p_box:.4g} "
          f"significant={p_box < 0.05}")
except Exception as e:
    print("-- Box's M --  (not computable:", e, ")")

# Structure matrix: correlation of each variable with each discriminant function
scalings = model.scalings_
print("-- Structure Matrix (variable vs discriminant function) --")
for j, name in enumerate(feature_cols):
    ldf_corr = []
    for c in range(scalings.shape[1]):
        scores = Xtr @ scalings[:, c]
        ldf_corr.append(np.corrcoef(Xtr[:, j], scores)[0, 1])
    print(f"  {name}: " + " ".join(f"LD{c+1}={v:.4f}" for c, v in enumerate(ldf_corr)))

# ANOVA F per variable (univariate group separation)
print("-- ANOVA F (per variable) --")
anova = []
for j, name in enumerate(feature_cols):
    grp = [Xtr[ytr == i, j] for i in range(n_classes)]
    f_stat, f_p = stats.f_oneway(*grp)
    anova.append((name, float(f_stat), float(f_p)))
for name, f_stat, f_p in sorted(anova, key=lambda t: -t[1]):
    print(f"  {name}: F={f_stat:.4f} p={f_p:.4g} significant={f_p < 0.05}")

# Group centroids in discriminant space
print("-- Group Centroids (discriminant space) --")
for i, label in enumerate(class_labels):
    cs = Xtr[ytr == i] @ scalings
    print(f"  {label}: " + " ".join(f"LD{c+1}={cs[:, c].mean():.4f}"
                                    for c in range(cs.shape[1])))

# ===== Visualizations =====
import os, tempfile
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from sklearn.metrics import ConfusionMatrixDisplay
_out = tempfile.gettempdir()

# Confusion matrix heatmap
fig, ax = plt.subplots()
ConfusionMatrixDisplay(confusion_matrix(yte, pred),
                       display_labels=class_labels).plot(ax=ax, cmap="Blues", colorbar=False)
ax.set_title("Confusion Matrix"); fig.tight_layout()
fig.savefig(os.path.join(_out, "confusion.png"))

# LDA projection (discriminant space), colored by class
Xall = np.vstack([Xtr, Xte]); yall = np.hstack([ytr, yte])
proj = model.transform(Xall)
fig, ax = plt.subplots(figsize=(9, 7))
for i, label in enumerate(class_labels):
    m = yall == i
    if proj.shape[1] >= 2:
        ax.scatter(proj[m, 0], proj[m, 1], label=label, alpha=0.7, s=45, edgecolors="white")
    else:
        ax.scatter(proj[m, 0], np.zeros(m.sum()), label=label, alpha=0.7, s=45)
ax.set_xlabel("LD1"); ax.set_ylabel("LD2" if proj.shape[1] >= 2 else "")
ax.set_title("LDA Projection"); ax.legend(); ax.grid(True, alpha=0.3)
fig.tight_layout(); fig.savefig(os.path.join(_out, "lda_projection.png"))

# Feature importance barplot
fig, ax = plt.subplots()
order = np.argsort(imp)
ax.barh([feature_cols[i] for i in order], imp[order], color="steelblue")
ax.set_xlabel("Importance"); ax.set_title("Feature Importance")
fig.tight_layout(); fig.savefig(os.path.join(_out, "importance.png"))

# ROC curves (one-vs-rest per class)
from sklearn.metrics import roc_curve, auc
from sklearn.preprocessing import label_binarize
fig, ax = plt.subplots()
yb = label_binarize(yte, classes=range(n_classes))
for i, label in enumerate(class_labels):
    fpr, tpr, _ = roc_curve(yb[:, i], proba[:, i])
    ax.plot(fpr, tpr, label=f"{label} (AUC={auc(fpr, tpr):.3f})")
ax.plot([0, 1], [0, 1], "k--", lw=1)
ax.set_xlabel("False Positive Rate"); ax.set_ylabel("True Positive Rate")
ax.set_title("ROC Curve (one-vs-rest)"); ax.legend(loc="lower right")
fig.tight_layout(); fig.savefig(os.path.join(_out, "roc.png"))

# Class separation: standardized class means per feature
fig, ax = plt.subplots(figsize=(10, 6))
xpos = np.arange(len(feature_cols)); width = 0.8 / n_classes
for i, label in enumerate(class_labels):
    means = model.means_[i]
    ax.bar(xpos + (i - n_classes / 2 + 0.5) * width, means, width, label=label)
ax.set_xticks(xpos); ax.set_xticklabels(feature_cols, rotation=45, ha="right")
ax.set_ylabel("Mean (standardized)"); ax.set_title("Class Means by Feature"); ax.legend()
fig.tight_layout(); fig.savefig(os.path.join(_out, "class_separation.png"))