# COMPLETE MODIFIED pacebeats_model.py - READY TO USE - last modified 4/7/2026
# This is your pacebeats_model.py file with ALL new training functions added
# Just copy-paste this entire file on Hugging Face

import os, sys, uuid, subprocess, pickle, time, logging
import numpy as np
import pandas as pd
import pytz

try:
    import lightgbm as lgb
    from pykalman import KalmanFilter
    from supabase import create_client
    from sklearn.linear_model import LogisticRegression
    from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
    from sklearn.preprocessing import StandardScaler, LabelEncoder
    from sklearn.model_selection import train_test_split, GroupKFold, TimeSeriesSplit, cross_val_score
    from sklearn.metrics import roc_auc_score, log_loss, precision_score, recall_score, f1_score, classification_report
    from sklearn.dummy import DummyClassifier
except Exception:
    print("Installing dependencies...")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy", "pandas", "pykalman", "supabase", "scikit-learn", "lightgbm"])
    import lightgbm as lgb
    from pykalman import KalmanFilter
    from supabase import create_client
    from sklearn.linear_model import LogisticRegression
    from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
    from sklearn.preprocessing import StandardScaler, LabelEncoder
    from sklearn.model_selection import train_test_split, GroupKFold, TimeSeriesSplit, cross_val_score
    from sklearn.metrics import roc_auc_score, log_loss, precision_score, recall_score, f1_score, classification_report
    from sklearn.dummy import DummyClassifier

try:
    from apscheduler.schedulers.background import BackgroundScheduler
    from apscheduler.triggers.cron import CronTrigger
except:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "apscheduler"])
    from apscheduler.schedulers.background import BackgroundScheduler
    from apscheduler.triggers.cron import CronTrigger

from datetime import datetime, timezone, timedelta
from typing import Dict, List, Optional, Tuple
import warnings

warnings.filterwarnings('ignore')

logger = logging.getLogger(__name__)
training_scheduler = None

# =========================
# Configuration
# =========================
MODEL_SAVE_PATH = "pacebeats_ml_model.pkl"

SUPABASE_URL = os.getenv("SUPABASE_URL", "https://mxhnswymqijymrwvsybm.supabase.co").strip()
SUPABASE_KEY = os.getenv("SUPABASE_KEY",
                         os.getenv("SUPABASE_SERVICE_KEY",
                                   os.getenv("SUPABASE_ANON_KEY",
                                             "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6Im14aG5zd3ltcWlqeW1yd3ZzeWJtIiwicm9sZSI6InNlcnZpY2Vfcm9sZSIsImlhdCI6MTc1MjgzMTg2NCwiZXhwIjoyMDY4NDA3ODY0fQ.bWiFaZZ1xVIyTz9dxtuyMY-odWj2gRT_yzv79FxDH3A"))).strip()

TABLE_SONGS    = "music"                 # catalog table
TABLE_EVENTS   = "listening_events"      # listening logs
TABLE_RECS     = "recommendation_served" # recs served logs

USER_ID = os.getenv("PACEBEATS_USER_ID", "00000000-0000-0000-0000-000000000000")

supabase = create_client(SUPABASE_URL, SUPABASE_KEY)

# =========================
# Pace smoothing & helpers
# =========================
kf = KalmanFilter([1], [1], 0.01, 1.0)
_km, _kc = 0.0, 1.0
def smooth_pace(raw):
    global _km, _kc
    _km, _kc = kf.filter_update(_km, _kc, observation=raw)
    return float(_km)

def compute_pace(dt_s, dist_m): return (dt_s / (dist_m / 1000.0)) if dist_m > 0 else np.inf
def sec_to_minpkm(s): return s / 60.0

PACE_BUCKETS = {
    "easy_walk": {"pace_min": 12.0, "pace_max": float("inf"), "bpm_center": 70,  "energy_target": 0.3,  "valence_target": 0.6},
    "recovery" : {"pace_min": 6.0,  "pace_max": 12.0,         "bpm_center": 90,  "energy_target": 0.4,  "valence_target": 0.55},
    "cruise"   : {"pace_min": 5.0,  "pace_max": 6.0,          "bpm_center": 130, "energy_target": 0.6,  "valence_target": 0.7},
    "tempo"    : {"pace_min": 4.0,  "pace_max": 5.0,          "bpm_center": 150, "energy_target": 0.75, "valence_target": 0.75},
    "interval" : {"pace_min": 3.0,  "pace_max": 4.0,          "bpm_center": 175, "energy_target": 0.85, "valence_target": 0.7},
    "sprint"   : {"pace_min": 0.0,  "pace_max": 3.0,          "bpm_center": 190, "energy_target": 0.9,  "valence_target": 0.6},
}
ALLOWED_MOODS = {"sad","happy","chill","hype","focus","angry"}

# =========================
# Catalog
# =========================
def fetch_catalog():
    res = supabase.table(TABLE_SONGS).select("*").execute()
    rows = res.data or []
    df = pd.DataFrame(rows)
    if df.empty:
        raise SystemExit("music table is empty. Load data first.")

    if "bpm" not in df.columns:
        df["bpm"] = pd.to_numeric(df.get("tempo", np.nan), errors="coerce")

    numeric_cols = ["bpm", "energy", "valence", "danceability", "acousticness",
                    "speechiness", "loudness", "liveness", "duration_ms"]
    for col in numeric_cols:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")

    if "duration_min" not in df.columns and "duration_ms" in df.columns:
        df["duration_min"] = df["duration_ms"] / (1000 * 60)
    if "title" not in df.columns:
        df["title"] = df.get("name", "")
    if "track_id" not in df.columns:
        raise SystemExit("music table needs a track_id column.")
    df["track_id"] = df["track_id"].astype(str)

    if "spotify_id" in df.columns:
        df["spotify_id"] = df["spotify_id"].astype(str)

    df = df.dropna(subset=["bpm"]).reset_index(drop=True)
    
    if "mood" not in df.columns or df["mood"].isna().all():
        df["mood"] = "unknown"
    
    def infer_mood(row):
        if pd.notna(row.get("mood")) and row["mood"] != "unknown":
            return row["mood"]
        
        energy = row.get("energy", 0.5)
        valence = row.get("valence", 0.5)
        
        if energy > 0.7 and valence > 0.6:
            return "hype"
        elif energy > 0.7 and valence < 0.4:
            return "angry"
        elif energy < 0.4 and valence < 0.4:
            return "sad"
        elif energy < 0.5 and valence > 0.5:
            return "chill"
        elif energy > 0.6 and valence > 0.5:
            return "happy"
        elif energy > 0.4 and energy < 0.6:
            return "focus"
        else:
            return "neutral"
    
    df["mood"] = df.apply(infer_mood, axis=1)
    
    return df

