# app/services/rl_engine.py
"""
Contextual Multi-Armed Bandit RL engine — v4.0

v4.0 changes:
  - Fixed get_next_day_signal() _thompson_sample unpack crash (was 2-tuple, needs 3)
  - Fixed consecutive_same_action counter (elif/pass logic bug)
  - Moved _compute_circadian_bonus() outside action loop (5x DB query reduction)
  - Added propensity_array to get_stable_signal() response (KeyError fix)
  - Fixed state_updated_at timezone comparison (midnight re-computation bug)
  - Corrected Normal mode debt cost: 1.0 -> 0.3 (sustainable mode)
  - Improved overwork signal: normalized weighted formula [0,1]
  - Improved burnout_warning: added single high-fatigue override
  - Removed dual-decay conflict between _load_bandit_params and _apply_memory_decay
  - Reordered _apply_safety_guard for deload -> fatigue ceiling -> circuit breaker clarity
  - Updated version docstring
"""

import math
import json
import random
import logging
from datetime import datetime, timezone, date, timedelta
from typing import Dict, List, Optional, Tuple

from sqlalchemy.orm import Session
from app.models.user import User

logger = logging.getLogger("hale.rl")

ACTIONS = ["Recovery", "Light Review", "Normal", "Deep Work", "Exploration"]

ACTION_INTENSITIES: Dict[str, Tuple[float, float]] = {
    "Recovery":     (0.10, 0.30),
    "Light Review": (0.25, 0.50),
    "Normal":       (0.45, 0.70),
    "Deep Work":    (0.65, 0.95),
    "Exploration":  (0.40, 0.65),
}

ACTION_FATIGUE_CEILING: Dict[str, float] = {
    "Deep Work":   0.75,
    "Exploration": 0.80,
    "Normal":      0.85,
}

PARAM_FLOOR = 1.0
PARAM_CAP   = 50.0

DEFAULT_BANDIT_PARAMS: Dict[str, Dict[str, float]] = {
    action: {"alpha": 2.0, "beta": 2.0} for action in ACTIONS
}
DEFAULT_BANDIT_PARAMS["Normal"]["alpha"] = 3.0

PLATEAU_WINDOW        = 7
PLATEAU_VAR_THRESHOLD = 0.04
COMPETENCE_DECAY_PER_DAY = 0.95
HISTORY_WINDOW_DAYS   = 14


def _get_default_params() -> Dict:
    return json.loads(json.dumps(DEFAULT_BANDIT_PARAMS))


def _clamp_params(params: Dict) -> Dict:
    for action in ACTIONS:
        if action in params:
            params[action]["alpha"] = max(PARAM_FLOOR, min(PARAM_CAP, params[action]["alpha"]))
            params[action]["beta"]  = max(PARAM_FLOOR, min(PARAM_CAP, params[action]["beta"]))
    return params


def _load_bandit_params(user: User) -> Dict:
    """Load stored bandit params. Does NOT apply time decay (handled by _apply_memory_decay weekly)."""
    params = _get_default_params()
    if getattr(user, "rl_bandit_params", None):
        try:
            stored = (
                user.rl_bandit_params
                if isinstance(user.rl_bandit_params, dict)
                else json.loads(user.rl_bandit_params)
            )
            for action in ACTIONS:
                if action in stored and "alpha" in stored[action] and "beta" in stored[action]:
                    params[action] = stored[action]
            # Carry over debt and metadata keys
            if "_recovery_debt" in stored:
                params["_recovery_debt"] = stored["_recovery_debt"]
            if "_last_updated" in stored:
                params["_last_updated"] = stored["_last_updated"]
        except (json.JSONDecodeError, TypeError, KeyError):
            logger.warning("Corrupted bandit params for %s — using defaults", user.user_id)

    # Natural debt recovery: debt decays passively each day without any session
    last_updated = params.get("_last_updated")
    if last_updated:
        try:
            last_dt = datetime.fromisoformat(last_updated)
            if last_dt.tzinfo is None:
                last_dt = last_dt.replace(tzinfo=timezone.utc)
            elapsed_days = max(0.0, (datetime.now(timezone.utc) - last_dt).total_seconds() / 86400.0)
            debt = params.get("_recovery_debt", 0.0)
            params["_recovery_debt"] = max(0.0, debt - (elapsed_days * 2.0))  # Slower natural recovery
        except Exception:
            pass

    return _clamp_params(params)


