modules/my_patterns.py

"""Module — My Patterns: aggregated insights, charts, and dynamic recommendations.

This tab reads entirely from session-state populated by the Thought Diary
(cognitive_journal) module. It does NOT write any journal data — it is a
read-only consumer.

Surfaces:
- Distortion pattern summary & bar chart
- PHQ-9 / GAD-7 screener history timeline
- Daily check-in trend lines (mood, sleep, stress)
- LLM-powered song & activity recommendations based on journal context
"""
from __future__ import annotations

from collections import Counter
from typing import Any, Dict, List, Optional, Sequence

import plotly.graph_objects as go
import streamlit as st

from backend.claude_client import chat
from backend.i18n import claude_language_name, t
from modules.cognitive_journal import (
    CHECKINS_KEY,
    ENTRIES_KEY,
    GAD7_HISTORY_KEY,
    PHQ9_HISTORY_KEY,
    get_cognitive_journal_context,
)

SONGS_KEY = "my_patterns_songs"
ACTIVITIES_KEY = "my_patterns_activities"

# Re-use the same distortion labels used in the journal
DISTORTION_LABELS: Dict[str, str] = {
    "catastrophizing": "Catastrophizing",
    "mind_reading": "Mind Reading",
    "all_or_nothing": "All-or-Nothing",
    "fortune_telling": "Fortune Telling",
    "personalization": "Personalization",
    "mental_filter": "Mental Filter",
    "emotional_reasoning": "Emotional Reasoning",
    "should_statements": "Should Statements",
}


# ── session helpers ──────────────────────────────────────────────────────────

def _init_state() -> None:
    if SONGS_KEY not in st.session_state:
        st.session_state[SONGS_KEY] = ""
    if ACTIVITIES_KEY not in st.session_state:
        st.session_state[ACTIVITIES_KEY] = ""


def _has_data() -> bool:
    """Return True if there is any journal / screener / check-in data."""
    return bool(
        st.session_state.get(ENTRIES_KEY)
        or st.session_state.get(CHECKINS_KEY)
        or st.session_state.get(PHQ9_HISTORY_KEY)
        or st.session_state.get(GAD7_HISTORY_KEY)
    )


# ── data extraction (read-only from session state) ───────────────────────────

def _extract_distortion_counts(entries: Sequence[Any]) -> Dict[str, int]:
    """Count distortion types across entries."""
    counts: Counter[str] = Counter()
    for entry in entries:
        distortions = []
        if isinstance(entry, dict):
            distortions = entry.get("distortions", [])
        for d in distortions:
            dtype = d.get("type") if isinstance(d, dict) else None
            if dtype:
                counts[dtype] += 1
    return dict(counts)


def _extract_mood_counts(entries: Sequence[Any]) -> Dict[str, int]:
    """Count mood labels across entries."""
    counts: Counter[str] = Counter()
    for entry in entries:
        mood = None
        if isinstance(entry, dict):
            mood = entry.get("overall_mood") or entry.get("mood")
        if mood:
            counts[mood] += 1
    return dict(counts)


def _extract_checkin_series(checkins: Sequence[Any]):
    """Return parallel lists of (timestamps, mood, sleep, stress) for chart."""
    timestamps, moods, sleeps, stresses = [], [], [], []
    for ci in checkins:
        if not isinstance(ci, dict):
            continue
        timestamps.append(ci.get("timestamp", ""))
        moods.append(ci.get("mood"))
        sleeps.append(ci.get("sleep"))
        stresses.append(ci.get("stress"))
    return timestamps, moods, sleeps, stresses


# ── chart builders ───────────────────────────────────────────────────────────

def _render_distortion_chart(entries: Sequence[Any]) -> None:
    """Horizontal bar chart of distortion frequency."""
    counts = _extract_distortion_counts(entries)
    if not counts:
        st.caption("No cognitive distortions identified yet.")
        return

    sorted_items = sorted(counts.items(), key=lambda kv: kv[1], reverse=True)
    labels = [DISTORTION_LABELS.get(k, k) for k, _ in sorted_items]
    values = [v for _, v in sorted_items]

    fig = go.Figure(
        go.Bar(
            x=values,
            y=labels,
            orientation="h",
            marker_color="#8B5CF6",
        )
    )
    fig.update_layout(
        title="Distortion Frequency",
        xaxis_title="Count",
        yaxis_title="",
        height=max(250, 40 * len(labels)),
        margin=dict(l=10, r=10, t=40, b=30),
    )
    st.plotly_chart(fig, use_container_width=True)