catalog = fetch_catalog()

# =========================
# Recommendation + Logging
# =========================
def log_listening_event(track_id, played_ms, skipped=False, liked=None, disliked=None, completed=False, session_id=None):
    data = {
        "user_id": USER_ID,
        "track_id": str(track_id),
        "session_id": session_id or str(uuid.uuid4()),
        "played_ms": int(played_ms),
        "skipped": bool(skipped) if skipped is not None else False,
        "liked": bool(liked) if liked is not None else None,
        "disliked": bool(disliked) if disliked is not None else None,
        "completed": bool(completed) if completed is not None else False,
    }
    data = {k: v for k, v in data.items() if v is not None}
    try:
        supabase.table(TABLE_EVENTS).insert(data).execute()
        print(f"   📝 Logged interaction: played {played_ms}ms")
    except Exception as e:
        print(f"   ⚠️ Failed to log event: {e}")

def log_recommendation_served(recommendations_df, session_id, bpm_center, pace_min,
                              user_mood=None, run_mode=None, target_pace_min=None):
    """Log all recommendations served to the user for ML training."""
    records = []
    for rank, (_, row) in enumerate(recommendations_df.iterrows()):
        rec = {
            "session_id": session_id,
            "user_id": USER_ID,
            "ts": datetime.now(timezone.utc).isoformat(),
            "track_id": str(row["track_id"]),
            "rank": rank + 1,
            "bpm_center": float(bpm_center),
            "pace_min": float(pace_min),
            "user_mood": user_mood,
            "candidate_score": float(row.get("score", row.get("rule_score", 0))),
        }
        if run_mode is not None:
            rec["run_mode"] = run_mode
        if target_pace_min is not None:
            rec["target_pace_min"] = float(target_pace_min)
        records.append(rec)

    if not records:
        return

    try:
        supabase.table(TABLE_RECS).insert(records).execute()
    except Exception as e:
        try:
            basic_records = [{k: v for k, v in r.items() if k not in ["run_mode", "target_pace_min"]} for r in records]
            supabase.table(TABLE_RECS).insert(basic_records).execute()
        except Exception as e2:
            print(f"   ⚠️ Could not log recommendations: {e2}")

# =========================
# Label logic (training)
# =========================
def compute_labels(events_df):
    labels = []
    for _, row in events_df.iterrows():
        if row.get('liked') == True:
            labels.append(1)
        elif row.get('disliked') == True:
            labels.append(0)
        elif row.get('skipped') == True and row.get('played_ms', 0) < 15000:
            labels.append(0)
        elif row.get('skipped') == False and row.get('played_ms', 0) >= 30000:
            labels.append(1)
        elif row.get('completed') == True:
            labels.append(1)
        else:
            labels.append(None)
    return labels

# =========================
# Training dataset builder
# =========================
from sklearn.preprocessing import MultiLabelBinarizer

def create_training_dataset():
    """Builds training rows by joining tables, ignoring session_id to fix mismatches"""
    try:
        recs_df = pd.DataFrame(supabase.table(TABLE_RECS).select("*").execute().data or [])
        events_df = pd.DataFrame(supabase.table(TABLE_EVENTS).select("*").execute().data or [])

        if recs_df.empty or events_df.empty:
            print("Missing data in tables.")
            return pd.DataFrame()

        # Ignore local files so the ML math does not crash
        events_df = events_df[~events_df['track_id'].astype(str).str.startswith('local:')]
        
        music_df = fetch_catalog()

        # Drop duplicates to prevent massive data multiplication
        events_df = events_df.drop_duplicates(subset=["user_id", "track_id"], keep="last")
        recs_df = recs_df.drop_duplicates(subset=["user_id", "track_id"], keep="last")

        users_df = pd.DataFrame(
            supabase.table("users").select("id,experience_duration,pace_band,unknown_pace,preferred_genres").execute().data or []
        )
        users_df.rename(columns={"id":"user_id"}, inplace=True)

        # THE FIX: Merge using ONLY user_id and track_id
        training_df = recs_df.merge(events_df, on=["user_id", "track_id"], how="inner") \
            .merge(music_df[["track_id","bpm","energy","valence","danceability","acousticness",
                             "speechiness","loudness","liveness","genre","mode","duration_min"]],
                   on="track_id", how="left") \
            .merge(users_df, on="user_id", how="left")

        training_df["label"] = compute_labels(training_df)
        training_df = training_df.dropna(subset=["label"])

        if training_df.empty:
            print("After merging, no matching rows were found with valid labels.")
            return pd.DataFrame()

        training_df["bpm_error"] = np.abs(training_df["bpm"] - training_df["bpm_center"])
        
        # Handle column renaming if both tables had a 'ts' column
        ts_col = "ts_x" if "ts_x" in training_df.columns else "ts"
        training_df["ts"] = pd.to_datetime(training_df[ts_col])
        
        training_df["hour_of_day"] = training_df["ts"].dt.hour
        training_df["day_of_week"] = training_df["ts"].dt.dayofweek
        training_df["user_total_plays"] = training_df.groupby("user_id")["track_id"].transform("count")

        training_df["experience_duration"] = training_df["experience_duration"].fillna("unknown")
        training_df["pace_band"] = training_df["pace_band"].fillna("unknown")
        training_df["unknown_pace"] = training_df["unknown_pace"].fillna(False)

        def ensure_list(x):
            if x is None or (isinstance(x, float) and np.isnan(x)):
                return []
            if isinstance(x, (list, tuple)):
                return list(x)
            return [x]
        
        training_df["preferred_genres"] = training_df.get("preferred_genres", []).apply(ensure_list)

        print(f"Training dataset created with {len(training_df)} rows")
        return training_df

    except Exception as e:
        print(f"Error creating training dataset: {e}")
        return pd.DataFrame()