def _compute_historical_features(user: User, db: Optional[Session] = None) -> Dict[str, float]:
    """Query Progress table for last 14 days and extract real behavioral patterns."""
    defaults = {
        "consecutive_hard_days":  0.0,
        "recent_completion_rate": 0.7,
        "avg_satisfaction_7d":    0.6,
        "fatigue_trend":          0.0,
        "overwork_signal":        0.0,
        "deep_work_streak":       0.0,
        "avg_daily_minutes":      45.0,
        "consecutive_same_action": 0.0,
        "ema_fatigue":            0.0,
        "ema_mood":               0.5,
        "last_action":            None,
        "max_goal_completion":    0.0,
        "weekend_rate":           0.5,
        "weekday_rate":           0.5,
    }
    if db is None:
        return defaults
    try:
        from app.models.progress import Progress
        cutoff = (date.today() - timedelta(days=HISTORY_WINDOW_DAYS)).isoformat()
        rows = (
            db.query(Progress)
            .filter(Progress.user_id == user.user_id, Progress.date >= cutoff)
            .order_by(Progress.date.asc())
            .all()
        )
        if not rows:
            return defaults

        by_date: Dict[str, list] = {}
        for r in rows:
            by_date.setdefault(r.date, []).append(r)

        sorted_dates = sorted(by_date.keys())
        today_str    = date.today().isoformat()

        # Feature 1: Exponential Recency Weighting
        recent_cut  = (date.today() - timedelta(days=7)).isoformat()
        recent_dates = [d for d in sorted_dates if d >= recent_cut]
        
        lambda_decay = 0.15
        total_weight_comp, done_weight = 0.0, 0.0
        total_weight_sat, sat_weight_sum = 0.0, 0.0
        
        for d in recent_dates:
            days_ago = max(0, (date.today() - date.fromisoformat(d)).days)
            weight = math.exp(-lambda_decay * days_ago)
            
            for r in by_date[d]:
                total_weight_comp += weight
                if r.completed:
                    done_weight += weight
                
                if r.satisfaction is not None and r.completed:
                    total_weight_sat += weight
                    sat_weight_sum += r.satisfaction * weight

        completion_rate = (done_weight / total_weight_comp) if total_weight_comp > 0 else 0.7
        avg_sat = (sat_weight_sum / total_weight_sat) if total_weight_sat > 0 else 0.6

        # Feature 3: Goal Completion Proximity Effect
        max_completion = 0.0
        try:
            from app.models.goal import Goal
            active_goals = db.query(Goal).filter(Goal.user_id == user.user_id, Goal.is_active == True).all()
            if active_goals:
                max_completion = max((g.completion_pct or 0.0) for g in active_goals)
        except Exception as e:
            logger.debug("Goal proximity fetch failed: %s", e)

        # Feature 4: Weekend vs Weekday Profile (over full history window)
        weekend_done, weekend_total = 0, 0
        weekday_done, weekday_total = 0, 0
        for d in sorted_dates:
            is_weekend = date.fromisoformat(d).weekday() >= 5
            for r in by_date[d]:
                if is_weekend:
                    weekend_total += 1
                    if r.completed: weekend_done += 1
                else:
                    weekday_total += 1
                    if r.completed: weekday_done += 1
                    
        weekend_rate = (weekend_done / weekend_total) if weekend_total > 0 else 0.5
        weekday_rate = (weekday_done / weekday_total) if weekday_total > 0 else 0.5

        # Consecutive hard days (actual study time >= 60 min)
        consecutive_hard = 0
        for d in reversed(sorted_dates):
            if d == today_str:
                continue
            total_min = sum(r.actual_duration_min or 0 for r in by_date[d] if r.completed)
            if total_min >= 60:
                consecutive_hard += 1
            else:
                break

        # Average daily minutes
        comp_rows     = [r for r in rows if r.completed and r.actual_duration_min]
        avg_daily_min = sum(r.actual_duration_min for r in comp_rows) / len(comp_rows) if comp_rows else 45.0

        # Fatigue trend (slope over last 7 days)
        fatigue_pts = [r.fatigue for d in sorted_dates[-7:] for r in by_date[d] if r.fatigue is not None]
        fatigue_trend = 0.0
        if len(fatigue_pts) >= 3:
            n = len(fatigue_pts)
            xs = list(range(n))
            mx, my = sum(xs) / n, sum(fatigue_pts) / n
            num = sum((xs[i] - mx) * (fatigue_pts[i] - my) for i in range(n))
            den = sum((xs[i] - mx) ** 2 for i in range(n))
            fatigue_trend = num / den if den != 0 else 0.0

        # Deep Work streak & consecutive_same_action counter — FIXED
        # Bug: elif/pass meant consecutive_same_action was never incremented beyond 1
        deep_streak = 0
        consecutive_same_action = 0
        last_action = None
        current_run_action = None
        current_run_count = 0
        for d in reversed(sorted_dates):
            if d == today_str:
                continue
            labels = [r.action_label for r in by_date[d] if r.action_label]
            if not labels:
                break
            daily_action = labels[0]
            # Track deep work streak
            if daily_action == "Deep Work":
                deep_streak += 1
            else:
                # Deep work streak only counts consecutive days
                if deep_streak > 0:
                    pass  # already broken, keep counting done
            # Track consecutive same action
            if current_run_action is None:
                current_run_action = daily_action
                current_run_count = 1
                last_action = daily_action
            elif daily_action == current_run_action:
                current_run_count += 1
            else:
                break  # Streak ended
        consecutive_same_action = current_run_count

        # Recovery debt acceleration heuristic (last 3 days)
        last_3_days = [d for d in sorted_dates if d >= (date.today() - timedelta(days=3)).isoformat()]
        deep_mins_3d = 0
        recov_mins_3d = 0
        for d in last_3_days:
            for r in by_date[d]:
                if r.completed:
                    if r.action_label == "Deep Work":
                        deep_mins_3d += (r.actual_duration_min or 0)
                    elif r.action_label in ("Recovery", "Light Review", "Rest"):
                        recov_mins_3d += (r.actual_duration_min or 0)
        
        recovery_debt_accel = deep_mins_3d - recov_mins_3d

        # Temporal State Memory (EMAs)
        ema_fatigue = 0.0
        ema_mood = 0.5
        alpha_ema = 0.3
        for d in sorted_dates:
            daily_f = [r.fatigue for r in by_date[d] if r.fatigue is not None]
            daily_m = [r.mood for r in by_date[d] if r.mood is not None]
            if daily_f: ema_fatigue = alpha_ema * (sum(daily_f)/len(daily_f)) + (1-alpha_ema) * ema_fatigue
            if daily_m: ema_mood = alpha_ema * (sum(daily_m)/len(daily_m)) + (1-alpha_ema) * ema_mood

        # Tuned Burnout Warning — less false positives, added single-signal override
        burnout_warning = (
            (fatigue_trend > 0.1 and recovery_debt_accel > 0 and ema_fatigue > 0.7)
            or ema_fatigue > 0.88   # Single-signal override: extreme fatigue alone is enough
        )

        # Composite overwork signal — properly normalized to [0, 1]
        # Weighted: hard_days (max ~5) contributes 35%, deep_streak (max ~5) 35%, fatigue_trend 30%
        hard_component  = min(1.0, consecutive_hard / 5.0) * 0.35
        streak_component = min(1.0, deep_streak / 5.0) * 0.35
        trend_component  = min(1.0, max(0.0, fatigue_trend) * 5.0) * 0.30
        overwork = round(hard_component + streak_component + trend_component, 3)

        return {
            "consecutive_hard_days":  float(consecutive_hard),
            "recent_completion_rate": float(completion_rate),
            "avg_satisfaction_7d":    float(avg_sat),
            "fatigue_trend":          float(fatigue_trend),
            "overwork_signal":        float(overwork),
            "deep_work_streak":       float(deep_streak),
            "avg_daily_minutes":      float(avg_daily_min),
            "consecutive_same_action": float(consecutive_same_action),
            "ema_fatigue":            float(ema_fatigue),
            "ema_mood":               float(ema_mood),
            "last_action":            last_action,
            "burnout_warning":        burnout_warning,
            "max_goal_completion":    float(max_completion),
            "weekend_rate":           float(weekend_rate),
            "weekday_rate":           float(weekday_rate),
        }
    except Exception as exc:
        logger.warning("Historical features failed for %s: %s", user.user_id, exc)
        return defaults


