basic_tools.py

"""
Mental Health Counselor — Tool-Calling Demo with Math Tools (Groq)
===================================================================
Adapted to the nested wellbeing-frame schema (categories of {value, description}).
The LLM decides WHAT data it needs and WHEN to fetch it via tool calls.

Setup:
    pip install groq
    export GROQ_API_KEY=gsk_...

Run:
    python basic_tools.py                         # uses data.json by default
    python basic_tools.py --data path/to/file.json
    python basic_tools.py --quiet                 # hide tool call logs
"""

import json
import os
import sys
import math
import argparse
from groq import Groq
from typing import Optional


# ─────────────────────────────────────────────────────────────────────────────
# SCHEMA-AWARE METRIC DEFINITIONS
# Each frame nests categories → sub-dimensions → {value: 0-5, description}.
# Polarity: "+" higher = better wellbeing, "-" higher = worse, "0" neutral.
# ─────────────────────────────────────────────────────────────────────────────

METRIC_CATEGORIES = {
    "emotions": {
        "calm_neutral":            "+",
        "happy_positive":          "+",
        "anxious_worried":         "-",
        "sad_low":                 "-",
        "angry_irritable":         "-",
        "lonely":                  "-",
        "overwhelmed":             "-",
        "numb_emotionally_flat":   "-",
    },
    "stresses": {
        "work_academic":              "-",
        "relationship":               "-",
        "health_related":             "-",
        "financial":                  "-",
        "time_pressure_overload":     "-",
        "uncertainty_future_anxiety": "-",
        "internal_pressure":          "-",
        "low_manageable":             "+",   # high score = stress feels manageable
    },
    "cognitive_patterns": {
        "balanced_realistic":       "+",
        "rumination":               "-",
        "catastrophizing":          "-",
        "black_and_white":          "-",
        "self_critical":            "-",
        "helplessness_low_control": "-",
        "overanalysis_indecision":  "-",
        "positive_reframing":       "+",
    },
    "motivation_values": {
        "highly_motivated_goal_driven": "+",
        "moderate_motivation":          "+",
        "low_motivation_disengaged":    "-",
        "anhedonia_loss_of_interest":   "-",
        "purpose_driven":               "+",
        "value_conflict":               "-",
        "directionless_unclear_goals":  "-",
    },
    "sleep": {
        "restful_healthy":      "+",
        "mild_disturbance":     "-",
        "insufficient_sleep":   "-",
        "insomnia":             "-",
        "irregular_schedule":   "-",
        "oversleeping_fatigue": "-",
    },
    "energy": {
        "high_energized":    "+",
        "stable_normal":     "+",
        "low_tired":         "-",
        "exhausted_drained": "-",
        "fluctuating":       "-",
        "restless_wired":    "-",
    },
    "personality": {
        "optimistic":                "+",
        "pessimistic":               "-",
        "self_confident":            "+",
        "self_doubting":             "-",
        "emotionally_reactive":      "-",
        "emotionally_stable":        "+",
        "introverted":               "0",   # neutral disposition
        "socially_expressive":       "+",
        "conscientious_disciplined": "+",
        "avoidant_tendency":         "-",
    },
    "habits": {
        "structured_healthy_routines": "+",
        "productive_habits":           "+",
        "inconsistent_routines":       "-",
        "procrastination":             "-",
        "avoidance_behaviors":         "-",
        "compulsive_behaviors":        "-",
        "self_care_present":           "+",
        "self_care_neglect":           "-",
    },
    "social": {
        "strong_support_system": "+",
        "moderate_support":      "+",
        "limited_support":       "-",
        "socially_isolated":     "-",
        "active_engagement":     "+",
        "relationship_conflict": "-",
        "help_seeking_behavior": "+",
        "withdrawing":           "-",
    },
}

# Flatten into dotted paths
ALL_METRICS = [f"{cat}.{sub}" for cat, subs in METRIC_CATEGORIES.items() for sub in subs]
POLARITY    = {f"{cat}.{sub}": pol for cat, subs in METRIC_CATEGORIES.items() for sub, pol in subs.items()}
NEGATIVE_POLARITY = {p for p, pol in POLARITY.items() if pol == "-"}

# Top-level numerics (not under a category)
TOP_LEVEL_NUMERIC = {"exercise_minutes": "+"}

# Value ranges
VALUE_MAX    = 5      # every sub-dim value
EXERCISE_MAX = 120    # minutes

# Category weights for composite wellbeing score (sum to 1.0)
CATEGORY_WEIGHTS = {
    "emotions":           0.20,
    "stresses":           0.15,
    "cognitive_patterns": 0.15,
    "sleep":              0.15,
    "motivation_values":  0.10,
    "energy":             0.10,
    "habits":             0.05,
    "social":             0.05,
    "personality":        0.05,
}


# ─────────────────────────────────────────────────────────────────────────────
# MATH HELPER FUNCTIONS
# ─────────────────────────────────────────────────────────────────────────────

def _safe_mean(values: list) -> float:
    return sum(values) / len(values) if values else 0.0


def _safe_std(values: list) -> float:
    if len(values) < 2:
        return 0.0
    mu = _safe_mean(values)
    return math.sqrt(sum((x - mu) ** 2 for x in values) / (len(values) - 1))