def get_training_features():
    return [
        "bpm","energy","valence","danceability","acousticness",
        "speechiness","loudness","liveness",
        "bpm_error",
        "pace_min","bpm_center","hour_of_day","day_of_week",
        "candidate_score",
        "user_total_plays",
        "unknown_pace",
    ]

def get_categorical_features():
    return ["genre","mode","user_mood","experience_duration","pace_band"]

# =========================
# Candidate gen & scoring
# =========================
def generate_candidates(pace_min, user_mood=None, max_candidates=200):
    """Generate candidate songs filtered by pace (BPM) and optionally mood."""
    if not np.isfinite(pace_min): 
        print("[ERROR] Invalid pace_min")
        return pd.DataFrame(), None, None
    
    pace_bucket_info = None
    for bucket_name, info in PACE_BUCKETS.items():
        if info["pace_min"] <= pace_min < info["pace_max"]:
            pace_bucket_info = info
            break
    
    if pace_bucket_info is None: 
        print(f"[ERROR] No pace bucket for {pace_min:.1f} min/km")
        return pd.DataFrame(), None, None

    bpm_center = pace_bucket_info["bpm_center"]
    
    df = catalog.copy()
    print(f"[CATALOG] Starting with {len(df)} total songs")
    
    mood = (user_mood or "").lower().strip()
    if mood and mood in ALLOWED_MOODS and "mood" in df.columns:
        original_count = len(df)
        df = df[df["mood"].astype(str).str.lower() == mood].copy()
        print(f"[MOOD] Filtered '{mood}': {original_count} → {len(df)} songs")
    else:
        print(f"[MOOD] No mood filter applied (mood='{mood}')")

    bpm_window = 10
    candidates = df[np.abs(df["bpm"] - bpm_center) <= bpm_window].copy()
    print(f"[BPM] Window ±{bpm_window}: {len(candidates)} songs in range [{bpm_center-bpm_window}, {bpm_center+bpm_window}]")
    
    if len(candidates) < 10:
        bpm_window = 20
        candidates = df[np.abs(df["bpm"] - bpm_center) <= bpm_window].copy()
        print(f"[BPM] Widened to ±{bpm_window}: {len(candidates)} songs")
    
    if len(candidates) < 5:
        bpm_window = 30
        candidates = df[np.abs(df["bpm"] - bpm_center) <= bpm_window].copy()
        print(f"[BPM] Widened to ±{bpm_window}: {len(candidates)} songs")
    
    if len(candidates) == 0:
        candidates = df.copy()
    
    if len(candidates) > max_candidates:
        candidates = candidates.sample(n=max_candidates, random_state=42)
        print(f"[SAMPLE] Reduced to {max_candidates} candidates")
    
    return candidates.reset_index(drop=True), pace_bucket_info, bpm_center

def compute_rule_scores(candidates_df, pace_bucket_info, bpm_center):
    W_BPM, W_ENERGY, W_VALENCE, W_DANCE = 1.0, 15.0, 12.0, 5.0
    LIKE_BONUS, DISLIKE_PENALTY = -5, 5

    c = candidates_df.copy()
    c["bpm_diff"] = np.abs(c["bpm"] - bpm_center)
    c["e_diff"]   = np.abs(c.get("energy", 0.5)  - pace_bucket_info["energy_target"]).fillna(0.5)
    c["v_diff"]   = np.abs(c.get("valence", 0.5) - pace_bucket_info["valence_target"]).fillna(0.5)
    c["d_diff"]   = np.abs(c.get("danceability", 0.5) - 0.5).fillna(0.5)

    try:
        events_res = supabase.table(TABLE_EVENTS).select("track_id, liked, disliked").eq("user_id", USER_ID).execute()
        fb = pd.DataFrame(events_res.data or [])
        if not fb.empty:
            fb_agg = fb.groupby("track_id").agg({
                "liked": lambda x: x.sum() > 0,
                "disliked": lambda x: x.sum() > 0
            }).reset_index()
        else:
            fb_agg = pd.DataFrame(columns=["track_id", "liked", "disliked"])
    except Exception as e:
        print(f"   ⚠️ Could not fetch feedback: {e}")
        fb_agg = pd.DataFrame(columns=["track_id", "liked", "disliked"])

    c["track_id"] = c["track_id"].astype(str)
    c = c.merge(fb_agg, on="track_id", how="left")

    c["rule_score"] = (
        W_BPM*c["bpm_diff"] +
        W_ENERGY*c["e_diff"] +
        W_VALENCE*c["v_diff"] +
        W_DANCE*c["d_diff"]
    )
    c.loc[c["liked"]==True,  "rule_score"] += LIKE_BONUS
    c.loc[c["disliked"]==True, "rule_score"] += DISLIKE_PENALTY
    return c

# =========================
# ML re-ranking
# =========================
def ml_rerank_candidates(candidates_df, pace_min, bpm_center, user_mood=None, alpha=0.3):
    global ml_model
    if candidates_df.empty: return candidates_df
    if 'ml_model' not in globals() or ml_model is None:
        ml_model = PaceBeatsMlModel()

    candidates = candidates_df.copy()
    candidates["pace_min"]    = pace_min
    candidates["bpm_center"]  = bpm_center
    candidates["user_mood"]   = user_mood or "none"
    now = datetime.now()
    candidates["hour_of_day"] = now.hour
    candidates["day_of_week"] = now.weekday()
    candidates["user_total_plays"] = 10
    candidates["bpm_error"] = candidates["bpm_diff"]
    candidates["candidate_score"] = candidates["rule_score"]

    ml_prob = ml_model.predict_proba(candidates)
    eps = 1e-8
    ml_prob = np.clip(ml_prob, eps, 1-eps)
    ml_logit = np.log(ml_prob/(1-ml_prob))

    rule = -candidates["rule_score"].values
    rule_norm = (rule - rule.mean()) / (rule.std() + 1e-6)

    final = alpha*rule_norm + (1-alpha)*ml_logit
    candidates["ml_probability"] = ml_prob
    candidates["final_score"]    = final
    return candidates.sort_values("final_score", ascending=False)