def get_user_state_vector(user: User) -> List[float]:
    return [
        user.competence    or 0.7,
        user.mood          or 0.5,
        user.fatigue       or 0.0,
        user.sleep_quality or 0.8,
        user.shock         or 0.0,
        user.sick          or 0.0,
        user.stress        or 0.1,
    ]


def _compute_context_features(state: List[float], user: User, hist: Dict[str, float]) -> Dict[str, float]:
    competence, mood, fatigue, sleep, shock, sick, stress = state[:7]
    energy     = mood * 0.3 + sleep * 0.3 + (1 - fatigue) * 0.25 + (1 - stress) * 0.15
    disruption = max(shock, sick, stress * 0.8)
    readiness  = energy * (1 - disruption * 0.6) * (0.5 + competence * 0.5)
    hour       = datetime.now().hour
    if   9 <= hour <= 11 or 14 <= hour <= 17: time_factor = 1.0
    elif 7 <= hour <=  9 or 11 <= hour <= 14: time_factor = 0.85
    else:                                      time_factor = 0.6
    day_of_week    = datetime.now().weekday()
    
    # Feature 4: Weekend Warrior dynamic factor
    if day_of_week >= 5:
        weekend_rate = hist.get("weekend_rate", 0.5)
        weekday_rate = hist.get("weekday_rate", 0.5)
        if weekend_rate > weekday_rate + 0.05:  # Noticeably better on weekends
            weekend_factor = 1.1
        else:
            weekend_factor = 0.85
    else:
        weekend_factor = 1.0
        
    streak         = user.current_streak or 0
    streak_factor  = min(1.0, 0.5 + streak * 0.05)
    idle_days      = user.days_since_last_login
    inactivity_factor = max(0.3, 1.0 - idle_days * 0.1)
    return {
        "energy": energy, "disruption": disruption, "readiness": readiness,
        "time_factor": time_factor, "weekend_factor": weekend_factor,
        "streak_factor": streak_factor, "inactivity_factor": inactivity_factor,
        "competence": competence, "mood": mood, "fatigue": fatigue,
        "idle_days": float(idle_days),
    }


def _detect_plateau(user: User) -> bool:
    history = getattr(user, "rl_state", None) or []
    if not isinstance(history, list):
        return False
    recent = [float(v) for v in history[-PLATEAU_WINDOW:] if isinstance(v, (int, float))]
    if len(recent) < PLATEAU_WINDOW:
        return False
    mean     = sum(recent) / len(recent)
    variance = sum((r - mean) ** 2 for r in recent) / len(recent)
    return variance < PLATEAU_VAR_THRESHOLD


# ─────────────────────────────────────────────────────────────────
# Feature 1: Circadian Rhythm Sensing
# Learns whether the user is a Morning Lark, Night Owl, etc.
# and nudges Deep Work toward their proven peak performance window.
# ─────────────────────────────────────────────────────────────────

CIRCADIAN_WINDOWS = {
    "morning":   (5,  12),   # 05:00–11:59
    "afternoon": (12, 17),   # 12:00–16:59
    "evening":   (17, 21),   # 17:00–20:59
    "night":     (21, 5),    # 21:00–04:59 (wraps)
}


def _compute_circadian_bonus(user: User, db: Optional[Session], current_action: str) -> float:
    """Return a nudge bonus (+) or penalty (-) based on user's proven peak performance hour.
    
    This is called ONCE per signal computation (outside the action loop) and the
    result is applied only to the Deep Work arm, avoiding redundant DB queries.
    """
    if db is None or current_action != "Deep Work":
        return 0.0
    try:
        from app.models.progress import Progress
        cutoff = (date.today() - timedelta(days=28)).isoformat()  # 4 weeks
        rows = (
            db.query(Progress)
            .filter(
                Progress.user_id == user.user_id,
                Progress.completed == True,
                Progress.completed_at.isnot(None),
                Progress.date >= cutoff,
                Progress.satisfaction.isnot(None),
            )
            .all()
        )
        if len(rows) < 5:  # Need at least 5 data points
            return 0.0

        # Aggregate satisfaction by time window
        window_scores: Dict[str, list] = {k: [] for k in CIRCADIAN_WINDOWS}
        for r in rows:
            completed_at = r.completed_at
            if completed_at.tzinfo is None:
                completed_at = completed_at.replace(tzinfo=timezone.utc)
            hour = completed_at.hour
            for wname, (wstart, wend) in CIRCADIAN_WINDOWS.items():
                if wstart < wend:
                    if wstart <= hour < wend:
                        window_scores[wname].append(r.satisfaction)
                else:  # wraps midnight
                    if hour >= wstart or hour < wend:
                        window_scores[wname].append(r.satisfaction)

        avgs = {
            k: sum(v) / len(v)
            for k, v in window_scores.items()
            if len(v) >= 2
        }
        if len(avgs) < 2:
            return 0.0

        best_window  = max(avgs, key=avgs.__getitem__)
        worst_window = min(avgs, key=avgs.__getitem__)

        current_hour = datetime.now().hour
        current_window = "afternoon"
        for wname, (wstart, wend) in CIRCADIAN_WINDOWS.items():
            if wstart < wend:
                if wstart <= current_hour < wend:
                    current_window = wname
            else:
                if current_hour >= wstart or current_hour < wend:
                    current_window = wname

        if current_window == best_window:
            bonus = 2.0 * (avgs[best_window] - avgs.get(worst_window, 0.5))
            logger.info("Circadian boost +%.2f for Deep Work (peak window: %s)", bonus, best_window)
            return min(3.0, bonus)
        elif current_window == worst_window:
            logger.info("Circadian penalty -1.5 for Deep Work (off-peak: %s)", worst_window)
            return -1.5
        return 0.0
    except Exception as exc:
        logger.debug("Circadian computation failed: %s", exc)
        return 0.0