def _safe_median(values: list) -> float:
    if not values:
        return 0.0
    s = sorted(values)
    n = len(s)
    if n % 2 == 1:
        return s[n // 2]
    return (s[n // 2 - 1] + s[n // 2]) / 2


def _pearson_r(xs: list, ys: list) -> Optional[float]:
    n = min(len(xs), len(ys))
    if n < 3:
        return None
    xs, ys = xs[:n], ys[:n]
    mx, my = _safe_mean(xs), _safe_mean(ys)
    num = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    dx = math.sqrt(sum((x - mx) ** 2 for x in xs))
    dy = math.sqrt(sum((y - my) ** 2 for y in ys))
    if dx * dy == 0:
        return None
    return round(num / (dx * dy), 4)


def _linear_regression(values: list) -> tuple:
    n = len(values)
    if n < 2:
        return (0.0, values[0] if values else 0.0)
    x_mean = (n - 1) / 2
    y_mean = _safe_mean(values)
    num = sum((i - x_mean) * (v - y_mean) for i, v in enumerate(values))
    den = sum((i - x_mean) ** 2 for i in range(n))
    slope = num / den if den != 0 else 0.0
    intercept = y_mean - slope * x_mean
    return (round(slope, 4), round(intercept, 4))


# ─────────────────────────────────────────────────────────────────────────────
# SCHEMA-AWARE EXTRACTION
# ─────────────────────────────────────────────────────────────────────────────

def _get_raw_value(frame: dict, path: str):
    """Walk path; return numeric leaf. Handles {value, description} unwrapping."""
    if "." not in path:
        v = frame.get(path)
        return v if isinstance(v, (int, float)) else None
    obj = frame
    for p in path.split("."):
        if not isinstance(obj, dict):
            return None
        obj = obj.get(p)
        if obj is None:
            return None
    if isinstance(obj, dict) and isinstance(obj.get("value"), (int, float)):
        return obj["value"]
    if isinstance(obj, (int, float)):
        return obj
    return None


def _compute_composite_score(frame: dict) -> Optional[float]:
    """0-100 composite per frame using category weights."""
    cat_total, weight_total = 0.0, 0.0
    for cat, subs in METRIC_CATEGORIES.items():
        contribs = []
        for sub, pol in subs.items():
            v = _get_raw_value(frame, f"{cat}.{sub}")
            if v is None or pol == "0":
                continue
            normed = v / VALUE_MAX
            if pol == "-":
                normed = 1.0 - normed
            contribs.append(normed)
        if contribs:
            wt = CATEGORY_WEIGHTS.get(cat, 0.0)
            cat_total    += _safe_mean(contribs) * wt
            weight_total += wt
    if weight_total == 0:
        return None
    return round((cat_total / weight_total) * 100, 1)


def _get_value(frame: dict, path: str):
    """Same as _get_raw_value, plus the virtual 'wellbeing_score' path."""
    if path == "wellbeing_score":
        return _compute_composite_score(frame)
    return _get_raw_value(frame, path)


def _extract_series(frames: list, path: str) -> list:
    out = []
    for f in frames:
        v = _get_value(f, path)
        if v is not None:
            out.append(v)
    return out


def _is_negative(path: str) -> bool:
    if path in NEGATIVE_POLARITY:
        return True
    if path == "wellbeing_score":
        return False
    if path in TOP_LEVEL_NUMERIC:
        return TOP_LEVEL_NUMERIC[path] == "-"
    return False


# ─────────────────────────────────────────────────────────────────────────────
# LOAD USER DATA FROM FILE
# ─────────────────────────────────────────────────────────────────────────────

def load_user_data(path: str) -> dict:
    with open(path, "r") as f:
        return json.load(f)


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 1: CORE DATA ACCESS TOOLS
# ─────────────────────────────────────────────────────────────────────────────

def get_personal_profile(user_data: dict) -> dict:
    return user_data["personal_memory"]


def get_recent_chat_history(user_data: dict, last_n: int = 6) -> list:
    history = user_data["recent_chat_memory"]
    return history[-min(last_n, 20):]


def get_wellbeing_trend(user_data: dict, last_n_weeks: int = None, fields: list = None) -> list:
    """Returns weekly frames. If `fields` is given (list of metric paths), each entry is reduced to
    {week, <each requested path>: value}."""
    frames = user_data["wellbeing_frames"]
    if last_n_weeks:
        frames = frames[-last_n_weeks:]
    if fields:
        reduced = []
        for frame in frames:
            row = {"week": frame.get("week")}
            for path in fields:
                row[path] = _get_value(frame, path)
            reduced.append(row)
        return reduced
    return frames


def get_wellbeing_snapshot(user_data: dict) -> dict:
    """Latest frame plus a flat numeric summary that's easy for the LLM to read."""
    frames = user_data["wellbeing_frames"]
    if not frames:
        return {"error": "No wellbeing frames"}
    latest = frames[-1]
    summary = {
        "week": latest.get("week"),
        "wellbeing_score": _compute_composite_score(latest),
        "coping_mechanisms": latest.get("coping_mechanisms", []),
        "exercise_minutes": latest.get("exercise_minutes"),
        "by_category": {},
    }
    for cat, subs in METRIC_CATEGORIES.items():
        cat_obj = latest.get(cat) or {}
        summary["by_category"][cat] = {
            sub: (cat_obj.get(sub) or {}).get("value") for sub in subs
        }
    return summary


def compare_wellbeing_weeks(user_data: dict, week_a: int, week_b: int) -> dict:
    frames = {f["week"]: f for f in user_data["wellbeing_frames"] if "week" in f}
    fa, fb = frames.get(week_a), frames.get(week_b)
    if not fa or not fb:
        return {"error": f"Week not found. Available weeks: {sorted(frames.keys())}"}

    diff = {}
    for path in ALL_METRICS + list(TOP_LEVEL_NUMERIC.keys()):
        va, vb = _get_value(fa, path), _get_value(fb, path)
        if va is None or vb is None or va == vb:
            continue
        diff[path] = {f"week_{week_a}": va, f"week_{week_b}": vb, "change": vb - va}

    composite_a = _compute_composite_score(fa)
    composite_b = _compute_composite_score(fb)
    top = sorted(diff.items(), key=lambda kv: abs(kv[1]["change"]), reverse=True)[:10]

    return {
        "wellbeing_score": {
            f"week_{week_a}": composite_a,
            f"week_{week_b}": composite_b,
            "change": (composite_b - composite_a) if (composite_a is not None and composite_b is not None) else None,
        },
        "top_changes": dict(top),
        "summary": f"Compared weeks {week_a} and {week_b}",
    }


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 2: ANALYTICS & MATH TOOLS (A1-A12)
# ─────────────────────────────────────────────────────────────────────────────

# A1. METRIC SUMMARY STATS
def get_metric_summary_stats(user_data: dict, metrics: list = None) -> dict:
    """Mean, median, std-dev, min, max for each metric path."""
    frames = user_data["wellbeing_frames"]
    targets = metrics or (ALL_METRICS + ["wellbeing_score", "exercise_minutes"])
    results = {}

    for path in targets:
        series = _extract_series(frames, path)
        if not series:
            continue
        results[path] = {
            "mean":    round(_safe_mean(series), 2),
            "median":  round(_safe_median(series), 2),
            "std":     round(_safe_std(series), 2),
            "min":     min(series),
            "max":     max(series),
            "latest":  series[-1],
            "n_weeks": len(series),
        }
    return results


# A2. DETECT PATTERN SHIFTS
def detect_pattern_shift(user_data: dict, window: int = 3, threshold: float = 1.0, top_n: int = 15) -> list:
    """Detects significant changes between two time windows. Returns top_n by |z|."""
    frames = user_data["wellbeing_frames"]
    if len(frames) < window * 2:
        mid = len(frames) // 2
        older, recent = frames[:mid], frames[mid:]
    else:
        older, recent = frames[-(window * 2):-window], frames[-window:]

    shifts = []
    for path in ALL_METRICS + ["wellbeing_score", "exercise_minutes"]:
        old_vals = _extract_series(older, path)
        new_vals = _extract_series(recent, path)
        if not old_vals or not new_vals:
            continue

        old_mean, new_mean = _safe_mean(old_vals), _safe_mean(new_vals)
        change   = new_mean - old_mean
        all_vals = _extract_series(frames, path)
        std = _safe_std(all_vals) or 1.0
        z = change / std

        direction = (
            ("worsening" if change > 0 else "improving")
            if _is_negative(path)
            else ("improving" if change > 0 else "worsening")
        )
        severity = "major" if abs(z) >= 2.0 else ("notable" if abs(z) >= 1.0 else "minor")

        shifts.append({
            "metric":    path,
            "old_mean":  round(old_mean, 2),
            "new_mean":  round(new_mean, 2),
            "change":    round(change, 2),
            "z_score":   round(z, 2),
            "direction": direction,
            "severity":  severity,
        })

    shifts.sort(key=lambda s: abs(s["z_score"]), reverse=True)
    return shifts[:top_n]


# A3. CORRELATION
def get_correlated_factors(user_data: dict, target_metric: str = "wellbeing_score", min_abs_r: float = 0.3, top_n: int = 15) -> list:
    """Find metric paths correlated with target_metric."""
    frames = user_data["wellbeing_frames"]
    target_series = _extract_series(frames, target_metric)
    if len(target_series) < 3:
        return [{"error": "Not enough data points"}]

    candidates = ALL_METRICS + ["exercise_minutes"]
    if target_metric in candidates:
        candidates = [c for c in candidates if c != target_metric]

    results = []
    for path in candidates:
        other = _extract_series(frames, path)
        r = _pearson_r(target_series, other)
        if r is not None and abs(r) >= min_abs_r:
            results.append({
                "metric":    path,
                "pearson_r": r,
                "strength":  "strong" if abs(r) >= 0.7 else ("moderate" if abs(r) >= 0.5 else "weak"),
            })

    results.sort(key=lambda x: abs(x["pearson_r"]), reverse=True)
    return results[:top_n]


# A4. VOLATILITY / STABILITY
def compute_volatility(user_data: dict, last_n_weeks: int = None, top_n: int = 15) -> dict:
    frames = user_data["wellbeing_frames"]
    if last_n_weeks:
        frames = frames[-last_n_weeks:]

    rows = []
    for path in ALL_METRICS + ["wellbeing_score", "exercise_minutes"]:
        series = _extract_series(frames, path)
        if len(series) < 2:
            continue
        mean = _safe_mean(series)
        std  = _safe_std(series)
        cv   = round(std / mean, 3) if mean != 0 else 0.0
        wow_changes = [abs(series[i] - series[i - 1]) for i in range(1, len(series))]
        avg_wow = round(_safe_mean(wow_changes), 2)
        max_wow = round(max(wow_changes), 2)
        stability = "stable" if cv < 0.15 else ("moderate" if cv < 0.30 else "volatile")

        rows.append({
            "metric":         path,
            "cv":             cv,
            "std":            round(std, 2),
            "avg_wow_change": avg_wow,
            "max_wow_change": max_wow,
            "stability":      stability,
        })

    rows.sort(key=lambda r: r["cv"], reverse=True)
    return {r["metric"]: r for r in rows[:top_n]}


# A5. ANOMALY DETECTION
def detect_anomalies(user_data: dict, z_threshold: float = 1.5) -> list:
    frames = user_data["wellbeing_frames"]
    stats = {}
    for path in ALL_METRICS + ["wellbeing_score", "exercise_minutes"]:
        series = _extract_series(frames, path)
        if len(series) >= 3:
            stats[path] = (_safe_mean(series), _safe_std(series))

    anomalies = []
    for frame in frames:
        week = frame.get("week", "?")
        for path, (mu, std) in stats.items():
            val = _get_value(frame, path)
            if val is None or std == 0:
                continue
            z = (val - mu) / std
            if abs(z) < z_threshold:
                continue
            if _is_negative(path):
                direction = "unusually high (worse)" if z > 0 else "unusually low (better)"
            else:
                direction = "unusually high (better)" if z > 0 else "unusually low (worse)"
            anomalies.append({
                "week":      week,
                "metric":    path,
                "value":     val,
                "mean":      round(mu, 2),
                "z_score":   round(z, 2),
                "direction": direction,
            })

    anomalies.sort(key=lambda a: (a["week"], abs(a["z_score"])), reverse=True)
    return anomalies


# A6. TREND PREDICTION
def predict_trend(user_data: dict, metric: str = "wellbeing_score", forecast_weeks: int = 2) -> dict:
    frames = user_data["wellbeing_frames"]
    series = _extract_series(frames, metric)
    if len(series) < 3:
        return {"error": f"Need at least 3 data points for '{metric}'."}

    slope, intercept = _linear_regression(series)
    n = len(series)
    projections = [
        {"weeks_from_now": i, "projected_value": round(intercept + slope * (n - 1 + i), 2)}
        for i in range(1, forecast_weeks + 1)
    ]

    if _is_negative(metric):
        direction = "worsening" if slope > 0.05 else ("improving" if slope < -0.05 else "flat")
    else:
        direction = "improving" if slope > 0.05 else ("worsening" if slope < -0.05 else "flat")

    return {
        "metric":      metric,
        "data_points": n,
        "slope":       slope,
        "direction":   direction,
        "current":     series[-1],
        "projections": projections,
    }


# A7. COMPOSITE WELLBEING SCORE
def compute_wellbeing_composite(user_data: dict, weights: dict = None) -> dict:
    """Per-week 0-100 wellbeing score using category weights (override via `weights`)."""
    frames = user_data["wellbeing_frames"]
    weights = weights or CATEGORY_WEIGHTS

    scored_weeks = []
    for frame in frames:
        cat_total, w_total = 0.0, 0.0
        for cat, subs in METRIC_CATEGORIES.items():
            contribs = []
            for sub, pol in subs.items():
                v = _get_raw_value(frame, f"{cat}.{sub}")
                if v is None or pol == "0":
                    continue
                normed = v / VALUE_MAX
                if pol == "-":
                    normed = 1.0 - normed
                contribs.append(normed)
            if contribs:
                wt = weights.get(cat, 0.0)
                cat_total += _safe_mean(contribs) * wt
                w_total   += wt
        score = round((cat_total / w_total) * 100, 1) if w_total > 0 else None
        scored_weeks.append({"week": frame.get("week"), "score": score})

    scores = [s["score"] for s in scored_weeks if s["score"] is not None]
    slope, _ = _linear_regression(scores) if len(scores) >= 2 else (0, 0)
    direction = "improving" if slope > 0.3 else ("declining" if slope < -0.3 else "stable")

    return {
        "weekly_scores": scored_weeks,
        "current_score": scored_weeks[-1]["score"] if scored_weeks else None,
        "average_score": round(_safe_mean(scores), 1) if scores else None,
        "trend":         direction,
    }


# A8. RATE OF CHANGE
def get_rate_of_change(user_data: dict, last_n_weeks: int = 4, top_n: int = 15) -> dict:
    frames = user_data["wellbeing_frames"][-(last_n_weeks + 1):]

    rows = []
    for path in ALL_METRICS + ["wellbeing_score", "exercise_minutes"]:
        series = _extract_series(frames, path)
        if len(series) < 2:
            continue
        deltas = [round(series[i] - series[i - 1], 2) for i in range(1, len(series))]
        if len(deltas) >= 2:
            accel = round(deltas[-1] - deltas[-2], 2)
            accel_label = "accelerating" if abs(accel) > 0.5 else "steady"
        else:
            accel, accel_label = 0.0, "insufficient"

        latest_delta = deltas[-1]
        if _is_negative(path):
            direction = "worsening" if latest_delta > 0 else ("improving" if latest_delta < 0 else "unchanged")
        else:
            direction = "improving" if latest_delta > 0 else ("worsening" if latest_delta < 0 else "unchanged")

        rows.append({
            "metric":      path,
            "deltas":      deltas,
            "latest":      latest_delta,
            "direction":   direction,
            "accel":       accel,
            "accel_label": accel_label,
        })

    rows.sort(key=lambda r: abs(r["latest"]), reverse=True)
    return {r["metric"]: r for r in rows[:top_n]}


# A9. CLUSTER SIMILAR WEEKS
def cluster_similar_weeks(user_data: dict) -> dict:
    frames = user_data["wellbeing_frames"]
    if len(frames) < 2:
        return {"error": "Need at least 2 weeks"}

    paths = ALL_METRICS
    vectors = []
    for frame in frames:
        vec = []
        for p in paths:
            v = _get_raw_value(frame, p)
            if v is None:
                vec.append(0.5)
                continue
            n = v / VALUE_MAX
            if _is_negative(p):
                n = 1.0 - n
            vec.append(n)
        vectors.append(vec)

    def dist(a, b):
        return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))

    threshold = 0.3 * math.sqrt(len(paths))
    clusters: list = []
    for i in range(len(frames)):
        placed = False
        for cluster in clusters:
            if any(dist(vectors[i], vectors[j]) < threshold for j in cluster):
                cluster.add(i)
                placed = True
                break
        if not placed:
            clusters.append({i})

    return {
        f"cluster_{ci}": {
            "weeks": [frames[i].get("week", i + 1) for i in sorted(c)],
            "size":  len(c),
        }
        for ci, c in enumerate(clusters)
    }


# A10. LAGGED CORRELATION
def get_lagged_correlation(user_data: dict, metric_a: str = "stresses.work_academic", metric_b: str = "wellbeing_score", max_lag: int = 3) -> dict:
    frames = user_data["wellbeing_frames"]
    sa = _extract_series(frames, metric_a)
    sb = _extract_series(frames, metric_b)
    n = min(len(sa), len(sb))
    if n < 4:
        return {"error": "Need at least 4 weeks"}

    results, best_r, best_lag = [], 0.0, 0
    for lag in range(-max_lag, max_lag + 1):
        if lag >= 0:
            a_slice = sa[: n - lag] if lag > 0 else sa[:n]
            b_slice = sb[lag:n]
        else:
            a_slice = sa[-lag:n]
            b_slice = sb[: n + lag]
        r = _pearson_r(a_slice, b_slice)
        if r is None:
            continue
        results.append({"lag": lag, "r": r})
        if abs(r) > abs(best_r):
            best_r, best_lag = r, lag

    return {
        "metric_a":      metric_a,
        "metric_b":      metric_b,
        "results":       results,
        "strongest_lag": best_lag,
        "strongest_r":   best_r,
    }


# A11. STREAKS
def get_streaks(user_data: dict, top_n: int = 15) -> dict:
    frames = user_data["wellbeing_frames"]
    rows = []

    for path in ALL_METRICS + ["wellbeing_score", "exercise_minutes"]:
        series = _extract_series(frames, path)
        if len(series) < 2:
            continue

        streak_type = None
        streak_len  = 0
        for i in range(len(series) - 1, 0, -1):
            diff = series[i] - series[i - 1]
            if diff > 0:
                d = "worsening" if _is_negative(path) else "improving"
            elif diff < 0:
                d = "improving" if _is_negative(path) else "worsening"
            else:
                d = "flat"
            if streak_type is None:
                streak_type, streak_len = d, 1
            elif d == streak_type:
                streak_len += 1
            else:
                break

        avg = _safe_mean(series)
        above_avg_streak = 0
        for v in reversed(series):
            if v >= avg:
                above_avg_streak += 1
            else:
                break

        rows.append({
            "metric":          path,
            "streak_type":     streak_type or "flat",
            "streak_weeks":    streak_len,
            "avg":             round(avg, 2),
            "weeks_above_avg": above_avg_streak,
            "latest":          series[-1],
        })

    rows.sort(key=lambda r: r["streak_weeks"], reverse=True)
    return {r["metric"]: r for r in rows[:top_n]}


# A12. COPING EFFECTIVENESS
def estimate_coping_effectiveness(user_data: dict) -> list:
    """Compares weeks where each coping mechanism is active vs. not, by composite wellbeing."""
    frames = user_data["wellbeing_frames"]
    if len(frames) < 3:
        return [{"error": "Need at least 3 weeks"}]

    all_coping = set()
    for f in frames:
        mechs = f.get("coping_mechanisms") or []
        if isinstance(mechs, str):
            mechs = [mechs]
        all_coping.update(mechs)

    if not all_coping:
        return [{"error": "No coping mechanism data"}]

    results = []
    for mech in sorted(all_coping):
        with_scores, without_scores = [], []
        for f in frames:
            active = f.get("coping_mechanisms") or []
            if isinstance(active, str):
                active = [active]
            score = _compute_composite_score(f)
            if score is None:
                continue
            (with_scores if mech in active else without_scores).append(score)

        if not with_scores:
            continue
        diff = round(_safe_mean(with_scores) - _safe_mean(without_scores), 2)
        results.append({
            "mechanism":     mech,
            "weeks_used":    len(with_scores),
            "avg_with":      round(_safe_mean(with_scores), 1),
            "avg_without":   round(_safe_mean(without_scores), 1) if without_scores else None,
            "effectiveness": diff,
            "verdict":       "helpful" if diff > 3 else ("neutral" if diff > -3 else "may not help"),
        })

    results.sort(key=lambda x: x["effectiveness"], reverse=True)
    return results


# ─────────────────────────────────────────────────────────────────────────────
# TOOL DEFINITIONS FOR LLM
# ─────────────────────────────────────────────────────────────────────────────

_METRIC_PATH_HINT = (
    "Dotted path into a wellbeing frame, e.g. 'emotions.anxious_worried', "
    "'stresses.work_academic', 'sleep.insufficient_sleep'. "
    "Special values: 'wellbeing_score' (0-100 composite), 'exercise_minutes'."
)

TOOLS = [
    # Core
    {"type": "function", "function": {"name": "get_personal_profile", "description": "Returns user's personal info: name, age, gender, occupation, traits, etc.", "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {"name": "get_recent_chat_history", "description": "Returns last N chat messages with sentiment scores.", "parameters": {"type": "object", "properties": {"last_n": {"type": "integer", "description": "Number of messages (default 6)", "default": 6}}, "required": []}}},
    {"type": "function", "function": {"name": "get_wellbeing_trend", "description": "Returns weekly wellbeing snapshots over time. If `fields` (list of metric paths) is given, each row is reduced to {week, <field>: value}.", "parameters": {"type": "object", "properties": {"last_n_weeks": {"type": "integer"}, "fields": {"type": "array", "items": {"type": "string", "description": _METRIC_PATH_HINT}}}, "required": []}}},
    {"type": "function", "function": {"name": "get_wellbeing_snapshot", "description": "Latest week's full snapshot (composite score + per-category sub-dimension values).", "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {"name": "compare_wellbeing_weeks", "description": "Compare two specific weeks; returns top 10 changed metrics + composite delta.", "parameters": {"type": "object", "properties": {"week_a": {"type": "integer"}, "week_b": {"type": "integer"}}, "required": ["week_a", "week_b"]}}},

    # Analytics
    {"type": "function", "function": {"name": "get_metric_summary_stats", "description": "Mean, median, std, min, max, latest for each metric path.", "parameters": {"type": "object", "properties": {"metrics": {"type": "array", "items": {"type": "string", "description": _METRIC_PATH_HINT}}}, "required": []}}},
    {"type": "function", "function": {"name": "detect_pattern_shift", "description": "Top-N significant changes between recent vs. older windows (z-score sorted).", "parameters": {"type": "object", "properties": {"window": {"type": "integer", "default": 3}, "threshold": {"type": "number", "default": 1.0}, "top_n": {"type": "integer", "default": 15}}, "required": []}}},
    {"type": "function", "function": {"name": "get_correlated_factors", "description": "Find metric paths correlated with a target metric.", "parameters": {"type": "object", "properties": {"target_metric": {"type": "string", "default": "wellbeing_score", "description": _METRIC_PATH_HINT}, "min_abs_r": {"type": "number", "default": 0.3}, "top_n": {"type": "integer", "default": 15}}, "required": []}}},
    {"type": "function", "function": {"name": "compute_volatility", "description": "Top-N most volatile metrics by coefficient of variation.", "parameters": {"type": "object", "properties": {"last_n_weeks": {"type": "integer"}, "top_n": {"type": "integer", "default": 15}}, "required": []}}},
    {"type": "function", "function": {"name": "detect_anomalies", "description": "Per-week outlier values across all metrics.", "parameters": {"type": "object", "properties": {"z_threshold": {"type": "number", "default": 1.5}}, "required": []}}},
    {"type": "function", "function": {"name": "predict_trend", "description": "Linear-regression projection for a metric path.", "parameters": {"type": "object", "properties": {"metric": {"type": "string", "default": "wellbeing_score", "description": _METRIC_PATH_HINT}, "forecast_weeks": {"type": "integer", "default": 2}}, "required": []}}},
    {"type": "function", "function": {"name": "compute_wellbeing_composite", "description": "Per-week 0-100 wellbeing score using category weights (override via `weights`: {category: weight}).", "parameters": {"type": "object", "properties": {"weights": {"type": "object"}}, "required": []}}},
    {"type": "function", "function": {"name": "get_rate_of_change", "description": "Top-N week-over-week deltas (velocity).", "parameters": {"type": "object", "properties": {"last_n_weeks": {"type": "integer", "default": 4}, "top_n": {"type": "integer", "default": 15}}, "required": []}}},
    {"type": "function", "function": {"name": "cluster_similar_weeks", "description": "Group similar weeks by multivariate similarity across all metric paths.", "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {"name": "get_lagged_correlation", "description": "Find if one metric leads/lags another.", "parameters": {"type": "object", "properties": {"metric_a": {"type": "string", "default": "stresses.work_academic", "description": _METRIC_PATH_HINT}, "metric_b": {"type": "string", "default": "wellbeing_score", "description": _METRIC_PATH_HINT}, "max_lag": {"type": "integer", "default": 3}}, "required": []}}},
    {"type": "function", "function": {"name": "get_streaks", "description": "Top-N improvement/decline streaks.", "parameters": {"type": "object", "properties": {"top_n": {"type": "integer", "default": 15}}, "required": []}}},
    {"type": "function", "function": {"name": "estimate_coping_effectiveness", "description": "Compares composite wellbeing on weeks where each coping mechanism is active vs. inactive.", "parameters": {"type": "object", "properties": {}, "required": []}}},
]


# ─────────────────────────────────────────────────────────────────────────────
# TOOL DISPATCHER
# ─────────────────────────────────────────────────────────────────────────────

def dispatch_tool(name: str, arguments: dict, user_data: dict) -> str:
    arguments = arguments or {}
    try:
        if name == "get_personal_profile":
            result = get_personal_profile(user_data)
        elif name == "get_recent_chat_history":
            result = get_recent_chat_history(user_data, **arguments)
        elif name == "get_wellbeing_trend":
            result = get_wellbeing_trend(user_data, **arguments)
        elif name == "get_wellbeing_snapshot":
            result = get_wellbeing_snapshot(user_data)
        elif name == "compare_wellbeing_weeks":
            result = compare_wellbeing_weeks(user_data, **arguments)
        elif name == "get_metric_summary_stats":
            result = get_metric_summary_stats(user_data, **arguments)
        elif name == "detect_pattern_shift":
            result = detect_pattern_shift(user_data, **arguments)
        elif name == "get_correlated_factors":
            result = get_correlated_factors(user_data, **arguments)
        elif name == "compute_volatility":
            result = compute_volatility(user_data, **arguments)
        elif name == "detect_anomalies":
            result = detect_anomalies(user_data, **arguments)
        elif name == "predict_trend":
            result = predict_trend(user_data, **arguments)
        elif name == "compute_wellbeing_composite":
            result = compute_wellbeing_composite(user_data, **arguments)
        elif name == "get_rate_of_change":
            result = get_rate_of_change(user_data, **arguments)
        elif name == "cluster_similar_weeks":
            result = cluster_similar_weeks(user_data, **arguments)
        elif name == "get_lagged_correlation":
            result = get_lagged_correlation(user_data, **arguments)
        elif name == "get_streaks":
            result = get_streaks(user_data, **arguments)
        elif name == "estimate_coping_effectiveness":
            result = estimate_coping_effectiveness(user_data, **arguments)
        else:
            result = {"error": f"Unknown tool: {name}"}
    except Exception as e:
        result = {"error": f"Tool execution failed: {str(e)}"}

    return json.dumps(result, indent=2, default=str)


# ─────────────────────────────────────────────────────────────────────────────
# SYSTEM PROMPT
# ─────────────────────────────────────────────────────────────────────────────

SYSTEM_PROMPT = """You are a warm, intuitive mental health counselor. You listen deeply and speak like a real person.

You have access to understand the user's history and patterns, but YOU NEVER MENTION THE DATA, METRICS, SCORES, OR TRACKING.
The user should feel like you know them naturally, not like you're reading a spreadsheet.

═══════════════════════════════════════════════════════════════
YOUR INTERNAL PROCESS (hidden from user):
═══════════════════════════════════════════════════════════════

1. Call tools to understand patterns, trends, strengths, and concerns
   - get_wellbeing_snapshot: current emotional state across categories
   - detect_pattern_shift, get_rate_of_change: how things are moving
   - get_correlated_factors: what's connected to wellbeing
   - get_streaks: positive momentum and resilience
   - estimate_coping_effectiveness: what actually helps them

2. Available metric paths follow the schema:
   - emotions.{calm_neutral, happy_positive, anxious_worried, sad_low, angry_irritable, lonely, overwhelmed, numb_emotionally_flat}
   - stresses.{work_academic, relationship, health_related, financial, time_pressure_overload, uncertainty_future_anxiety, internal_pressure, low_manageable}
   - cognitive_patterns.{balanced_realistic, rumination, catastrophizing, black_and_white, self_critical, helplessness_low_control, overanalysis_indecision, positive_reframing}
   - sleep.{restful_healthy, mild_disturbance, insufficient_sleep, insomnia, irregular_schedule, oversleeping_fatigue}
   - energy.{high_energized, stable_normal, low_tired, exhausted_drained, fluctuating, restless_wired}
   - habits.{structured_healthy_routines, productive_habits, inconsistent_routines, procrastination, avoidance_behaviors, compulsive_behaviors, self_care_present, self_care_neglect}
   - social.{strong_support_system, moderate_support, limited_support, socially_isolated, active_engagement, relationship_conflict, help_seeking_behavior, withdrawing}
   - personality.{optimistic, pessimistic, self_confident, self_doubting, emotionally_reactive, emotionally_stable, introverted, socially_expressive, conscientious_disciplined, avoidant_tendency}
   - motivation_values.{highly_motivated_goal_driven, moderate_motivation, low_motivation_disengaged, anhedonia_loss_of_interest, purpose_driven, value_conflict, directionless_unclear_goals}
   - Plus: 'wellbeing_score' (composite 0-100), 'exercise_minutes'

3. Analyze deeply but SILENTLY:
   - What patterns do I see?
   - What strengths emerge from their history?
   - Where do they need support?
   - What's the underlying story here?

4. Respond with human warmth, not data:
   - Use observations, not numbers
   - Say "looking at your patterns" not "your wellbeing_score is 62"
   - Say "I notice you tend to struggle when..." not "correlation is 0.68"
   - Say "there's momentum here" not "3-week improving streak"

═══════════════════════════════════════════════════════════════
RESPONSE GUIDELINES:
═══════════════════════════════════════════════════════════════

✓ DO:
  • Speak naturally, like a caring friend
  • Use what you learn to personalize responses
  • Reference their actual struggles and wins without mentioning numbers
  • Validate feelings first, then offer insight
  • One insight per response, max

✗ DON'T:
  • Mention scores, metrics, percentages, or measurements
  • Reference field names or paths
  • Talk about "data", "tracking", or "analysis"
  • List statistics
  • Be clinical or robotic

═══════════════════════════════════════════════════════════════
EXAMPLES:
═══════════════════════════════════════════════════════════════

❌ BAD: "Your stresses.work_academic has increased to 5/5. Sleep.insufficient_sleep correlates with mood at r=0.75."
✓ GOOD: "Things have been heavier lately, especially with work. And it sounds like when you're not sleeping well, your whole perspective shifts."

❌ BAD: "Your coping mechanism 'exercise' has effectiveness 0.82."
✓ GOOD: "I notice that when you've been moving your body—even just walks—something shifts in how you feel."

═══════════════════════════════════════════════════════════════
HARD RULES:
═══════════════════════════════════════════════════════════════

- NEVER let data language slip into your response
- NEVER mention you're analyzing, measuring, or tracking
- Use data ONLY to inform your intuition—not to explain it
- Weave references naturally ("I remember you mentioning..." not "in week 3 you said...")
- If crisis indicators emerge, suggest professional support warmly—not clinically
"""


# ─────────────────────────────────────────────────────────────────────────────
# AGENTIC LOOP
# ─────────────────────────────────────────────────────────────────────────────

def run_agentic_turn(client: Groq, messages: list, user_data: dict, verbose: bool = True) -> str:
    while True:
        response = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=messages,
            tools=TOOLS,
            tool_choice="auto",
            max_tokens=1024,
        )
        choice = response.choices[0]
        msg = choice.message

        msg_dict = {"role": "assistant", "content": msg.content or ""}
        if msg.tool_calls:
            msg_dict["tool_calls"] = [
                {
                    "id": tc.id,
                    "type": "function",
                    "function": {"name": tc.function.name, "arguments": tc.function.arguments},
                }
                for tc in msg.tool_calls
            ]
        messages.append(msg_dict)

        if not msg.tool_calls:
            return msg.content or ""

        for tc in msg.tool_calls:
            fn_name = tc.function.name
            try:
                fn_args = json.loads(tc.function.arguments or "{}")
            except json.JSONDecodeError:
                fn_args = {}

            if verbose:
                args_str = ", ".join(f"{k}={v!r}" for k, v in fn_args.items()) if fn_args else ""
                print(f"  🔧 {fn_name}({args_str})")

            result_str = dispatch_tool(fn_name, fn_args, user_data)
            messages.append({
                "role": "tool",
                "tool_call_id": tc.id,
                "content": result_str,
            })


# ─────────────────────────────────────────────────────────────────────────────
# BANNER & MAIN LOOP
# ─────────────────────────────────────────────────────────────────────────────

def print_banner(user_data: dict):
    name   = user_data["personal_memory"]["name"]
    frames = user_data["wellbeing_frames"]
    weeks  = len(frames)
    latest = frames[-1] if frames else {}

    composite = _compute_composite_score(latest) if latest else None
    happy   = _get_value(latest, "emotions.happy_positive") if latest else None
    anxious = _get_value(latest, "emotions.anxious_worried") if latest else None
    work    = _get_value(latest, "stresses.work_academic") if latest else None

    print("\n" + "═" * 70)
    print("   🧠  Mental Health Counselor — Analytics-Enhanced (Groq)")
    print("═" * 70)
    print(f"   User: {name}, {user_data['personal_memory'].get('age', '?')}y")
    print(f"   Data: {weeks} weeks tracked")
    if composite is not None:
        print(f"   Wellbeing score (latest): {composite}/100")
    if happy is not None:
        print(f"   Happy/positive: {happy}/5    Anxious/worried: {anxious}/5    Work stress: {work}/5")
    print("═" * 70)
    print("   Available tools: 5 core + 12 analytics")
    print("   Tool calls shown with 🔧 (use --quiet to hide)")
    print("   Type 'quit' to exit.\n")


def main():
    parser = argparse.ArgumentParser(description="Mental Health Counselor (Groq) with Analytics")
    parser.add_argument("--data", default="data3.json", help="Path to user data JSON")
    parser.add_argument("--quiet", action="store_true", help="Hide tool calls")
    parser.add_argument("--model", default="llama-3.3-70b-versatile", help="Groq model")
    args = parser.parse_args()

    if not os.path.exists(args.data):
        print(f"Error: data file '{args.data}' not found.")
        sys.exit(1)

    user_data = load_user_data(args.data)

    api_key = "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"
    if not api_key:
        print("Error: GROQ_API_KEY environment variable not set.")
        sys.exit(1)

    client = Groq(api_key=api_key)

    print_banner(user_data)

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    while True:
        try:
            user_input = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nTake care. Goodbye! 🌿")
            break

        if not user_input:
            continue
        if user_input.lower() in ("quit", "exit", "bye"):
            print("\nTake care of yourself. Goodbye! 🌿")
            break

        messages.append({"role": "user", "content": user_input})

        print("\nCounselor (thinking...):")
        reply = run_agentic_turn(client, messages, user_data, verbose=not args.quiet)
        print(f"\nCounselor: {reply}\n")


if __name__ == "__main__":
    main()