# =========================
# Recommend (with onboarding personalization)
# =========================
def recommend_tracks_ml(pace_min, user_mood=None, top_n=5, session_id=None, use_ml=True, alpha=0.3, run_mode=None, target_pace_min=None):
    """Main recommendation pipeline."""
    global ml_model
    if 'ml_model' not in globals() or ml_model is None:
        ml_model = PaceBeatsMlModel()

    print(f"\n{'='*60}")
    print(f"[RECOMMEND] pace={pace_min:.1f} min/km, mood={user_mood}, top_n={top_n}, use_ml={use_ml}")
    
    candidates, pace_bucket_info, bpm_center = generate_candidates(pace_min, user_mood, max_candidates=200)
    
    if candidates.empty:
        print("[ERROR] No candidate tracks found")
        return pd.DataFrame()
    
    print(f"[CANDIDATES] {len(candidates)} songs to score")
    
    candidates_scored = compute_rule_scores(candidates, pace_bucket_info, bpm_center)
    
    print(f"[RULE SCORES] Min: {candidates_scored['rule_score'].min():.2f}, Max: {candidates_scored['rule_score'].max():.2f}, Mean: {candidates_scored['rule_score'].mean():.2f}")
    
    if use_ml and ml_model.is_trained:
        print("[RANKING] Using ML re-ranking")
        final_candidates = ml_rerank_candidates(candidates_scored, pace_min, bpm_center, user_mood, alpha=alpha)
    else:
        print("[RANKING] Using rule-based scoring only")
        final_candidates = candidates_scored.copy()
        final_candidates["final_score"] = 1.0 / (final_candidates["rule_score"] + 0.01)
        max_score = final_candidates["final_score"].max()
        min_score = final_candidates["final_score"].min()
        
        if max_score > min_score + 0.001:
            final_candidates["final_score"] = (final_candidates["final_score"] - min_score) / (max_score - min_score)
        
        final_candidates = final_candidates.sort_values("final_score", ascending=False)
        print(f"[FINAL SCORES] Min: {final_candidates['final_score'].min():.3f}, Max: {final_candidates['final_score'].max():.3f}")

    recs = final_candidates.head(top_n).copy()
    
    if len(recs) > 0:
        print(f"[RESULTS] Returning {len(recs)} tracks:")
        for idx, row in recs.head(3).iterrows():
            print(f"  {row.get('title', 'Unknown')} - BPM:{row['bpm']:.0f}, Score:{row['final_score']:.2f}")
    
    if session_id and not recs.empty:
        target_pace = target_pace_min if (run_mode and str(run_mode).lower() == "goal") else None
        log_recommendation_served(
            recs, session_id, bpm_center, pace_min, user_mood,
            run_mode=run_mode, target_pace_min=target_pace
        )
    
    print(f"{'='*60}\n")
    return recs