# ─────────────────────────────────────────────────────────────────
# Feature 2: Dynamic Memory Decay (Exponential Forgetting)
# Applied weekly — allows RL to forget old habits and adapt to
# lifestyle changes quickly without manual resets.
# ─────────────────────────────────────────────────────────────────

MEMORY_DECAY_GAMMA   = 0.88   # 88% retention per week (12% forgotten)
MEMORY_DECAY_INTERVAL = 7     # Apply once per 7 days


def _apply_memory_decay(user: User, params: Dict, db: Session) -> Dict:
    """Decay old bandit params toward the prior (defaults) once per week."""
    try:
        today_str = date.today().isoformat()
        last_decay = getattr(user, "rl_last_decay_date", None)
        if last_decay:
            days_since = (date.today() - date.fromisoformat(last_decay)).days
            if days_since < MEMORY_DECAY_INTERVAL:
                return params  # Not yet time

        defaults = _get_default_params()
        for action in ACTIONS:
            if action not in params:
                continue
            # Blend alpha/beta toward default by (1-gamma) each week
            params[action]["alpha"] = (
                params[action]["alpha"] * MEMORY_DECAY_GAMMA
                + defaults[action]["alpha"] * (1 - MEMORY_DECAY_GAMMA)
            )
            params[action]["beta"] = (
                params[action]["beta"] * MEMORY_DECAY_GAMMA
                + defaults[action]["beta"] * (1 - MEMORY_DECAY_GAMMA)
            )

        user.rl_last_decay_date = today_str
        db.commit()
        logger.info("Memory decay applied for %s (gamma=%.2f)", user.user_id, MEMORY_DECAY_GAMMA)
        return _clamp_params(params)
    except Exception as exc:
        logger.debug("Memory decay skipped: %s", exc)
        return params


# ─────────────────────────────────────────────────────────────────
# Feature 3: Goal-Specific Bandit Params
# Each goal has its own alpha/beta profile so the AI learns that
# the user may handle Coding at 80% intensity but Music at 30%.
# ─────────────────────────────────────────────────────────────────


def _load_goal_bandit_params(user: User, goal_id: Optional[int]) -> Dict:
    """Load per-goal bandit params, blended with global params."""
    if goal_id is None:
        return _get_default_params()
    try:
        stored = (
            user.rl_bandit_params
            if isinstance(user.rl_bandit_params, dict)
            else json.loads(user.rl_bandit_params or "{}")
        )
        goal_key = f"goal_{goal_id}"
        goal_data = stored.get(goal_key, {})
        global_params = _load_bandit_params(user)

        # Blend goal-specific with global (60/40)
        blended = {}
        for action in ACTIONS:
            g_alpha = goal_data.get(action, {}).get("alpha", global_params[action]["alpha"])
            g_beta  = goal_data.get(action, {}).get("beta",  global_params[action]["beta"])
            blended[action] = {
                "alpha": 0.6 * g_alpha + 0.4 * global_params[action]["alpha"],
                "beta":  0.6 * g_beta  + 0.4 * global_params[action]["beta"],
            }
        return _clamp_params(blended)
    except Exception:
        return _get_default_params()


def _save_goal_bandit_params(user: User, goal_id: int, action: str, reward: float) -> None:
    """Update per-goal bandit params after a task completion."""
    try:
        stored = (
            user.rl_bandit_params
            if isinstance(user.rl_bandit_params, dict)
            else json.loads(user.rl_bandit_params or "{}")
        )
        goal_key = f"goal_{goal_id}"
        goal_data = stored.get(goal_key, {a: {"alpha": 2.0, "beta": 2.0} for a in ACTIONS})
        if action in goal_data:
            if reward > 0.5:
                goal_data[action]["alpha"] = min(PARAM_CAP, goal_data[action]["alpha"] + reward)
            else:
                goal_data[action]["beta"]  = min(PARAM_CAP, goal_data[action]["beta"] + (1 - reward))
        stored[goal_key] = goal_data
        user.rl_bandit_params = stored
    except Exception as exc:
        logger.debug("Goal bandit save failed: %s", exc)


# ─────────────────────────────────────────────────────────────────
# Feature 4: 1-Step Lookahead (Transition Dynamics)
# Estimates tomorrow's fatigue based on today's action and pre-
# emptively penalizes actions that will cause tomorrow's burnout.
# ─────────────────────────────────────────────────────────────────

ACTION_FATIGUE_COST: Dict[str, float] = {
    "Deep Work":    0.25,
    "Exploration":  0.15,
    "Normal":       0.10,
    "Light Review": 0.05,
    "Recovery":    -0.15,  # Recovery reduces fatigue
}


def _compute_lookahead_penalty(action: str, current_fatigue: float) -> float:
    """If repeating this action for 3 days pushes fatigue over burnout threshold, penalize it."""
    cost = ACTION_FATIGUE_COST.get(action, 0.0)
    # Feature 5: 3-Day Rolling Lookahead
    predicted_3_day_fatigue = min(1.0, current_fatigue + (3 * cost))
    burnout_threshold = ACTION_FATIGUE_CEILING.get("Normal", 0.85)

    if predicted_3_day_fatigue > burnout_threshold and action in ("Deep Work", "Exploration", "Normal"):
        severity = (predicted_3_day_fatigue - burnout_threshold) * 3.0
        logger.debug("Lookahead penalty %.2f for %s (pred_3d_fatigue=%.2f)", severity, action, predicted_3_day_fatigue)
        return -severity
    return 0.0


