from collections import Counter
from pathlib import Path

import joblib
from sklearn.metrics import (
    accuracy_score,
    balanced_accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
)
from sklearn.model_selection import GridSearchCV, StratifiedKFold, cross_val_predict
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC

from features import LABEL_PD, LABEL_TPD, check_dataset_quality, load_dataset


BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
MODEL_DIR = BASE_DIR / "models"
MODEL_PATH = MODEL_DIR / "svm_voice_confidence_model.joblib"


def build_pipeline():
    """
    Pipeline wajib:
    1. StandardScaler
    2. SVM classifier
    """
    return Pipeline(
        [
            ("scaler", StandardScaler()),
            (
                "svm",
                SVC(
                    probability=True,
                    class_weight="balanced",
                    random_state=42,
                ),
            ),
        ]
    )


def build_cv(y):
    label_counts = Counter(y)
    min_class_count = min(label_counts.values())
    n_splits = min(5, min_class_count)

    if n_splits < 2:
        raise ValueError("Minimal perlu 2 data pada setiap kelas untuk Stratified K-Fold.")

    return StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)


def build_grid_search(cv):
    param_grid = {
        "svm__C": [0.1, 1, 10, 100],
        "svm__gamma": ["scale", 0.01, 0.001, 0.0001],
        "svm__kernel": ["rbf"],
    }

    return GridSearchCV(
        estimator=build_pipeline(),
        param_grid=param_grid,
        scoring="f1_macro",
        cv=cv,
        n_jobs=1,
        refit=True,
        verbose=1,
    )


def print_wrong_predictions(files, y_true, y_pred):
    print("\n=== File yang Salah Prediksi ===")
    has_wrong_prediction = False

    for file_path, true_label, predicted_label in zip(files, y_true, y_pred):
        if true_label != predicted_label:
            has_wrong_prediction = True
            print(f"{Path(file_path).name} | {true_label} | {predicted_label}")

    if not has_wrong_prediction:
        print("Tidak ada file yang salah prediksi pada cross-validation.")


def evaluate_model(model, X, y, files, cv):
    """
    Evaluasi memakai prediksi out-of-fold agar lebih realistis untuk dataset kecil.
    """
    y_pred = cross_val_predict(model, X, y, cv=cv, n_jobs=1)
    labels = [LABEL_PD, LABEL_TPD]

    print("\n=== Evaluasi Cross-Validation ===")
    print(f"Accuracy          : {accuracy_score(y, y_pred):.4f}")
    print(f"Balanced Accuracy : {balanced_accuracy_score(y, y_pred):.4f}")
    print(f"Precision Macro   : {precision_score(y, y_pred, average='macro', zero_division=0):.4f}")
    print(f"Recall Macro      : {recall_score(y, y_pred, average='macro', zero_division=0):.4f}")
    print(f"F1 Macro          : {f1_score(y, y_pred, average='macro', zero_division=0):.4f}")

    print("\n=== Classification Report ===")
    print(
        classification_report(
            y,
            y_pred,
            labels=labels,
            target_names=["PD - Percaya Diri", "TPD - Tidak Percaya Diri"],
            zero_division=0,
        )
    )

    print("=== Confusion Matrix ===")
    matrix = confusion_matrix(y, y_pred, labels=labels)
    print("Urutan label:", labels)
    print(matrix)

    print("\n=== Ringkasan Benar/Salah per Kelas ===")
    for index, label in enumerate(labels):
        total = int(matrix[index].sum())
        correct = int(matrix[index, index])
        wrong = total - correct
        print(f"{label}: benar={correct}, salah={wrong}, total={total}")

    print_wrong_predictions(files, y, y_pred)


def main():
    print("Mengecek kualitas dataset...")
    check_dataset_quality(DATA_DIR)

    print("\nMembaca dataset dan mengekstraksi fitur...")
    X, y, files = load_dataset(DATA_DIR)

    print(f"\nTotal data valid: {len(files)}")
    print(f"Jumlah fitur per audio: {X.shape[1]}")
    print("Distribusi label:", dict(Counter(y)))

    cv = build_cv(y)
    grid_search = build_grid_search(cv)

    print("\nMelakukan GridSearchCV SVM dengan scoring f1_macro...")
    grid_search.fit(X, y)

    best_model = grid_search.best_estimator_

    print("\n=== Hasil GridSearchCV ===")
    print("Best params:", grid_search.best_params_)
    print(f"Best CV f1_macro: {grid_search.best_score_:.4f}")
    print("Urutan kelas model:", list(best_model.classes_))

    evaluate_model(best_model, X, y, files, cv)

    print("\nMelatih ulang model terbaik dengan seluruh dataset...")
    best_model.fit(X, y)
    print("Urutan kelas model final:", list(best_model.classes_))

    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    joblib.dump(best_model, MODEL_PATH)

    print(f"Model terbaik berhasil disimpan ke: {MODEL_PATH}")


if __name__ == "__main__":
    main()