# =========================
# ML Model Class
# =========================
class PaceBeatsMlModel:
    """Trainable ML ranker for PaceBeats"""

    def __init__(self):
        self.model = None
        self.model_type = None
        self.scaler = StandardScaler()
        self.label_encoders: Dict[str, LabelEncoder] = {}
        self.feature_names: List[str] = []
        self.categorical_features: List[str] = []
        self.is_trained = False
        self.training_metrics: Dict = {}
        self.evaluation_results: Dict = {}
        self.load_model()

    def save_model(self):
        if self.is_trained:
            with open(MODEL_SAVE_PATH, "wb") as f:
                pickle.dump({
                    "model": self.model,
                    "model_type": self.model_type,
                    "scaler": self.scaler,
                    "label_encoders": self.label_encoders,
                    "feature_names": self.feature_names,
                    "categorical_features": self.categorical_features,
                    "training_metrics": self.training_metrics,
                }, f)
            print(f"✅ Model saved to {MODEL_SAVE_PATH}")

    def load_model(self):
        if os.path.exists(MODEL_SAVE_PATH):
            try:
                with open(MODEL_SAVE_PATH, "rb") as f:
                    state = pickle.load(f)
                self.model = state.get("model")
                self.model_type = state.get("model_type")
                self.scaler = state.get("scaler", StandardScaler())
                self.label_encoders = state.get("label_encoders", {})
                self.feature_names = state.get("feature_names", [])
                self.categorical_features = state.get("categorical_features", [])
                self.training_metrics = state.get("training_metrics", {})
                self.is_trained = True
                print(f"✅ Model loaded from {MODEL_SAVE_PATH}")
                return True
            except Exception as e:
                print(f"⚠️ Could not load model: {e}")
        return False

    def prepare_features(self, df: pd.DataFrame, is_training: bool=False):
        numeric_features     = get_training_features()
        categorical_features = get_categorical_features()
        self.categorical_features = categorical_features

        feat_df = df.copy()

        for feat in numeric_features:
            if feat not in feat_df.columns:
                feat_df[feat] = 0

        for cat in categorical_features:
            if cat not in feat_df.columns:
                feat_df[cat] = "unknown"
            
            if is_training:
                self.label_encoders[cat] = LabelEncoder()
                feat_df[f"{cat}_encoded"] = self.label_encoders[cat].fit_transform(feat_df[cat].astype(str))
            else:
                if cat in self.label_encoders:
                    feat_df[f"{cat}_encoded"] = self.label_encoders[cat].transform(feat_df[cat].astype(str))
                else:
                    feat_df[f"{cat}_encoded"] = 0

        pref_cols = [c for c in feat_df.columns if c.startswith("pref_genre_")]
        final_cols = numeric_features + [f"{c}_encoded" for c in categorical_features] + pref_cols

        if is_training:
            self.feature_names = final_cols

        X = feat_df.reindex(columns=self.feature_names, fill_value=0)

        if is_training:
            Xs = self.scaler.fit_transform(X)
        else:
            Xs = self.scaler.transform(X)
        return Xs

    def _baseline_models(self):
        return {
            'dummy_most_frequent': DummyClassifier(strategy='most_frequent', random_state=42),
            'dummy_uniform'      : DummyClassifier(strategy='uniform', random_state=42),
            'logistic_regression': LogisticRegression(max_iter=2000, class_weight='balanced', random_state=42),
            'random_forest'      : RandomForestClassifier(n_estimators=100, max_depth=10, class_weight='balanced', random_state=42),
            'gradient_boosting'  : GradientBoostingClassifier(n_estimators=100, max_depth=6, random_state=42),
            'lightgbm'           : lgb.LGBMClassifier(n_estimators=100, max_depth=6, class_weight='balanced', random_state=42, verbose=-1),
        }

    def _precision_at_k(self, y_true, y_scores, k=5):
        idx = np.argsort(y_scores)[::-1][:k]
        return float(np.sum(y_true[idx])) / max(k,1)

    def _ndcg_at_k(self, y_true, y_scores, k=5):
        idx = np.argsort(y_scores)[::-1][:k]
        rel = y_true[idx]
        dcg = np.sum(rel / np.log2(np.arange(2, len(rel)+2)))
        ideal_idx = np.argsort(y_true)[::-1][:k]
        ideal = y_true[ideal_idx]
        idcg = np.sum(ideal / np.log2(np.arange(2, len(ideal)+2)))
        return float(dcg / idcg) if idcg > 0 else 0.0

    def _evaluate(self, model, X, y, name="model"):
        prob = model.predict_proba(X)[:,1]
        pred = (prob>=0.5).astype(int)
        # roc_auc_score requires at least 2 classes in y_true
        auc = roc_auc_score(y, prob) if len(np.unique(y)) > 1 else 0.5
        return {
            "model_name": name,
            "auc": auc,
            "logloss": log_loss(y, prob),
            "precision": precision_score(y, pred, zero_division=0),
            "recall": recall_score(y, pred, zero_division=0),
            "f1": f1_score(y, pred, zero_division=0),
            "precision_at_5": self._precision_at_k(y, prob, k=5),
            "ndcg_at_5": self._ndcg_at_k(y, prob, k=5),
            "n_samples": len(y),
        }

    def _time_split(self, df, test_size=0.2):
        if 'ts' not in df.columns:
            return train_test_split(df, test_size=test_size, random_state=42)
        d = df.sort_values('ts')
        cut = int(len(d)*(1-test_size))
        return d.iloc[:cut], d.iloc[cut:]

    def _user_split(self, df, test_size=0.2):
        if 'user_id' not in df.columns:
            return train_test_split(df, test_size=test_size, random_state=42)
        users = df['user_id'].unique()
        ntest = max(1, int(len(users)*test_size))
        np.random.seed(42)
        test_users = np.random.choice(users, ntest, replace=False)
        train_df = df[~df['user_id'].isin(test_users)]
        test_df  = df[df['user_id'].isin(test_users)]
        if test_df.empty or test_df['label'].nunique()<2:
            return train_test_split(df, test_size=test_size, random_state=42)
        return train_df, test_df

    def train_with_evaluation(self, training_df: pd.DataFrame, model_type='lightgbm', test_size=0.2, cv_splits=5):
        if training_df.empty:
            print("❌ Training dataframe is empty")
            return False

        print(f"🤖 Training {model_type} with comprehensive evaluation...")
        X_all = self.prepare_features(training_df, is_training=True)
        y_all = training_df['label'].values

        if 'ts' in training_df.columns:
            tr_df, te_df = self._time_split(training_df, test_size)
        elif 'user_id' in training_df.columns:
            tr_df, te_df = self._user_split(training_df, test_size)
        else:
            tr_df, te_df = train_test_split(training_df, test_size=test_size, random_state=42)

        X_train = self.prepare_features(tr_df, is_training=False)
        X_test  = self.prepare_features(te_df, is_training=False)
        y_train = tr_df['label'].values
        y_test  = te_df['label'].values

        models = self._baseline_models()
        if model_type not in models:
            print(f"⚠️ Unknown model type {model_type}, using lightgbm")
            model_type = 'lightgbm'
        self.model = models[model_type]
        self.model_type = model_type

        print(f"Training {model_type}...")
        self.model.fit(X_train, y_train)
        
        eval_result = self._evaluate(self.model, X_test, y_test, model_type)
        self.training_metrics = eval_result
        print(f"Test AUC: {eval_result['auc']:.3f}, F1: {eval_result['f1']:.3f}")

        self.is_trained = True
        self.save_model()
        return True

    def train(self, training_df: pd.DataFrame):
        return self.train_with_evaluation(training_df)

    def predict_proba(self, candidates_df: pd.DataFrame):
        if not self.is_trained or self.model is None:
            return np.ones(len(candidates_df)) * 0.5
        X = self.prepare_features(candidates_df, is_training=False)
        prob = self.model.predict_proba(X)[:, 1]
        return prob

    def get_feature_importance(self):
        if not hasattr(self.model, 'feature_importances_'):
            return None
        return dict(zip(self.feature_names, self.model.feature_importances_))


# =========================
# NEW: HYBRID TRAINING FUNCTIONS
# =========================

def update_user_preference_cache(user_id: str):
    """Update user preference cache after each run (lightweight)"""
    try:
        events = supabase.table("listening_events") \
            .select("track_id, liked, disliked, played_ms") \
            .eq("user_id", user_id) \
            .limit(500) \
            .execute()
        
        if not events.data:
            return None
        
        events_df = pd.DataFrame(events.data)
        
        track_ids = events_df["track_id"].unique()
        music_data = supabase.table("music") \
            .select("track_id, bpm, energy, valence, genre") \
            .in_("track_id", list(track_ids)) \
            .execute()
        
        if not music_data.data:
            return None
        
        music_df = pd.DataFrame(music_data.data)
        merged = events_df.merge(music_df, on="track_id", how="left")
        
        avg_bpm = merged["bpm"].mean()
        avg_energy = merged["energy"].mean()
        avg_valence = merged["valence"].mean()
        
        likes = merged[merged["liked"] == True]
        if len(likes) > 0:
            liked_energy = likes["energy"].mean()
            liked_valence = likes["valence"].mean()
            
            if liked_energy > 0.7 and liked_valence > 0.6:
                preferred_mood = "hype"
            elif liked_energy > 0.7 and liked_valence < 0.4:
                preferred_mood = "angry"
            elif liked_energy < 0.4 and liked_valence < 0.4:
                preferred_mood = "sad"
            elif liked_energy < 0.5 and liked_valence > 0.5:
                preferred_mood = "chill"
            else:
                preferred_mood = "neutral"
        else:
            preferred_mood = "neutral"
        
        total_feedback = len(merged[merged["liked"].notna()])
        like_ratio = len(likes) / total_feedback if total_feedback > 0 else 0.5
        
        cache_data = {
            "user_id": user_id,
            "avg_bpm": float(avg_bpm) if pd.notna(avg_bpm) else None,
            "preferred_mood": preferred_mood,
            "avg_energy": float(avg_energy) if pd.notna(avg_energy) else None,
            "avg_valence": float(avg_valence) if pd.notna(avg_valence) else None,
            "total_runs": len(events_df.groupby("session_id")),
            "total_feedback_count": total_feedback,
            "last_like_dislike_ratio": float(like_ratio),
            "updated_at": datetime.now(timezone.utc).isoformat(),
        }
        
        supabase.table("user_preference_cache") \
            .upsert(cache_data) \
            .execute()
        
        print(f"✅ Updated user {user_id} preference cache")
        return cache_data
        
    except Exception as e:
        print(f"⚠️ Failed to update preference cache: {e}")
        return None