def _context_adjusted_params(
    base_params: Dict, ctx: Dict[str, float], hist: Dict[str, float],
    is_cold_start: bool, is_plateau: bool,
    db: Optional[Session] = None, user: Optional[User] = None,
    goal_id: Optional[int] = None,
) -> Dict[str, Dict[str, float]]:
    adjusted = {}
    overwork   = hist["overwork_signal"]
    hard_days  = hist["consecutive_hard_days"]
    comp_rate  = hist["recent_completion_rate"]
    avg_sat    = hist["avg_satisfaction_7d"]
    deep_stk   = hist["deep_work_streak"]
    ema_fatigue = hist.get("ema_fatigue", 0.0)
    cons_action = hist.get("consecutive_same_action", 0.0)
    last_action = hist.get("last_action")

    # Behavioral Archetype heuristic
    is_high_achiever = comp_rate > 0.8 and hard_days >= 3 and ema_fatigue < 0.6
    is_burnout_prone = ema_fatigue > 0.6 and overwork > 0.6 and avg_sat < 0.5
    is_avoidant      = comp_rate < 0.4 and ctx["idle_days"] > 2

    # Feature 3: Circadian bonus computed ONCE outside the loop (avoids 5x DB queries)
    circadian_deep_work_bonus = 0.0
    if user is not None and db is not None:
        circadian_deep_work_bonus = _compute_circadian_bonus(user, db, "Deep Work")

    for action in ACTIONS:
        alpha = base_params[action]["alpha"]
        beta  = base_params[action]["beta"]
        nudge = 0.0
        
        # Feature 3: Goal Completion Proximity Effect
        if action == "Deep Work" and hist.get("max_goal_completion", 0.0) >= 0.8:
            nudge += 2.0
            logger.debug("Finish Line Boost: +2.0 to Deep Work")

        # Action Diversity Entropy Regularization
        if action == last_action and cons_action >= 3:
            nudge -= 1.5 * (cons_action - 2)

        # Archetype Global Adjustments
        if is_burnout_prone and action in ("Deep Work", "Normal"):
            nudge -= 2.5
        if is_high_achiever and action == "Deep Work" and ctx["readiness"] > 0.5:
            nudge += 1.5
        if is_avoidant and action in ("Light Review", "Exploration"):
            nudge += 2.0

        if action == "Recovery":
            if ctx["disruption"] > 0.5 or ctx["fatigue"] > 0.7:
                nudge += 3.0 * ctx["disruption"]
            elif ctx["energy"] < 0.3:
                nudge += 2.0
            if ctx["idle_days"] >= 3:
                nudge += 3.0
            if hard_days >= 3:
                nudge += 3.0 + hard_days * 0.5
            if overwork > 0.8:
                nudge += overwork * 6.0
            elif overwork > 0.6:
                nudge += overwork * 4.0

        elif action == "Light Review":
            if 0.3 < ctx["energy"] < 0.55 and ctx["disruption"] < 0.5:
                nudge += 1.5
            if ctx["mood"] < 0.4:
                nudge += 1.0
            if comp_rate < 0.5:
                nudge += 2.0
            if hard_days >= 2:
                nudge += hard_days * 0.8

        elif action == "Normal":
            if 0.45 < ctx["energy"] < 0.75 and ctx["disruption"] < 0.3:
                nudge += 1.0
            if comp_rate > 0.6 and avg_sat > 0.5:
                nudge += 0.8

        elif action == "Deep Work":
            if ctx["readiness"] > 0.65 and ctx["time_factor"] > 0.8:
                nudge += 2.0 * ctx["readiness"]
            if ctx["competence"] > 0.7 and ctx["mood"] > 0.6 and ctx["fatigue"] < 0.3:
                nudge += 1.5
            if hard_days >= 3:
                nudge -= 2.0 + hard_days * 0.4
            if deep_stk >= 3:
                nudge -= 1.5 * deep_stk
            if overwork > 0.7:
                nudge -= overwork * 3.0
            # Apply pre-computed circadian bonus (already DB-query free)
            nudge += circadian_deep_work_bonus

        elif action == "Exploration":
            if ctx["competence"] < 0.4:
                nudge += 2.0
            elif ctx["mood"] > 0.6 and ctx["energy"] > 0.5:
                nudge += 0.8
            if avg_sat < 0.4:
                nudge += 2.5
            if is_plateau:
                nudge += 3.0

        # Feature 4: 1-Step Lookahead Penalty (applied to all actions)
        nudge += _compute_lookahead_penalty(action, ctx["fatigue"])

        if is_cold_start:
            alpha += 0.5
            beta  += 0.5

        adjusted[action] = {
            "alpha": max(PARAM_FLOOR, alpha + nudge),
            "beta":  max(PARAM_FLOOR, beta),
        }
    return adjusted


def _thompson_sample(params: Dict) -> Tuple[str, Dict[str, float], Dict[str, float]]:
    samples = {
        action: random.betavariate(max(params[action]["alpha"], 0.1), max(params[action]["beta"], 0.1))
        for action in ACTIONS
    }
    
    # Calculate propensity array (normalized samples to approximate P(A))
    total_sample = sum(samples.values())
    propensity = {action: round(samples[action] / max(0.001, total_sample), 3) for action in ACTIONS}
    
    return max(samples, key=samples.__getitem__), samples, propensity


