"""Clasificación de dígitos manuscritos — de vectores de píxeles a un pipeline con MLP.

Datos: Optical Recognition of Handwritten Digits (UCI, CC BY 4.0), incluido en
scikit-learn como `load_digits` (1.797 imágenes de 8x8 píxeles, 10 clases).
Reproduce los resultados del notebook del proyecto y genera data/digits.json.

Uso:
    python analysis/digits_classification.py
"""
import json
import warnings
from pathlib import Path

import numpy as np
from sklearn.datasets import load_digits
from sklearn.decomposition import PCA
from sklearn.ensemble import RandomForestClassifier
from sklearn.exceptions import ConvergenceWarning
from sklearn.linear_model import LogisticRegression
from sklearn.manifold import TSNE
from sklearn.metrics import confusion_matrix
from sklearn.model_selection import GridSearchCV, cross_val_score, train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC

warnings.filterwarnings("ignore", category=ConvergenceWarning)

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "data" / "digits.json"
ARCHITECTURES = [(50,), (100,), (100, 100), (100, 100, 100), (200, 100, 50), (50, 50, 50, 50)]


def mlp(layers):
    return MLPClassifier(hidden_layer_sizes=layers, max_iter=500, random_state=42)


def n_params(model):
    return int(sum(w.size for w in model.coefs_) + sum(b.size for b in model.intercepts_))


def main():
    digits = load_digits()
    X, y = digits.images.reshape(len(digits.images), -1), digits.target

    # El test se aparta antes de cualquier transformación: los escaladores se ajustan solo con train
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.3, random_state=42, stratify=y)
    scaler = StandardScaler().fit(X_tr)
    X_tr_s, X_te_s = scaler.transform(X_tr), scaler.transform(X_te)

    # ---------- Reducción de dimensionalidad (visualización, sobre todo el dataset) ----------
    X_all_s = StandardScaler().fit_transform(X)
    cumvar = np.cumsum(PCA().fit(X_all_s).explained_variance_ratio_) * 100
    pca2 = PCA(n_components=2).fit(X_all_s)
    idx = np.random.default_rng(42).choice(len(X), 1000, replace=False)
    xy_pca = pca2.transform(X_all_s[idx])
    xy_tsne = TSNE(n_components=2, random_state=42, perplexity=30).fit_transform(X_all_s[idx])

    # ---------- Arquitecturas del MLP ----------
    archs = []
    for layers in ARCHITECTURES:
        m = mlp(layers).fit(X_tr_s, y_tr)
        archs.append({"k": str(layers).replace(",)", ")"), "test": round(m.score(X_te_s, y_te) * 100, 2),
                      "params": n_params(m)})

    # ---------- Modelo elegido: MLP (100,) ----------
    best = mlp((100,)).fit(X_tr_s, y_tr)
    pred = best.predict(X_te_s)
    cm = confusion_matrix(y_te, pred)
    errors = [{"px": X_te[i].astype(int).tolist(), "real": int(y_te[i]), "pred": int(pred[i])}
              for i in np.where(pred != y_te)[0]]

    # ---------- Pipeline completo: escalado + PCA + MLP, con CV y GridSearch solo sobre train ----------
    pipe = Pipeline([("scaler", StandardScaler()), ("pca", PCA(n_components=50)),
                     ("mlp", mlp((100, 100)))])
    cv = cross_val_score(pipe, X_tr, y_tr, cv=5, scoring="accuracy")
    grid = GridSearchCV(pipe, {"pca__n_components": [20, 40, 50],
                               "mlp__hidden_layer_sizes": [(50,), (100,), (100, 100)],
                               "mlp__alpha": [0.0001, 0.001, 0.01]},
                        cv=3, scoring="accuracy", n_jobs=-1).fit(X_tr, y_tr)
    pca_step = grid.best_estimator_.named_steps["pca"]

    # ---------- Comparativa final ----------
    models = {
        "SVM (RBF)": SVC(kernel="rbf", random_state=42),
        "Regresión logística": LogisticRegression(max_iter=1000, random_state=42),
        "MLP (100,)": mlp((100,)),
        "MLP (100, 100, 100)": mlp((100, 100, 100)),
        "KNN (k=5)": KNeighborsClassifier(n_neighbors=5),
        "Random Forest": RandomForestClassifier(n_estimators=100, random_state=42),
    }
    comparison = []
    for name, model in models.items():
        model.fit(X_tr_s, y_tr)
        comparison.append({"k": name, "train": round(model.score(X_tr_s, y_tr) * 100, 2),
                           "test": round(model.score(X_te_s, y_te) * 100, 2)})

    out = {
        "kpis": {
            "imagenes": int(len(X)), "train": int(len(y_tr)), "test": int(len(y_te)),
            "pipeline_cv": round(cv.mean() * 100, 1), "pipeline_cv_std": round(cv.std() * 100, 1),
            "grid_best": {k.split("__")[1]: (str(v) if isinstance(v, tuple) else v) for k, v in grid.best_params_.items()},
            "grid_cv": round(grid.best_score_ * 100, 1),
            "pipeline_test": round(grid.best_estimator_.score(X_te, y_te) * 100, 1),
            "pca_var": round(float(pca_step.explained_variance_ratio_.sum()) * 100, 1),
            "n_95": int(np.argmax(cumvar >= 95) + 1),
            "pca2_var": round(float(pca2.explained_variance_ratio_.sum()) * 100, 1),
            "errores": len(errors),
        },
        "ejemplos": [X[np.where(y == d)[0][0]].astype(int).tolist() for d in range(10)],
        "pca_cumvar": [round(v, 2) for v in cumvar],
        "proyeccion": {"label": y[idx].tolist(),
                       "pca": np.round(xy_pca, 2).tolist(), "tsne": np.round(xy_tsne, 2).tolist()},
        "arquitecturas": archs,
        "cm": cm.tolist(),
        "recall_digito": [round(cm[d, d] / cm[d].sum() * 100, 1) for d in range(10)],
        "errores": errors,
        "comparativa": comparison,
    }
    OUT.parent.mkdir(exist_ok=True)
    OUT.write_text(json.dumps(out, ensure_ascii=False, separators=(",", ":")))
    print(json.dumps(out["kpis"], ensure_ascii=False, indent=2))
    print("arquitecturas", archs)
    print("comparativa", comparison)
    print("recall", out["recall_digito"])
    print("errores", [(e["real"], e["pred"]) for e in errors])


if __name__ == "__main__":
    main()