def _render_mood_chart(entries: Sequence[Any]) -> None:
    """Pie chart of mood distribution."""
    counts = _extract_mood_counts(entries)
    if not counts:
        return

    mood_colors = {
        "anxious": "#F59E0B",
        "sad": "#6366F1",
        "frustrated": "#EF4444",
        "hopeful": "#10B981",
        "neutral": "#9CA3AF",
        "overwhelmed": "#EC4899",
    }

    labels = list(counts.keys())
    values = list(counts.values())
    colors = [mood_colors.get(m, "#8B5CF6") for m in labels]

    fig = go.Figure(
        go.Pie(
            labels=[m.capitalize() for m in labels],
            values=values,
            marker=dict(colors=colors),
            hole=0.4,
        )
    )
    fig.update_layout(
        title="Mood Distribution",
        height=300,
        margin=dict(l=10, r=10, t=40, b=10),
    )
    st.plotly_chart(fig, use_container_width=True)


def _render_checkin_trends(checkins: Sequence[Any]) -> None:
    """Line chart of mood / stress trends over check-ins."""
    timestamps, moods, sleeps, stresses = _extract_checkin_series(checkins)
    if len(timestamps) < 2:
        return

    indices = list(range(1, len(timestamps) + 1))

    fig = go.Figure()
    if any(m is not None for m in moods):
        fig.add_trace(go.Scatter(x=indices, y=moods, mode="lines+markers", name="Mood (/10)"))
    if any(s is not None for s in sleeps):
        fig.add_trace(go.Scatter(x=indices, y=sleeps, mode="lines+markers", name="Sleep (h)"))
    if any(s is not None for s in stresses):
        fig.add_trace(go.Scatter(x=indices, y=stresses, mode="lines+markers", name="Stress (/10)"))

    fig.update_layout(
        title="Daily Check-in Trends",
        xaxis_title="Check-in #",
        height=320,
        margin=dict(l=10, r=10, t=40, b=30),
    )
    st.plotly_chart(fig, use_container_width=True)


def _render_screener_timeline(kind: str, history: Sequence[Any]) -> None:
    """Timeline of screener scores."""
    if not history or len(history) < 1:
        return

    scores = []
    times = []
    for idx, record in enumerate(history):
        if isinstance(record, dict):
            scores.append(record.get("score", 0))
            times.append(record.get("taken_at", f"Attempt {idx + 1}"))
        elif isinstance(record, (int, float)):
            scores.append(record)
            times.append(f"Attempt {idx + 1}")

    if not scores:
        return

    max_score = 27 if kind == "phq9" else 21
    label = "PHQ-9 (Depression)" if kind == "phq9" else "GAD-7 (Anxiety)"

    fig = go.Figure(
        go.Scatter(
            x=list(range(1, len(scores) + 1)),
            y=scores,
            mode="lines+markers",
            name=label,
            marker=dict(color="#6366F1" if kind == "phq9" else "#EC4899"),
        )
    )
    fig.update_layout(
        title=label,
        xaxis_title="Attempt",
        yaxis_title="Score",
        yaxis_range=[0, max_score],
        height=280,
        margin=dict(l=10, r=10, t=40, b=30),
    )
    st.plotly_chart(fig, use_container_width=True)


# ── LLM-powered recommendations ─────────────────────────────────────────────

_SONG_PROMPT = """You are Saathi's music recommender. Based on the user's mental-health journal context below, recommend 5 songs (mix of Indian and International) that match their current emotional state.

For each song provide:
- Song name and artist
- One line saying why this song fits their mood
- A YouTube or Spotify search term they can use to find it

Format: numbered list. Respond in {language_name}.

Journal context:
{cross_module_memory}
"""

_ACTIVITY_PROMPT = """You are Saathi's wellness activity recommender. Based on the user's mental-health journal context below, suggest 5 evidence-based coping activities matched to their emotional state.

Activities should be:
- Immediately actionable (can do right now, at home, for free)
- Evidence-based (CBT, mindfulness, behavioral activation)
- Appropriate to their mood/stress level
- Mix of physical, creative, and mindful activities

Format: numbered list with a brief description for each. Respond in {language_name}.

Journal context:
{cross_module_memory}
"""