def update_model_incrementally(new_events_df: pd.DataFrame):
    """Add new training data without full retrain (fast warm-start)"""
    global ml_model
    
    if ml_model is None:
        ml_model = PaceBeatsMlModel()
    
    if not ml_model.is_trained:
        print("⚠️ Model not trained yet - cannot do incremental update. Full train required.")
        return False
    
    if new_events_df.empty:
        print("⚠️ No new events to learn from")
        return False
    
    try:
        print("\n🔄 Incremental Model Update (warm_start)")
        print("="*60)
        
        training_id = str(uuid.uuid4())
        
        supabase.table("model_training_logs").insert({
            "training_id": training_id,
            "status": "in_progress",
            "training_type": "incremental",
            "started_at": datetime.now(timezone.utc).isoformat(),
        }).execute()
        
        X_new = ml_model.prepare_features(new_events_df, is_training=False)
        y_new = new_events_df["label"].values
        
        print(f"📊 New samples: {len(X_new)} | Positive: {int(y_new.sum())} | Negative: {int((1-y_new).sum())}")
        
        if hasattr(ml_model.model, "warm_start"):
            ml_model.model.warm_start = True
            start_time = time.time()
            ml_model.model.fit(X_new, y_new)
            elapsed = time.time() - start_time
            
            print(f"✅ Warm-start fit completed in {elapsed:.2f}s")
            
            if hasattr(ml_model.model, "predict_proba"):
                proba = ml_model.model.predict_proba(X_new)
                pred = (proba[:, 1] >= 0.5).astype(int)
                f1 = f1_score(y_new, pred, zero_division=0)
                print(f"📈 F1-Score on new data: {f1:.3f}")
            
            ml_model.save_model()
            
            supabase.table("model_training_metadata").update({
                "last_incremental_update_at": datetime.now(timezone.utc).isoformat(),
                "run_count_since_incremental": 0,
                "updated_at": datetime.now(timezone.utc).isoformat(),
            }).eq("id", 1).execute()
            
            supabase.table("model_training_logs").update({
                "status": "completed",
                "completed_at": datetime.now(timezone.utc).isoformat(),
                "duration_seconds": int(elapsed),
                "training_samples": len(X_new),
                "metrics": {"f1": f1} if hasattr(ml_model.model, "predict_proba") else None,
            }).eq("training_id", training_id).execute()
            
            print(f"{'='*60}\n")
            return True
        else:
            raise Exception(f"Model {ml_model.model_type} does not support warm_start")
            
    except Exception as e:
        print(f"❌ Incremental update failed: {e}")
        supabase.table("model_training_logs").update({
            "status": "failed",
            "completed_at": datetime.now(timezone.utc).isoformat(),
            "error_message": str(e),
        }).eq("training_id", training_id).execute()
        return False


def get_training_data_since_last_update():
    """Fetch only NEW training data since last training"""
    try:
        metadata = supabase.table("model_training_metadata") \
            .select("last_incremental_update_at, last_trained_at") \
            .execute()
        
        if not metadata.data:
            since_time = None
        else:
            since_time = metadata.data[0].get("last_incremental_update_at") or \
                        metadata.data[0].get("last_trained_at")
        
        if since_time:
            events_df = pd.DataFrame(
                supabase.table("listening_events") \
                    .select("*") \
                    .gte("ts_start", since_time) \
                    .execute().data or []
            )
        else:
            events_df = pd.DataFrame(
                supabase.table("listening_events").select("*").execute().data or []
            )
        
        if events_df.empty:
            return pd.DataFrame()
        
        recs_df = pd.DataFrame(
            supabase.table("recommendation_served") \
                .select("*") \
                .in_("track_id", list(events_df["track_id"].unique())) \
                .execute().data or []
        )
        
        music_df = fetch_catalog()
        
        training_df = events_df.merge(recs_df, on=["user_id", "track_id", "session_id"], how="left") \
            .merge(music_df[["track_id", "bpm", "energy", "valence", "danceability", "acousticness",
                             "speechiness", "loudness", "liveness", "genre", "mode", "duration_min"]],
                   on="track_id", how="left")
        
        training_df["label"] = compute_labels(training_df)
        training_df = training_df.dropna(subset=["label"])
        
        print(f"📊 New training data: {len(training_df)} samples")
        return training_df
        
    except Exception as e:
        print(f"⚠️ Error fetching new training data: {e}")
        return pd.DataFrame()


