"""Recovery — sleep, HRV, RHR, body battery, plus training↔recovery view.

Port of notebook 03. Layout:
    Row 1: 4 tiles (last sleep hours / RHR / HRV last-night with status / BB high)
    Row 2: Sleep composition (stacked area over last 60 days: deep / light / rem)
    Row 3: HRV last-night + RHR overlay (twin axis), 90-day trend with rolling avg
    Row 4: "Does yesterday's training affect today's HRV?" — scatter + Pearson r
    Row 5: Weekly summary table
"""

from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import streamlit as st

from openrun.web._helpers import (
    cached_sleep_stages,
    cached_wellness,
    empty_state,
    show_sidebar,
)


st.set_page_config(page_title="Recovery · openrun", page_icon="🛌", layout="wide")
show_sidebar()
st.title("🛌 Recovery")

wellness = cached_wellness()
stages = cached_sleep_stages()

if wellness.empty:
    empty_state(
        "No wellness data yet.",
        "Run `openrun-sync` or import a Garmin export with daily sleep/HRV/RHR.",
    )
    st.stop()


# ---------------------------------------------------------------------------
# Row 1 — at-a-glance tiles (most recent non-NaN per metric)
# ---------------------------------------------------------------------------

def _latest(series: pd.Series) -> tuple[pd.Timestamp | None, float | None]:
    s = series.dropna()
    if s.empty:
        return None, None
    return s.index[-1], s.iloc[-1]


sleep_date, sleep_total_s = _latest(wellness.get("sleep_total_s", pd.Series(dtype=float)))
sleep_h = (sleep_total_s / 3600.0) if sleep_total_s else None

rhr_date, rhr_val = _latest(wellness.get("resting_hr", pd.Series(dtype=float)))
hrv_date, hrv_val = _latest(wellness.get("hrv_last_night", pd.Series(dtype=float)))
hrv_status = wellness.get("hrv_status", pd.Series(dtype=object)).dropna()
hrv_status_val = hrv_status.iloc[-1] if not hrv_status.empty else None
bb_date, bb_high = _latest(wellness.get("bb_highest", pd.Series(dtype=float)))

cols = st.columns(4)
cols[0].metric(
    "Sleep last night",
    f"{sleep_h:.1f}h" if sleep_h is not None else "—",
    sleep_date.strftime("%b %d") if sleep_date is not None else None,
)
cols[1].metric(
    "Resting HR",
    f"{rhr_val:.0f}" if rhr_val is not None else "—",
    rhr_date.strftime("%b %d") if rhr_date is not None else None,
)
cols[2].metric(
    "HRV last night",
    f"{hrv_val:.0f}" if hrv_val is not None else "—",
    f"{hrv_status_val}" if hrv_status_val else (hrv_date.strftime("%b %d") if hrv_date else None),
)
cols[3].metric(
    "Body battery (high)",
    f"{bb_high:.0f}" if bb_high is not None else "—",
    bb_date.strftime("%b %d") if bb_date is not None else None,
)


# ---------------------------------------------------------------------------
# Row 2 — Sleep composition over time
# ---------------------------------------------------------------------------

st.subheader("Sleep composition")
if stages.empty or stages[["deep_s", "light_s", "rem_s"]].dropna(how="all").empty:
    empty_state("No sleep-stage data yet.")
else:
    window = st.select_slider(
        "Window",
        options=[14, 30, 60, 90, 180],
        value=60,
        format_func=lambda d: f"{d}d",
        key="sleep_window",
    )
    s = stages.tail(window).copy()
    s["deep_h"] = s["deep_s"] / 3600
    s["light_h"] = s["light_s"] / 3600
    s["rem_h"] = s["rem_s"] / 3600
    fig, ax = plt.subplots(figsize=(11, 3.6))
    ax.stackplot(
        s.index,
        s["deep_h"].fillna(0),
        s["light_h"].fillna(0),
        s["rem_h"].fillna(0),
        labels=["deep", "light", "rem"],
        colors=["#264653", "#2a9d8f", "#e9c46a"],
        alpha=0.9,
    )
    ax.set_ylabel("hours")
    ax.legend(loc="upper left", fontsize=8, ncol=3)
    fig.tight_layout()
    st.pyplot(fig, clear_figure=True)


# ---------------------------------------------------------------------------
# Row 3 — HRV + RHR overlay
# ---------------------------------------------------------------------------