def _get_song_recommendations(lang: str) -> str:
    """Call the LLM for personalised song recommendations."""
    context = get_cognitive_journal_context() or "No journal data yet."
    return chat(
        module="soothe_poetry",  # reuse soothe module's routing
        user_text="Recommend songs based on my journal patterns.",
        language_name=claude_language_name(lang),
        max_tokens=800,
        extra_context={
            "cross_module_memory": context,
        },
        history=[{
            "role": "user",
            "content": _SONG_PROMPT.format(
                language_name=claude_language_name(lang),
                cross_module_memory=context,
            ),
        }],
    )


def _get_activity_recommendations(lang: str) -> str:
    """Call the LLM for personalised activity recommendations."""
    context = get_cognitive_journal_context() or "No journal data yet."
    return chat(
        module="soothe_poetry",
        user_text="Suggest coping activities based on my journal patterns.",
        language_name=claude_language_name(lang),
        max_tokens=800,
        extra_context={
            "cross_module_memory": context,
        },
        history=[{
            "role": "user",
            "content": _ACTIVITY_PROMPT.format(
                language_name=claude_language_name(lang),
                cross_module_memory=context,
            ),
        }],
    )


# ── main render ──────────────────────────────────────────────────────────────

def render(lang: str) -> None:
    """Top-level render for the My Patterns tab."""
    _init_state()

    st.header(t("patterns_header", lang))
    st.caption(t("patterns_sub", lang))

    if not _has_data():
        st.info(t("patterns_empty", lang))
        st.info(t("patterns_soothe_nudge", lang))
        return

    entries = list(st.session_state.get(ENTRIES_KEY, []) or [])
    checkins = list(st.session_state.get(CHECKINS_KEY, []) or [])
    phq9_history = list(st.session_state.get(PHQ9_HISTORY_KEY, []) or [])
    gad7_history = list(st.session_state.get(GAD7_HISTORY_KEY, []) or [])

    # ── Charts row ────────────────────────────────────────────────────────
    if entries:
        col_dist, col_mood = st.columns(2)
        with col_dist:
            _render_distortion_chart(entries)
        with col_mood:
            _render_mood_chart(entries)

    # ── Check-in trends ──────────────────────────────────────────────────
    if checkins:
        _render_checkin_trends(checkins)

    # ── Screener timelines ───────────────────────────────────────────────
    if phq9_history or gad7_history:
        col_phq, col_gad = st.columns(2)
        with col_phq:
            if phq9_history:
                _render_screener_timeline("phq9", phq9_history)
        with col_gad:
            if gad7_history:
                _render_screener_timeline("gad7", gad7_history)

    st.divider()

    # ── Recommendation section ───────────────────────────────────────────
    rec_col1, rec_col2 = st.columns(2)

    with rec_col1:
        st.subheader(t("patterns_songs_heading", lang))
        st.caption(t("patterns_songs_sub", lang))
        if st.button(
            t("patterns_songs_button", lang),
            key="my_patterns_songs_button",
            use_container_width=True,
            type="primary",
        ):
            with st.spinner("🎵 …"):
                try:
                    st.session_state[SONGS_KEY] = _get_song_recommendations(lang)
                except Exception as e:
                    st.session_state[SONGS_KEY] = f"(Could not reach the model: {e})"

        if st.session_state[SONGS_KEY]:
            with st.container(border=True):
                st.markdown(st.session_state[SONGS_KEY])

    with rec_col2:
        st.subheader(t("patterns_activities_heading", lang))
        st.caption(t("patterns_activities_sub", lang))
        if st.button(
            t("patterns_activities_button", lang),
            key="my_patterns_activities_button",
            use_container_width=True,
            type="primary",
        ):
            with st.spinner("🧘 …"):
                try:
                    st.session_state[ACTIVITIES_KEY] = _get_activity_recommendations(lang)
                except Exception as e:
                    st.session_state[ACTIVITIES_KEY] = f"(Could not reach the model: {e})"

        if st.session_state[ACTIVITIES_KEY]:
            with st.container(border=True):
                st.markdown(st.session_state[ACTIVITIES_KEY])

    st.divider()
    st.info(t("patterns_soothe_nudge", lang))