def scheduled_full_retraining():
    """Full retraining with CV + evaluation (nightly)"""
    global ml_model
    
    if ml_model is None:
        ml_model = PaceBeatsMlModel()
    
    training_id = str(uuid.uuid4())
    
    try:
        print("\n🤖 Scheduled Full Model Retraining (Nightly)")
        print("============================================================")
        print(f"Training ID: {training_id}")
        print(f"Time: {datetime.now(timezone.utc).isoformat()}")
        
        supabase.table("model_training_logs").insert({
            "training_id": training_id,
            "status": "in_progress",
            "training_type": "full_cv_all_models",
            "started_at": datetime.now(timezone.utc).isoformat(),
        }).execute()
        
        df = create_training_dataset()
        
        if df.empty or len(df) < 50:
            print(f"⚠️ Insufficient training data ({len(df)} rows). Skipping full retrain.")
            supabase.table("model_training_logs").update({
                "status": "completed",
                "completed_at": datetime.now(timezone.utc).isoformat(),
                "error_message": f"Insufficient data: {len(df)} rows",
            }).eq("training_id", training_id).execute()
            return False
        
        print(f"📊 Training dataset: {len(df)} samples")
        
        algorithms = {
            1: 'lightgbm',
            2: 'random_forest',
            3: 'gradient_boosting',
            4: 'logistic_regression'
        }
        
        best_f1 = -1
        best_algo = None
        
        for db_id, algo in algorithms.items():
            start_time = time.time()
            print(f"Training {algo}...")
            
            ok = ml_model.train_with_evaluation(df, algo, test_size=0.2, cv_splits=5)
            elapsed = time.time() - start_time
            
            if ok:
                current_f1 = ml_model.training_metrics.get("f1", 0)
                
                supabase.table("model_training_metadata").upsert({
                    "id": db_id,
                    "is_trained": True,
                    "model_type": algo,
                    "last_trained_at": datetime.now(timezone.utc).isoformat(),
                    "last_cv_metrics": ml_model.training_metrics,
                    "training_duration_seconds": int(elapsed),
                    "training_samples_count": len(df),
                    "run_count_since_incremental": 0,
                    "updated_at": datetime.now(timezone.utc).isoformat(),
                }).execute()
                
                if current_f1 > best_f1:
                    best_f1 = current_f1
                    best_algo = algo
        
        if best_algo:
            ml_model.train_with_evaluation(df, best_algo, test_size=0.2, cv_splits=5)
            ml_model.save_model()
            
            supabase.table("model_training_logs").update({
                "status": "completed",
                "completed_at": datetime.now(timezone.utc).isoformat(),
                "metrics": {"best_model": best_algo, "best_f1": best_f1},
            }).eq("training_id", training_id).execute()
            
            print(f"✅ Full retraining completed. Best model: {best_algo}")
            return True
        else:
            raise Exception("All model trainings failed")
            
    except Exception as e:
        import traceback
        print(f"❌ Full retraining failed: {e}")
        supabase.table("model_training_logs").update({
            "status": "failed",
            "completed_at": datetime.now(timezone.utc).isoformat(),
            "error_message": str(e),
        }).eq("training_id", training_id).execute()
        return False


def record_feedback(track_id: str, liked: bool):
    """Record user feedback on a track"""
    try:
        data = {
            "user_id": USER_ID,
            "track_id": str(track_id),
            "liked": bool(liked),
        }
        supabase.table(TABLE_EVENTS).insert(data).execute()
        print(f"✅ Feedback recorded: {track_id} - {'👍' if liked else '👎'}")
    except Exception as e:
        print(f"⚠️ Failed to record feedback: {e}")

# =========================
# MODEL COMPARISON FOR THESIS
# =========================

def compare_all_models(training_df: pd.DataFrame, model_types=None):
    """
    Train all models and compare accuracy metrics for thesis presentation.
    Returns DataFrame with side-by-side comparison.
    """
    if model_types is None:
        model_types = ['logistic_regression', 'random_forest', 'gradient_boosting', 'lightgbm']
    
    if training_df.empty:
        print("❌ No training data available")
        return None
    
    print("\n" + "="*80)
    print("🎓 MODEL COMPARISON FOR THESIS - Training All Algorithms")
    print("="*80)
    print(f"Training Data: {len(training_df)} samples\n")
    
    results = []
    
    # Prepare data once
    model_obj = PaceBeatsMlModel()
    X_all = model_obj.prepare_features(training_df, is_training=True)
    y_all = training_df['label'].values
    
    # Train/test split
    if 'ts' in training_df.columns:
        tr_df, te_df = model_obj._time_split(training_df, 0.2)
    else:
        tr_df, te_df = model_obj._user_split(training_df, 0.2)
    
    X_train = model_obj.prepare_features(tr_df, is_training=False)
    X_test = model_obj.prepare_features(te_df, is_training=False)
    y_train = tr_df['label'].values
    y_test = te_df['label'].values
    
    # Get all models
    all_models_dict = model_obj._baseline_models()
    
    for model_type in model_types:
        if model_type not in all_models_dict:
            print(f"⚠️ Skipping {model_type} (not found)")
            continue
        
        print(f"\n📊 Training {model_type.upper()}...")
        try:
            model = all_models_dict[model_type]
            model.fit(X_train, y_train)
            
            # Evaluate
            eval_result = model_obj._evaluate(model, X_test, y_test, model_type)
            results.append(eval_result)
            
            print(f"   ✅ AUC: {eval_result['auc']:.4f} | F1: {eval_result['f1']:.4f} | Precision: {eval_result['precision']:.4f} | Recall: {eval_result['recall']:.4f}")
            
        except Exception as e:
            print(f"   ❌ Error: {e}")
    
    # Create comparison DataFrame
    if results:
        comparison_df = pd.DataFrame(results)
        comparison_df = comparison_df[['model_name', 'auc', 'logloss', 'precision', 'recall', 'f1', 'precision_at_5', 'ndcg_at_5']]
        comparison_df = comparison_df.round(4)
        
        # Rank by F1 score
        comparison_df['rank'] = comparison_df['f1'].rank(ascending=False).astype(int)
        comparison_df = comparison_df.sort_values('f1', ascending=False)
        
        print("\n" + "="*80)
        print("🏆 FINAL COMPARISON TABLE")
        print("="*80)
        print(comparison_df.to_string(index=False))
        print("="*80)
        
        # Winner
        best_model = comparison_df.iloc[0]
        print(f"\n🥇 BEST MODEL: {best_model['model_name'].upper()}")
        print(f"   F1-Score: {best_model['f1']:.4f}")
        print(f"   AUC-ROC: {best_model['auc']:.4f}")
        print(f"   Precision@5: {best_model['precision_at_5']:.4f}")
        print(f"   NDCG@5: {best_model['ndcg_at_5']:.4f}\n")
        
        return comparison_df
    else:
        print("❌ No models trained successfully")
        return None


