import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import time
import tracemalloc
import warnings
import os
warnings.filterwarnings('ignore')

from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import (train_test_split, KFold,
                                     cross_val_score, learning_curve)
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.metrics import (classification_report, confusion_matrix,
                              accuracy_score, roc_auc_score)
import joblib

# Load Dataset
df = pd.read_csv("data/features_30_sec.csv", index_col=0)
print(f"Dataset shape: {df.shape}")
print(f"Class distribution:\n{df['label'].value_counts()}\n")

# Preprocess Data
le = LabelEncoder()
y = le.fit_transform(df["label"])
X = df.drop("label", axis=1)

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled  = scaler.transform(X_test)

X_full = np.vstack([X_train_scaled, X_test_scaled])
y_full = np.hstack([y_train, y_test])

# Comutational Cost Measurement Function
def measure_cost(model, X_tr, y_tr, X_te):
    tracemalloc.start()
    t0 = time.time()
    model.fit(X_tr, y_tr)
    train_time = time.time() - t0
    _, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()

    t0 = time.time()
    model.predict(X_te)
    infer_time = (time.time() - t0) / len(X_te)

    return round(train_time, 3), round(peak / 1024**2, 3), round(infer_time * 1000, 4)

# Models
models = {
    "Logistic Regression": LogisticRegression(max_iter=1000, random_state=42),
    "Random Forest":       RandomForestClassifier(n_estimators=300, max_depth=20, random_state=42),
    "SVM (RBF)":           SVC(kernel="rbf", C=10, gamma="scale", probability=True, random_state=42),
}

results  = {}
cost_rows = []

# Train, evaluate, and measure cost for each model
for name, model in models.items():
    train_time, mem_mb, infer_ms = measure_cost(model, X_train_scaled, y_train, X_test_scaled)
    y_pred = model.predict(X_test_scaled)
    y_prob = model.predict_proba(X_test_scaled)

    acc     = accuracy_score(y_test, y_pred)
    roc_auc = roc_auc_score(y_test, y_prob, multi_class='ovr', average='macro')

    results[name] = {
        "model":   model,
        "y_pred":  y_pred,
        "y_prob":  y_prob,
        "acc":     acc,
        "roc_auc": roc_auc,
    }
    cost_rows.append({
        "Model":                name,
        "Train Time (s)":       train_time,
        "Peak Memory (MB)":     mem_mb,
        "Inference/sample (ms)": infer_ms,
    })

    print(f"\n{'='*55}")
    print(f"  {name}")
    print(f"  Accuracy : {acc*100:.2f}%   ROC-AUC : {roc_auc:.4f}")
    print(classification_report(y_test, y_pred, target_names=le.classes_))

# K-Fold Cross Validation
print("\n" + "="*55)
print("  K-FOLD CROSS VALIDATION  (K = 10)")
print("="*55)

kf = KFold(n_splits=10, shuffle=True, random_state=42)

for name, info in results.items():
    cv = cross_val_score(info["model"], X_full, y_full, cv=kf, scoring="accuracy", n_jobs=-1)
    results[name]["cv"] = cv
    print(f"\n{name}")
    print(f"  Fold scores : {cv.round(3)}")
    print(f"  Mean ± Std  : {cv.mean():.4f} ± {cv.std():.4f}")
    print(f"  Test acc    : {info['acc']:.4f}  "
          f"({'overfit' if info['acc'] > cv.mean() + 0.05 else 'generalises well'})")

# Visualizations
os.makedirs('plots', exist_ok=True)

# Heatmaps
fig, axes = plt.subplots(1, 3, figsize=(22, 7))
for ax, (name, info) in zip(axes, results.items()):
    cm = confusion_matrix(y_test, info["y_pred"])
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
                xticklabels=le.classes_, yticklabels=le.classes_, ax=ax,
                linewidths=0.5, cbar=False)
    ax.set_title(f"{name}\nAcc: {info['acc']*100:.1f}%  ROC-AUC: {info['roc_auc']:.3f}",
                 fontsize=11, fontweight="bold")
    ax.set_xlabel("Predicted Label")
    ax.set_ylabel("True Label")
    ax.tick_params(axis="x", rotation=45)
plt.suptitle("Confusion Matrices — All Models", fontsize=14, fontweight="bold", y=1.02)
plt.tight_layout()
plt.savefig("plots/confusion_matrices.png", dpi=150, bbox_inches="tight")
plt.show()

# CV Score Boxplot
fig, ax = plt.subplots(figsize=(10, 5))
cv_vals  = [info["cv"] for info in results.values()]
bp = ax.boxplot(cv_vals, labels=results.keys(), patch_artist=True, widths=0.4)
colors = ["#4e9e94", "#2a7d74", "#1a5c55"]
for patch, color in zip(bp["boxes"], colors):
    patch.set_facecolor(color)
    patch.set_alpha(0.7)
ax.set_title("K-Fold CV Score Distribution (K = 10)", fontweight="bold")
ax.set_ylabel("Accuracy")
ax.set_ylim(0.5, 1.0)
ax.axhline(y=0.8, color="red", linestyle="--", alpha=0.4, label="80% threshold")
ax.legend()
ax.grid(axis="y", linestyle="--", alpha=0.4)
plt.tight_layout()
plt.savefig("plots/cv_boxplot.png", dpi=150, bbox_inches="tight")
plt.show()