def _calculate_weekly_strain(user: User, db: Optional[Session]) -> Tuple[float, bool]:
    """Calculate 7-day Strain Ratio (Deep Work vs Recovery) and return (strain_ratio, is_deload_state)."""
    if not db:
        return 0.0, False
        
    from app.models.progress import Progress
    from datetime import timedelta
    cutoff = (datetime.now(timezone.utc) - timedelta(days=7)).date().isoformat()
    
    rows = db.query(Progress).filter(
        Progress.user_id == user.user_id,
        Progress.date >= cutoff,
        Progress.completed == True
    ).all()
    
    deep_work_mins = 0
    recovery_mins = 0
    
    for r in rows:
        mins = r.actual_duration_min or 0
        if r.action_label == "Deep Work":
            deep_work_mins += mins
        elif r.action_label in ("Recovery", "Light Review", "Rest"):
            # Reward Hacking Mitigation: 
            # Micro-sessions (<20 mins) do not clear deep strain debt.
            if mins >= 20:
                recovery_mins += mins
            
    strain_ratio = deep_work_mins / max(1.0, float(recovery_mins))
    
    # Deload Hysteresis Logic
    currently_deload = getattr(user, "_tmp_is_deload", False)
    if currently_deload:
        is_deload = strain_ratio > 2.5  # Only exit if ratio drops below 2.5
    else:
        is_deload = strain_ratio > 4.0  # Enter if ratio spikes above 4.0

    return round(strain_ratio, 2), is_deload


def _apply_safety_guard(action: str, fatigue: float, recovery_debt: float = 0.0, is_deload: bool = False, disruption: float = 0.0) -> str:
    """Apply ordered safety rules: disruption → deload → fatigue ceiling → circuit breaker."""

    # Rule 0: Disruption (Sick / Extreme Stress)
    if disruption >= 0.9 and action not in ("Recovery", "Light Review"):
        logger.info("High disruption (%.2f): Forcing %s → Recovery", disruption, action)
        return "Recovery"

    # Rule 1: Deload State — weekly strain too high, cap intensive actions
    if is_deload and action in ("Deep Work", "Exploration", "Normal"):
        if fatigue > 0.6:
            logger.info("Deload+HighFatigue: Downshifting %s → Light Review", action)
            action = "Light Review"
        elif action in ("Deep Work", "Exploration"):
            logger.info("Deload: Downshifting %s → Normal", action)
            action = "Normal"

    # Rule 2: Fatigue Ceiling — action-specific hard cap
    ceiling = ACTION_FATIGUE_CEILING.get(action)
    if ceiling is not None and fatigue > ceiling:
        order = ["Deep Work", "Exploration", "Normal", "Light Review", "Recovery"]
        idx   = order.index(action) if action in order else 0
        for safer in order[idx + 1:]:
            if fatigue <= ACTION_FATIGUE_CEILING.get(safer, 1.0):
                logger.info("Fatigue ceiling: %s → %s (fatigue=%.2f)", action, safer, fatigue)
                return safer
        return "Recovery"

    # Rule 3: Circuit Breaker — lockout Deep Work on critical debt
    if recovery_debt > 15.0 and action == "Deep Work":
        logger.info("Circuit Breaker: debt=%.1f, forcing Recovery", recovery_debt)
        return "Recovery"

    return action


def apply_competence_decay(db: Session, user: User) -> float:
    idle_days = user.days_since_last_login
    if idle_days <= 1:
        return user.competence or 0.7
    decayed = round(max(0.1, min(1.0, (user.competence or 0.7) * (COMPETENCE_DECAY_PER_DAY ** idle_days))), 4)
    if decayed != user.competence:
        user.competence = decayed
        user.competence_updated_at = datetime.now(timezone.utc)
        db.commit()
    return decayed


def get_behavior_signal(
    user: User, state_override: Optional[List[float]] = None, db: Optional[Session] = None,
) -> Dict:
    """Compute adaptive behavior signal using Contextual Thompson Sampling + historical data."""
    try:
        if db is not None and user.days_since_last_login > 1:
            apply_competence_decay(db, user)
            _apply_inactivity_adjustment(user, db)

        state = state_override or get_user_state_vector(user)
        while len(state) < 7: state.append(0.0)
        state = [max(0.0, min(1.0, float(v))) for v in state[:7]]
        competence, mood, fatigue = state[0], state[1], state[2]

        base_params   = _load_bandit_params(user)
        # Feature 2: Apply Memory Decay (once per week)
        if db is not None:
            base_params = _apply_memory_decay(user, base_params, db)
        hist          = _compute_historical_features(user, db)
        ctx           = _compute_context_features(state, user, hist)
        history_count = getattr(user, "rl_history_count", 0) or 0
        is_cold_start = history_count < 10
        is_plateau    = _detect_plateau(user)

        adjusted_params = _context_adjusted_params(
            base_params, ctx, hist, is_cold_start, is_plateau,
            db=db, user=user
        )
        
        # Smart Adaptive Entropy / Repetition Penalty
        # If stuck in a loop > 3 days AND (mood is dropping OR fatigue is rising), penalize.
        # This protects highly disciplined users who are thriving in Deep Work loops.
        consecutive = hist.get("consecutive_same_action", 0)
        last_action = hist.get("last_action")
        if consecutive >= 3 and last_action in adjusted_params:
            if hist.get("ema_mood", 0.5) < 0.6 or hist.get("fatigue_trend", 0.0) > 0:
                logger.info("Adaptive Entropy Triggered: Penalizing %s due to %d day repetition with poor state", last_action, consecutive)
                adjusted_params[last_action]["alpha"] = max(0.1, adjusted_params[last_action]["alpha"] - 2.0)
            
        action_label, samples, propensity = _thompson_sample(adjusted_params)
        recovery_debt = base_params.get("_recovery_debt", 0.0)
        strain_ratio, is_deload = _calculate_weekly_strain(user, db)
        
        action_label = _apply_safety_guard(action_label, fatigue, recovery_debt, is_deload, ctx.get("disruption", 0.0))

        lo, hi = ACTION_INTENSITIES[action_label]
        raw_intensity = ctx["energy"] * ctx["time_factor"] * ctx["weekend_factor"] * ctx["inactivity_factor"]
        if hist["overwork_signal"] > 0.5:
            raw_intensity *= (1.0 - hist["overwork_signal"] * 0.4)
        intensity  = lo + (hi - lo) * max(0.0, min(1.0, raw_intensity))
        intensity += random.uniform(-0.03, 0.03)
        intensity  = round(max(0.0, min(1.0, intensity)), 3)

    except Exception:
        logger.exception("RL signal failed, using fallback")
        state = [0.7, 0.5, 0.0, 0.8, 0.0, 0.0, 0.1]
        competence, mood, fatigue = state[0], state[1], state[2]
        action_label = "Normal"
        intensity    = 0.5
        samples      = {a: 0.5 for a in ACTIONS}
        propensity   = {a: 0.2 for a in ACTIONS}
        history_count, is_cold_start, is_plateau = 0, True, False
        hist = _compute_historical_features(user, None)

    if db is not None:
        try:
            user.current_action_label = action_label
            user.current_intensity    = intensity
            db.commit()
        except Exception:
            pass

    logger.info("RL signal for %s: mode=%s intensity=%.2f overwork=%.2f", user.user_id, action_label, intensity, hist.get("overwork_signal", 0))

    try:
        import os
        import json
        trace_dir = "data"
        os.makedirs(trace_dir, exist_ok=True)
        trace_file = os.path.join(trace_dir, "traces.jsonl")
        prop_arr = locals().get("propensity", {a: 0.2 for a in ACTIONS})
        trace_payload = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "user_id": user.user_id,
            "state_vector": state,
            "debt_level": locals().get("recovery_debt", 0.0),
            "weekly_strain_ratio": locals().get("strain_ratio", 0.0),
            "is_deload": locals().get("is_deload", False),
            "selected_action": action_label,
            "intensity": intensity,
            "propensity_array": prop_arr,
            "burnout_warning": hist.get("burnout_warning", False),
            "weekend_warrior_flag": hist.get("weekend_rate", 0) > hist.get("weekday_rate", 0) + 0.05,
            "max_goal_completion": hist.get("max_goal_completion", 0.0),
        }
        with open(trace_file, "a") as f:
            f.write(json.dumps(trace_payload) + "\n")
    except Exception as e:
        logger.error("Failed to write RL trace: %s", e)

    return {
        "action_label":      action_label,
        "intensity":         intensity,
        "mood":              mood,
        "fatigue":           fatigue,
        "raw_action_vector": state,
        "rl_samples":        {k: round(v, 3) for k, v in samples.items()},
        "rl_updates":        history_count,
        "cold_start":        is_cold_start,
        "plateau_detected":  is_plateau,
        "historical": {
            "consecutive_hard_days":  hist["consecutive_hard_days"],
            "recent_completion_rate": hist["recent_completion_rate"],
            "avg_satisfaction_7d":    hist["avg_satisfaction_7d"],
            "overwork_signal":        hist["overwork_signal"],
            "deep_work_streak":       hist["deep_work_streak"],
        },
        "weekly_strain_ratio": getattr(user, "_tmp_strain_ratio", 0.0) if 'strain_ratio' not in locals() else strain_ratio,
        "is_deload_state": getattr(user, "_tmp_is_deload", False) if 'is_deload' not in locals() else is_deload,
        "propensity_array": propensity,
        "burnout_warning": hist.get("burnout_warning", False)
    }