def export_comparison_to_csv(comparison_df, filename="model_comparison_results.csv"):
    """Export comparison results to CSV for thesis"""
    if comparison_df is not None:
        comparison_df.to_csv(filename, index=False)
        print(f"✅ Results exported to {filename}")
        return filename
    return None


def generate_thesis_report(training_df: pd.DataFrame):
    """
    Generate complete thesis report with:
    1. Model comparison
    2. Feature importance
    3. Recommendation accuracy breakdown
    """
    print("\n📋 GENERATING THESIS REPORT...")
    
    # 1. Compare models
    comparison_df = compare_all_models(training_df)
    
    if comparison_df is not None:
        # 2. Train best model and get feature importance
        best_model_name = comparison_df.iloc[0]['model_name']
        print(f"\n🔍 Training best model ({best_model_name}) for feature importance...")
        
        model_obj = PaceBeatsMlModel()
        ok = model_obj.train_with_evaluation(training_df, best_model_name)
        
        if ok and hasattr(model_obj.model, 'feature_importances_'):
            feature_importance = model_obj.get_feature_importance()
            if feature_importance:
                importance_df = pd.DataFrame(
                    list(feature_importance.items()),
                    columns=['feature', 'importance']
                ).sort_values('importance', ascending=False)
                
                print("\n📊 TOP 10 MOST IMPORTANT FEATURES:")
                print(importance_df.head(10).to_string(index=False))
        
        # 3. Export
        csv_file = export_comparison_to_csv(comparison_df)
        
        return {
            "comparison_df": comparison_df,
            "best_model": best_model_name,
            "csv_file": csv_file
        }
    
    return None

# Initialize - always create a model instance at startup
# This loads from saved .pkl if it exists, otherwise starts fresh (is_trained=False)
ml_model = PaceBeatsMlModel()

# =========================
# NEW: MULTI-MODEL TRAINING FUNCTIONS
# =========================

def scheduled_full_retraining():
    """Full retraining for all 4 algorithms (Nightly comparison for Thesis)"""
    global ml_model
    if ml_model is None:
        ml_model = PaceBeatsMlModel()
    
    training_id = str(uuid.uuid4())
    try:
        print(f"\n🤖 Starting Multi-Model Retraining | ID: {training_id}")
        print("============================================================")
        
        # 1. Fetch Training Data
        df = create_training_dataset()
        
        # 2. Check if data exists (Removed the 50 sample limit)
        if df.empty:
            print("⚠️ No training data found. Cannot train models.")
            return False

        # 3. Define the 4 Algorithms and their Supabase Row IDs
        algorithms = {
            1: 'lightgbm',
            2: 'random_forest',
            3: 'gradient_boosting',
            4: 'logistic_regression'
        }
        
        best_f1 = -1
        best_algo = None

        for db_id, algo_name in algorithms.items():
            print(f"📊 Training {algo_name.upper()} (Assigning to Row ID: {db_id})...")
            start_time = time.time()
            
            # Train and Evaluate specifically for this algorithm
            success = ml_model.train_with_evaluation(df, algo_name, test_size=0.2, cv_splits=5)
            elapsed = time.time() - start_time
            
            if success:
                current_metrics = ml_model.training_metrics
                if current_metrics.get("f1", 0) > best_f1:
                    best_f1 = current_metrics["f1"]
                    best_algo = algo_name
                
                # Upsert into specific row ID (1, 2, 3, or 4) to show all 4 statuses
                # 1. Update the metadata table
                try:
                    supabase.table("model_training_metadata").upsert({
                        "id": db_id,
                        "is_trained": True,
                        "model_type": algo_name,
                        "last_trained_at": datetime.now(timezone.utc).isoformat(),
                        "last_cv_metrics": current_metrics,
                        "training_duration_seconds": int(elapsed),
                        "training_samples_count": len(df),
                        "run_count_since_incremental": 0, # THIS FIXES THE CREATION ERROR
                        "updated_at": datetime.now(timezone.utc).isoformat(),
                    }).execute()
                    
                    # 2. Add a record to the logs table
                    supabase.table("model_training_logs").insert({
                        "training_id": str(uuid.uuid4()),
                        "status": "completed",
                        "training_type": f"cv_{algo_name}",
                        "completed_at": datetime.now(timezone.utc).isoformat(),
                        "duration_seconds": int(elapsed),
                        "training_samples": len(df),
                        "metrics": current_metrics
                    }).execute()
                    
                    print(f"✅ Saved {algo_name} to metadata and logs")
                except Exception as db_err:
                    print(f"⚠️ SUPABASE ERROR for {algo_name}: {db_err}")

        # 4. Finalize with the best model for active use in the app
        if best_algo:
            ml_model.train_with_evaluation(df, best_algo)
            ml_model.save_model()
            print(f"🏆 BEST MODEL SELECTED: {best_algo} (F1: {best_f1:.4f})")
            print("============================================================\n")
            return True
        return False

    except Exception as e:
        print(f"❌ Retraining Loop Failed: {e}")
        return False

# =========================
# Scheduler Setup
# =========================
def start_scheduler():
    # 1. Setup Timezone
    ph_tz = pytz.timezone('Asia/Manila')
    
    # 2. Initialize Scheduler
    scheduler = BackgroundScheduler()
    
    # 3. Add the 1:00 AM Job
    scheduler.add_job(
        scheduled_full_retraining,
        CronTrigger(hour=1, minute=15, timezone=ph_tz),
        id="nightly_retrain_1am",
        replace_existing=True
    )
    
    scheduler.start()
    print("🚀 Scheduler active: Models will train every day at 1:00 AM Manila Time.")

if __name__ == "__main__":
    print("------------------------------------------------------------")
    print("STARTING PACEBEATS MODEL SYSTEM")
    print("------------------------------------------------------------")
    
    # 4. IMMEDIATE TRIGGER (Run once right now to verify Supabase updates)
    print("⚡ Running immediate training cycle to verify database connection...")
    scheduled_full_retraining()
    
    # 5. START BACKGROUND SCHEDULER
    start_scheduler()
    
    try:
        while True:
            time.sleep(60)
    except (KeyboardInterrupt, SystemExit):
        print("Stopping scheduler...")