# Learning Curves — all 3 models
def plot_learning_curve(model, X, y, title):
    train_sz, train_sc, val_sc = learning_curve(
        model, X, y, cv=5, n_jobs=-1,
        train_sizes=np.linspace(0.1, 1.0, 10),
        scoring="accuracy"
    )
    t_mean, t_std = train_sc.mean(axis=1), train_sc.std(axis=1)
    v_mean, v_std = val_sc.mean(axis=1),   val_sc.std(axis=1)

    gap = t_mean[-1] - v_mean[-1]
    diagnosis = ("Overfitting"  if gap > 0.08 else
                 "Underfitting" if v_mean[-1] < 0.65 else
                 "Generalises well")

    plt.figure(figsize=(9, 5))
    plt.plot(train_sz, t_mean, "o-", color="teal",  label="Training Score")
    plt.fill_between(train_sz, t_mean - t_std, t_mean + t_std, alpha=0.15, color="teal")
    plt.plot(train_sz, v_mean, "o-", color="coral", label="Validation Score")
    plt.fill_between(train_sz, v_mean - v_std, v_mean + v_std, alpha=0.15, color="coral")
    plt.title(f"Learning Curve — {title}\nDiagnosis: {diagnosis}", fontweight="bold")
    plt.xlabel("Training Set Size")
    plt.ylabel("Accuracy")
    plt.legend(loc="lower right")
    plt.grid(linestyle="--", alpha=0.4)
    plt.ylim(0.3, 1.05)
    plt.tight_layout()
    plt.savefig(f"plots/learning_curve_{title.replace(' ','_')}.png", dpi=150, bbox_inches="tight")
    plt.show()

for name, info in results.items():
    plot_learning_curve(info["model"], X_full, y_full, name)

# Feature Importance (Random Forest only)
rf          = results["Random Forest"]["model"]
importances = pd.Series(rf.feature_importances_, index=X.columns).sort_values(ascending=True).tail(10)
fig, ax = plt.subplots(figsize=(11, 6))
importances.plot(kind="barh", color="teal", ax=ax, edgecolor="white")
ax.set_title("Top 10 Most Important Audio Features (Random Forest)", fontweight="bold")
ax.set_xlabel("Feature Importance Score")
ax.grid(axis="x", linestyle="--", alpha=0.4)
plt.tight_layout()
plt.savefig("plots/feature_importance.png", dpi=150, bbox_inches="tight")
plt.show()

# Model Accuracy Comparison Bar Chart
fig, ax = plt.subplots(figsize=(8, 5))
names  = list(results.keys())
accs   = [info["acc"] * 100 for info in results.values()]
aucs   = [info["roc_auc"] for info in results.values()]
x      = np.arange(len(names))
bars   = ax.bar(x - 0.2, accs,  0.35, label="Accuracy (%)",  color="teal",  alpha=0.8)
bars2  = ax.bar(x + 0.2, [a*100 for a in aucs], 0.35, label="ROC-AUC × 100", color="coral", alpha=0.8)
ax.set_xticks(x)
ax.set_xticklabels(names)
ax.set_ylim(60, 100)
ax.set_title("Model Comparison — Accuracy vs ROC-AUC", fontweight="bold")
ax.legend()
ax.grid(axis="y", linestyle="--", alpha=0.4)
for bar in bars:
    ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.3,
            f"{bar.get_height():.1f}", ha="center", fontsize=9)
plt.tight_layout()
plt.savefig("plots/model_comparison.png", dpi=150, bbox_inches="tight")
plt.show()

# Computational Cost Analysis
cost_df = pd.DataFrame(cost_rows)
print("\n" + "="*55)
print("  COMPUTATIONAL COST ANALYSIS")
print("="*55)
print(cost_df.to_string(index=False))

# Save the best model + scaler + label encoder
best_name = max(results, key=lambda k: results[k]["acc"])
joblib.dump(results[best_name]["model"], "music_genre_classifier.pkl")
joblib.dump(scaler, "scaler.pkl")
joblib.dump(le,     "label_encoder.pkl")
print(f"\nBest model ({best_name}) + scaler + label encoder saved.")

# Class Distribution Bar Chart
plt.figure(figsize=(10, 5))
df['label'].value_counts().sort_index().plot(kind='bar', color='teal', edgecolor='white')
plt.title("Class Distribution Across Music Genres", fontweight='bold')
plt.xlabel("Genre")
plt.ylabel("Number of Samples")
plt.xticks(rotation=45)
plt.tight_layout()
plt.savefig("plots/class_distribution.png", dpi=150)
plt.show()

# Corelation Heatmap
plt.figure(figsize=(12, 8))
top_features = X.columns[:15]
sns.heatmap(X[top_features].corr(), cmap='coolwarm', annot=False, linewidths=0.5)
plt.title("Feature Correlation Heatmap", fontweight='bold')
plt.tight_layout()
plt.savefig("plots/correlation_heatmap.png", dpi=150)
plt.show()

# Scalability Analysis — Training Time vs. Dataset Size
fractions = np.linspace(0.1, 1.0, 10)
scalability = {name: [] for name in models}

for frac in fractions:
    n = int(len(X_train_scaled) * frac)
    X_sub = X_train_scaled[:n]
    y_sub = y_train[:n]
    for name, model in models.items():
        t0 = time.time()
        model.fit(X_sub, y_sub)
        scalability[name].append(time.time() - t0)

plt.figure(figsize=(10, 5))
for name, times in scalability.items():
    plt.plot([int(len(X_train_scaled) * f) for f in fractions], times, marker='o', label=name)
plt.title("Scalability — Training Time vs. Dataset Size", fontweight='bold')
plt.xlabel("Training Samples")
plt.ylabel("Training Time (seconds)")
plt.legend()
plt.grid(linestyle='--', alpha=0.4)
plt.tight_layout()
plt.savefig("plots/scalability_curve.png", dpi=150)
plt.show()