def get_stable_signal(user: User, db: Optional[Session] = None) -> Dict:
    """Return today's RL signal. Recomputes fresh if stale (from a previous day)."""
    today = datetime.now(timezone.utc).date()
    persisted_label     = getattr(user, "current_action_label", None)
    persisted_intensity = getattr(user, "current_intensity",    None)
    state_updated_at    = getattr(user, "state_updated_at",     None)

    if persisted_label and persisted_intensity is not None and state_updated_at is not None:
        # Fix: always compare as UTC dates to avoid midnight re-computation bugs
        if hasattr(state_updated_at, "date"):
            if state_updated_at.tzinfo is None:
                state_updated_at = state_updated_at.replace(tzinfo=timezone.utc)
            updated_date = state_updated_at.date()
        else:
            updated_date = today
        if updated_date == today:
            state = get_user_state_vector(user)
            hist  = _compute_historical_features(user, db)
            strain_ratio, is_deload = _calculate_weekly_strain(user, db)
            
            return {
                "action_label":      persisted_label,
                "intensity":         persisted_intensity,
                "mood":              state[1],
                "fatigue":           state[2],
                "raw_action_vector": state,
                "rl_samples":        None,
                "rl_updates":        getattr(user, "rl_history_count", 0) or 0,
                "cold_start":        (getattr(user, "rl_history_count", 0) or 0) < 10,
                "plateau_detected":  False,
                "historical": {
                    "consecutive_hard_days":  hist["consecutive_hard_days"],
                    "recent_completion_rate": hist["recent_completion_rate"],
                    "avg_satisfaction_7d":    hist["avg_satisfaction_7d"],
                    "overwork_signal":        hist["overwork_signal"],
                    "deep_work_streak":       hist["deep_work_streak"],
                },
                "weekly_strain_ratio": strain_ratio,
                "is_deload_state":    is_deload,
                # Fixed: include propensity_array so frontend never gets KeyError
                "propensity_array":   {a: round(1.0 / len(ACTIONS), 3) for a in ACTIONS},
                "burnout_warning":    hist.get("burnout_warning", False),
            }
        else:
            # Stale signal from a previous day — recompute fresh
            logger.info("Stale RL signal for %s (from %s), recomputing fresh", user.user_id, updated_date)

    # No persisted signal or it's stale — compute fresh
    return get_behavior_signal(user, db=db)


def refresh_signal(user: User, db: Session) -> Dict:
    user.state_updated_at = datetime.now(timezone.utc)
    db.flush()
    return get_behavior_signal(user, db=db)