st.subheader("HRV & resting HR")
if "hrv_last_night" in wellness or "resting_hr" in wellness:
    window2 = st.select_slider(
        "Window",
        options=[30, 60, 90, 180, 365],
        value=90,
        format_func=lambda d: f"{d}d",
        key="hrv_window",
    )
    w = wellness.tail(window2)
    fig, ax1 = plt.subplots(figsize=(11, 3.8))
    if "hrv_last_night" in w:
        ax1.plot(w.index, w["hrv_last_night"], color="#2a9d8f", lw=1.2, alpha=0.55, label="HRV")
        ax1.plot(
            w.index,
            w["hrv_last_night"].rolling(7).mean(),
            color="#2a9d8f", lw=2.0, label="HRV · 7d",
        )
        ax1.set_ylabel("HRV (ms)", color="#2a9d8f")
    ax2 = ax1.twinx()
    if "resting_hr" in w:
        ax2.plot(w.index, w["resting_hr"], color="#e76f51", lw=1.0, alpha=0.4, label="RHR")
        ax2.plot(
            w.index,
            w["resting_hr"].rolling(7).mean(),
            color="#e76f51", lw=2.0, label="RHR · 7d",
        )
        ax2.set_ylabel("Resting HR (bpm)", color="#e76f51")
    ax1.legend(loc="upper left", fontsize=8)
    fig.tight_layout()
    st.pyplot(fig, clear_figure=True)


# ---------------------------------------------------------------------------
# Row 4 — Yesterday's training vs today's HRV
# ---------------------------------------------------------------------------

st.subheader("Training load → next-morning HRV")
st.caption(
    "If hard sessions actually tax your nervous system, expect lower next-morning HRV. "
    "Pearson r reported below; weak signals (|r| < 0.2) usually mean other factors dominate."
)

if "hrv_last_night" in wellness:
    from openrun.web._helpers import get_conn

    daily_tl = pd.read_sql(
        """SELECT date(start_time_local) AS d, SUM(training_load) AS tl
           FROM activities
           WHERE activity_type IN ('running','trail_running','cycling','strength_training')
             AND training_load IS NOT NULL
           GROUP BY d""",
        get_conn(),
        parse_dates=["d"],
    ).set_index("d")["tl"]

    if not daily_tl.empty:
        # Next-day HRV: shift HRV index back one day so HRV(today) lines up with TL(yesterday).
        hrv_next = wellness["hrv_last_night"].copy()
        hrv_next.index = hrv_next.index - pd.Timedelta(days=1)
        joined = pd.DataFrame({"tl": daily_tl, "hrv_next": hrv_next}).dropna()

        if len(joined) >= 10:
            r = joined["tl"].corr(joined["hrv_next"])
            st.caption(f"Pearson r between yesterday's TL and tomorrow's HRV: **{r:+.2f}** (n={len(joined)})")

            fig, ax = plt.subplots(figsize=(8, 4))
            ax.scatter(joined["tl"], joined["hrv_next"], alpha=0.45, color="#264653", s=18)
            if joined["tl"].std() > 0:
                a, b = np.polyfit(joined["tl"], joined["hrv_next"], 1)
                xs = np.linspace(joined["tl"].min(), joined["tl"].max(), 50)
                ax.plot(xs, a * xs + b, color="#e76f51", lw=1.5)
            ax.set_xlabel("Yesterday's training load")
            ax.set_ylabel("Next-night HRV (ms)")
            fig.tight_layout()
            st.pyplot(fig, clear_figure=True)
        else:
            st.info(f"Need ≥10 matched day-pairs; have {len(joined)}.")
    else:
        st.info("No training-load data to compare against.")
else:
    st.info("No HRV data available.")


# ---------------------------------------------------------------------------
# Row 5 — Weekly summary
# ---------------------------------------------------------------------------

st.subheader("Weekly summary — last 12 weeks")
weekly = wellness.copy()
weekly["week"] = weekly.index.to_period("W-SUN").start_time
agg = weekly.groupby("week").agg(
    sleep_h_mean=("sleep_hours", "mean"),
    hrv_mean=("hrv_last_night", "mean"),
    rhr_mean=("resting_hr", "mean"),
    stress_mean=("avg_stress", "mean"),
).tail(12).round(1)
agg.index = agg.index.strftime("%Y-%m-%d")
st.dataframe(agg, width="stretch")