def update_reward(db: Session, user: User, action_taken: str, reward: float, goal_id: Optional[int] = None) -> None:
    if action_taken not in ACTIONS:
        logger.warning("Unknown action '%s' — skipping RL update", action_taken)
        return
    reward = max(0.0, min(1.0, float(reward)))
    
    hist = _compute_historical_features(user, db)
    params = _load_bandit_params(user)

    # 1. Delayed Burnout Penalty / Psychological Safety Dampening
    if action_taken == "Deep Work" and hist.get("overwork_signal", 0.0) > 0.7:
        reward *= 0.6  # Immediate safety dampening
        
    # Feature 6: Outlier / Reward Hacking Detection
    if reward > 0.8 and hist.get("ema_fatigue", 0.0) > 0.8:
        # User claims extreme satisfaction despite extreme fatigue
        reward = 0.5 + (reward - 0.5) * 0.5  # Squash toward 0.5
        logger.info("Reward hacked? High satisfaction despite extreme fatigue. Discounted reward to %.2f", reward)
        
    # Longitudinal Reinforcement: Penalize yesterday's intense action if today is crushed
    if hist.get("last_action") == "Deep Work" and hist.get("ema_fatigue", 0.0) > 0.7:
        params["Deep Work"]["alpha"] = max(1.0, params["Deep Work"]["alpha"] - 1.5)
        
    # Recovery Budget Accounting — corrected costs per mode
    debt = params.get("_recovery_debt", 0.0)
    if action_taken == "Deep Work":
        debt += 3.0    # High cost: intense focus drains reserves
    elif action_taken == "Recovery":
        debt = max(0.0, debt - 4.0)  # Slightly higher recovery credit
    elif action_taken in ("Light Review", "Exploration"):
        debt += 0.5    # Light burden
    elif action_taken == "Normal":
        debt += 0.3    # Sustainable mode — was wrongly set to 1.0 before
    # (anything else adds 0)
    params["_recovery_debt"] = debt

    arm = params[action_taken]
    if reward > 0.5:
        arm["alpha"] += reward
    else:
        arm["beta"] += (1.0 - reward)

    # 2. Soft Normalization (Bounded Confidence) to preserve variance
    MAX_CONFIDENCE = 20.0
    for action in ACTIONS:
        total = params[action]["alpha"] + params[action]["beta"]
        if total > MAX_CONFIDENCE:
            scale = MAX_CONFIDENCE / total
            params[action]["alpha"] = max(1.0, 1.0 + (params[action]["alpha"] - 1.0) * scale)
            params[action]["beta"]  = max(1.0, 1.0 + (params[action]["beta"] - 1.0) * scale)

    # Feature 3: Update Goal-Specific Bandit if goal_id provided
    if goal_id is not None:
        _save_goal_bandit_params(user, goal_id, action_taken, reward)

    params["_last_updated"] = datetime.now(timezone.utc).isoformat()
    params = _clamp_params(params)
    user.rl_bandit_params = json.dumps(params)
    user.rl_history_count = (user.rl_history_count or 0) + 1

    history = getattr(user, "rl_state", None) or []
    if not isinstance(history, list): history = []
    history.append(round(reward, 4))
    user.rl_state = history[-50:]
    db.commit()

    logger.info("RL update for %s: action=%s goal=%s reward=%.2f -> alpha=%.2f beta=%.2f [total %d]",
        user.user_id, action_taken, goal_id, reward,
        params[action_taken]["alpha"], params[action_taken]["beta"], user.rl_history_count)


def _apply_inactivity_adjustment(user: User, db: Session) -> None:
    """If user was inactive 3+ days, decay bandit params toward defaults to prevent overfit."""
    idle_days = user.days_since_last_login
    if idle_days < 3:
        return

    params   = _load_bandit_params(user)
    defaults = _get_default_params()
    # Blend toward defaults: 20% per idle day beyond 2, capped at 80%
    blend = min(0.80, (idle_days - 2) * 0.20)

    for action in ACTIONS:
        params[action]["alpha"] = params[action]["alpha"] * (1 - blend) + defaults[action]["alpha"] * blend
        params[action]["beta"]  = params[action]["beta"]  * (1 - blend) + defaults[action]["beta"]  * blend

    params = _clamp_params(params)
    user.rl_bandit_params = json.dumps(params)
    db.commit()
    logger.info("Inactivity adjustment for %s: %d idle days, blend=%.0f%%", user.user_id, idle_days, blend * 100)


def reset_daily_state(db: Session, user: User) -> None:
    """Clear persisted RL decision so next GET /api/plan recomputes fresh.
    Called after day-close or manual completion."""
    user.current_action_label = None
    user.current_intensity    = None
    user.state_updated_at     = None
    db.commit()
    logger.info("Daily state reset for %s", user.user_id)


def get_next_day_signal(user: User, db: Optional[Session] = None) -> Dict:
    """Compute a preview signal for tomorrow (doesn't persist).
    Used to show 'tomorrow's plan preview' after day-close."""
    try:
        if db is not None and user.days_since_last_login > 1:
            apply_competence_decay(db, user)

        state = get_user_state_vector(user)
        while len(state) < 7: state.append(0.0)
        state = state[:7]

        # Simulate slightly recovered state for tomorrow
        # Reduce fatigue by 30%, boost mood slightly
        state[1] = min(1.0, state[1] + 0.05)   # mood bump
        state[2] = max(0.0, state[2] * 0.70)    # fatigue recovery

        base_params   = _load_bandit_params(user)
        hist          = _compute_historical_features(user, db)
        ctx           = _compute_context_features(state, user, hist)
        history_count = getattr(user, "rl_history_count", 0) or 0
        is_cold_start = history_count < 10
        is_plateau    = _detect_plateau(user)

        adjusted_params = _context_adjusted_params(base_params, ctx, hist, is_cold_start, is_plateau)
        action_label, samples, propensity = _thompson_sample(adjusted_params)  # Fixed: unpack 3 values
        action_label = _apply_safety_guard(action_label, state[2])

        lo, hi = ACTION_INTENSITIES[action_label]
        raw_intensity = ctx["energy"] * ctx["time_factor"] * ctx["weekend_factor"] * ctx["inactivity_factor"]
        if hist["overwork_signal"] > 0.5:
            raw_intensity *= (1.0 - hist["overwork_signal"] * 0.4)
        intensity  = lo + (hi - lo) * max(0.0, min(1.0, raw_intensity))
        intensity  = round(max(0.0, min(1.0, intensity)), 3)

        return {
            "action_label": action_label,
            "intensity":    intensity,
            "preview":      True,
        }
    except Exception:
        logger.exception("Next-day signal preview failed")
        return {"action_label": "Normal", "intensity": 0.5, "preview": True}