"""
Streamlit dashboard for the data-scaling study.

Run from the experiments/data_scaling_study/ directory:
    streamlit run dashboard/app.py

Three sections:
  1. Learning curves       — per-epoch metrics for every (model, share) run
  2. Data share vs final   — best val mIoU/Dice/IoU/PixelAcc as a function of data share
  3. Inference             — upload an image, see all 6 segmentations side-by-side

Reads logs from ../logs and checkpoints from ../checkpoints. Sections gracefully
degrade when runs are missing — useful while training is still in flight.

Note on the 100% rows
─────────────────────
The 100% checkpoints are not retrained here — they are bootstrapped from the
existing pv_panel_models/ baselines via bootstrap_100.py. Per-epoch metrics at
100% are parsed from the old text logs (per-batch averaging) and mIoU is null
per epoch (the old trainer didn't compute it). The single comparable number on
the scaling chart for 100% is read from the bootstrap's `recomputed_val_metrics`
field, which uses the same global confusion-matrix metric code as the 25/50%
runs. Banners in each tab explain.
"""
import io
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import streamlit as st
import torch
from PIL import Image
from torchvision import transforms

THIS_DIR = Path(__file__).resolve().parent
EXP_DIR = THIS_DIR.parent
LOGS_DIR = EXP_DIR / "logs"
CKPT_DIR = EXP_DIR / "checkpoints"

sys.path.insert(0, str(EXP_DIR))
from models import MODEL_REGISTRY  # noqa: E402

MODELS = ["unet", "segformer_b0"]
SHARES = [25, 50, 100]
PRETTY_MODEL = {"unet": "U-Net", "segformer_b0": "SegFormer-B0"}
METRICS = [
    ("dice", "Dice"),
    ("miou", "mIoU"),
    ("iou", "Foreground IoU"),
    ("pixel_acc", "Pixel Accuracy"),
    ("loss", "Loss"),
]

st.set_page_config(page_title="Data Scaling Study", layout="wide")


# ── Loaders ────────────────────────────────────────────────────────────────
@st.cache_data(show_spinner=False)
def load_log(model: str, share: int):
    p = LOGS_DIR / f"{model}_{share}.json"
    if not p.is_file():
        return None
    with open(p) as f:
        return json.load(f)


def log_to_df(log):
    df = pd.DataFrame(log["epochs"])
    df["model"] = log["model"]
    df["share"] = log["share"]
    return df


def is_bootstrapped(log):
    return log.get("bootstrapped_from") is not None


@st.cache_data(show_spinner=False)
def load_all_logs():
    logs = {}
    for m in MODELS:
        for s in SHARES:
            log = load_log(m, s)
            if log is not None:
                logs[(m, s)] = log
    return logs


def fmt_hms(seconds):
    if seconds is None:
        return "—"
    seconds = int(round(seconds))
    h, rem = divmod(seconds, 3600)
    m, s = divmod(rem, 60)
    return f"{h:d}:{m:02d}:{s:02d}" if h else f"{m:d}:{s:02d}"


def scaling_row(log):
    """Best-checkpoint val metrics for the scaling chart.

    For trained 25/50% runs: read the per-epoch maximum from the JSON.
    For bootstrapped 100% runs: read from `recomputed_val_metrics` so the
    metric definition matches the 25/50% runs.
    """
    epochs = log["epochs"]
    row = {
        "model": PRETTY_MODEL[log["model"]],
        "share": log["share"],
        "source": "bootstrapped" if is_bootstrapped(log) else "trained",
    }
    if is_bootstrapped(log) and log.get("recomputed_val_metrics") is not None:
        rv = log["recomputed_val_metrics"]
        row.update({
            "best_val_dice": rv.get("dice"),
            "best_val_miou": rv.get("miou"),
            "best_val_iou": rv.get("iou"),
            "best_val_pixel_acc": rv.get("pixel_acc"),
        })
    elif epochs:
        # Best-by-Dice index (matches the saved best.pth selection)
        idx = max(range(len(epochs)), key=lambda i: epochs[i].get("val_dice", -1) or -1)
        best = epochs[idx]
        row.update({
            "best_val_dice": best.get("val_dice"),
            "best_val_miou": best.get("val_miou"),
            "best_val_iou": best.get("val_iou"),
            "best_val_pixel_acc": best.get("val_pixel_acc"),
        })
    else:
        row.update({k: None for k in (
            "best_val_dice", "best_val_miou", "best_val_iou", "best_val_pixel_acc"
        )})

    # Wall-clock timing
    wall = log.get("wall_clock_seconds")  # trained runs
    if wall is None:
        wall = log.get("val_recompute_seconds")  # bootstrapped runs
    row["wall_clock_seconds"] = wall
    row["wall_clock"] = fmt_hms(wall)
    if epochs:
        per_epoch = [e.get("epoch_seconds") for e in epochs if e.get("epoch_seconds") is not None]
        row["sec_per_epoch"] = (sum(per_epoch) / len(per_epoch)) if per_epoch else None
    else:
        row["sec_per_epoch"] = None
    return row


# ── Inference helpers ──────────────────────────────────────────────────────
@st.cache_resource(show_spinner=False)
def load_best(model_name: str, share: int, device: str):
    p = CKPT_DIR / f"{model_name}_{share}_best.pth"
    if not p.is_file():
        return None
    builder = MODEL_REGISTRY[model_name]
    model, _ = builder()
    state = torch.load(p, map_location=device, weights_only=False)
    model.load_state_dict(state["model_state_dict"])
    model.to(device).eval()
    return model


def preprocess(image: Image.Image, image_size: int = 128):
    tf = transforms.Compose([
        transforms.Resize((image_size, image_size)),
        transforms.ToTensor(),
    ])
    return tf(image.convert("RGB")).unsqueeze(0)


def run_inference(model, image_tensor, device, threshold=0.5):
    """Returns (probs_2d, mask_2d) both as 2-D float numpy arrays in [0,1]."""
    with torch.no_grad():
        logits = model(image_tensor.to(device))
        probs = torch.sigmoid(logits).squeeze().cpu().numpy()
    if probs.ndim != 2:
        probs = probs.reshape(probs.shape[-2], probs.shape[-1])
    mask = (probs > threshold).astype(np.float32)
    return probs, mask


def overlay(rgb: np.ndarray, mask: np.ndarray, color=(0, 255, 0), alpha=0.45):
    out = rgb.copy()
    m = mask.astype(bool)
    out[m] = (alpha * np.array(color) + (1 - alpha) * out[m]).astype(np.uint8)
    return out


def heatmap(probs: np.ndarray) -> np.ndarray:
    """Map a [0,1] probability map to a 3-channel uint8 RGB image (red→hot)."""
    p = np.clip(probs, 0.0, 1.0)
    rgb = np.zeros((p.shape[0], p.shape[1], 3), dtype=np.uint8)
    rgb[..., 0] = (p * 255).astype(np.uint8)             # R
    rgb[..., 1] = (np.maximum(0, 1 - 2 * np.abs(p - 0.5)) * 255).astype(np.uint8)  # G
    rgb[..., 2] = ((1 - p) * 255).astype(np.uint8)       # B
    return rgb


# ── UI ─────────────────────────────────────────────────────────────────────
st.title("📊 Data-Scaling Study — U-Net vs SegFormer-B0")
st.caption(
    "How does training-set size affect segmentation quality? "
    "Two architectures, three data shares (25 / 50 / 100 %), shared validation set. "
    "100% checkpoints are bootstrapped from the existing pv_panel_models baselines."
)

logs = load_all_logs()

if not logs:
    st.warning(
        "No logs found in `../logs/`. "
        "Run training first (`./run_all.sh`) and bootstrap "
        "the 100% point (`python bootstrap_100.py`)."
    )

tab_curves, tab_scaling, tab_infer = st.tabs(
    ["1 · Learning curves", "2 · Data share vs final", "3 · Inference"]
)


# ── Tab 1: Learning curves ─────────────────────────────────────────────────
with tab_curves:
    st.subheader("Per-epoch metrics")
    if any(is_bootstrapped(l) for l in logs.values()):
        st.info(
            "**Note on 100%:** per-epoch metrics are parsed from the existing text logs "
            "and use the old per-batch averaging (Dice/IoU/PixelAcc only). "
            "mIoU is null per epoch and is omitted from the chart for 100%. "
            "Use the scaling chart in tab 2 for fair cross-share comparisons."
        )

    if not logs:
        st.info("Waiting for training logs.")
    else:
        col_m, col_split = st.columns([2, 2])
        with col_m:
            metric_key, metric_label = st.selectbox(
                "Metric",
                METRICS,
                format_func=lambda x: x[1],
            )
        with col_split:
            split = st.radio("Split", ["val", "train", "both"], horizontal=True, index=0)

        for model in MODELS:
            available = [s for s in SHARES if (model, s) in logs]
            if not available:
                continue
            st.markdown(f"#### {PRETTY_MODEL[model]}")
            fig = go.Figure()

            for share in available:
                df = log_to_df(logs[(model, share)])
                bootstrapped = is_bootstrapped(logs[(model, share)])

                if split in ("val", "both"):
                    col = f"val_{metric_key}"
                    if col in df.columns and df[col].notna().any():
                        sub = df.dropna(subset=[col])
                        suffix = " val (old-def)" if bootstrapped else " val"
                        fig.add_trace(go.Scatter(
                            x=sub["epoch"], y=sub[col],
                            mode="lines",
                            name=f"{share}%{suffix}",
                        ))

                if split in ("train", "both"):
                    col = f"train_{metric_key}"
                    if col in df.columns and df[col].notna().any():
                        sub = df.dropna(subset=[col])
                        suffix = " train (old-def)" if bootstrapped else " train"
                        fig.add_trace(go.Scatter(
                            x=sub["epoch"], y=sub[col],
                            mode="lines", line=dict(dash="dot"),
                            name=f"{share}%{suffix}",
                        ))

            fig.update_layout(
                xaxis_title="Epoch",
                yaxis_title=metric_label,
                height=380,
                margin=dict(l=10, r=10, t=10, b=10),
                legend=dict(orientation="h", y=-0.2),
            )
            st.plotly_chart(fig, use_container_width=True)


# ── Tab 2: Data share vs final ─────────────────────────────────────────────
with tab_scaling:
    st.subheader("Best-checkpoint val metrics vs data share")
    st.caption(
        "Each point is the best-epoch validation score for one (model, share) run. "
        "All numbers use the same global confusion-matrix metric code, including the 100% "
        "points (recomputed via bootstrap_100.py)."
    )

    if not logs:
        st.info("Waiting for training logs.")
    else:
        rows = [scaling_row(log) for log in logs.values()]
        df = pd.DataFrame(rows).sort_values(["model", "share"]).reset_index(drop=True)

        # Display table — show formatted wall clock; hide raw seconds.
        display_df = df.drop(columns=["wall_clock_seconds", "sec_per_epoch"])
        st.dataframe(display_df, use_container_width=True, hide_index=True)

        # Timing summary
        trained_seconds = df.loc[df["source"] == "trained", "wall_clock_seconds"].sum()
        if trained_seconds:
            st.caption(
                f"⏱  Total training wall-clock across the four 25/50% runs: "
                f"**{fmt_hms(trained_seconds)}**  ({trained_seconds:,.0f} s)"
            )

        col1, col2 = st.columns(2)
        with col1:
            fig1 = px.line(
                df.dropna(subset=["best_val_miou"]),
                x="share", y="best_val_miou", color="model",
                markers=True, title="Best val mIoU",
                labels={"share": "Training data (%)", "best_val_miou": "Val mIoU"},
            )
            fig1.update_xaxes(tickvals=SHARES)
            st.plotly_chart(fig1, use_container_width=True)
        with col2:
            fig2 = px.line(
                df.dropna(subset=["best_val_dice"]),
                x="share", y="best_val_dice", color="model",
                markers=True, title="Best val Dice",
                labels={"share": "Training data (%)", "best_val_dice": "Val Dice"},
            )
            fig2.update_xaxes(tickvals=SHARES)
            st.plotly_chart(fig2, use_container_width=True)

        col3, col4 = st.columns(2)
        with col3:
            fig3 = px.bar(
                df.dropna(subset=["best_val_iou"]),
                x="share", y="best_val_iou", color="model", barmode="group",
                title="Best val foreground IoU",
                labels={"share": "Training data (%)", "best_val_iou": "Val IoU (foreground)"},
            )
            fig3.update_xaxes(tickvals=SHARES)
            st.plotly_chart(fig3, use_container_width=True)
        with col4:
            fig4 = px.bar(
                df.dropna(subset=["best_val_pixel_acc"]),
                x="share", y="best_val_pixel_acc", color="model", barmode="group",
                title="Best val pixel accuracy",
                labels={"share": "Training data (%)", "best_val_pixel_acc": "Val pixel acc"},
            )
            fig4.update_xaxes(tickvals=SHARES)
            st.plotly_chart(fig4, use_container_width=True)

        st.markdown("##### Training time")
        time_df = df[df["source"] == "trained"].dropna(subset=["wall_clock_seconds"])
        if not time_df.empty:
            time_df = time_df.assign(wall_minutes=time_df["wall_clock_seconds"] / 60.0)
            tcol1, tcol2 = st.columns(2)
            with tcol1:
                fig_t1 = px.bar(
                    time_df,
                    x="share", y="wall_minutes", color="model", barmode="group",
                    title="Total training time (minutes)",
                    labels={"share": "Training data (%)", "wall_minutes": "Wall clock (min)"},
                )
                fig_t1.update_xaxes(tickvals=SHARES)
                st.plotly_chart(fig_t1, use_container_width=True)
            with tcol2:
                fig_t2 = px.bar(
                    time_df.dropna(subset=["sec_per_epoch"]),
                    x="share", y="sec_per_epoch", color="model", barmode="group",
                    title="Average seconds per epoch",
                    labels={"share": "Training data (%)", "sec_per_epoch": "Seconds / epoch"},
                )
                fig_t2.update_xaxes(tickvals=SHARES)
                st.plotly_chart(fig_t2, use_container_width=True)
        else:
            st.caption("No timing data yet — runs in progress will populate this once the first one finishes.")


# ── Tab 3: Inference ───────────────────────────────────────────────────────
with tab_infer:
    st.subheader("Upload an image — see all 6 segmentations")
    st.caption(
        "Each cell uses the best-epoch checkpoint of one (model, data-share) combination. "
        "The 100% cells use the bootstrapped checkpoint (existing pv_panel_models baseline)."
    )

    col_a, col_b, col_c = st.columns([2, 2, 2])
    with col_a:
        threshold = st.slider("Threshold", 0.0, 1.0, 0.5, 0.05, key="infer_thr")
    with col_b:
        view = st.radio(
            "View",
            ["mask", "overlay", "heatmap"],
            horizontal=True,
            key="infer_view",
        )
    with col_c:
        cell_w = st.select_slider(
            "Cell size (px)", options=[160, 200, 240, 280, 320], value=240, key="infer_cell"
        )

    uploaded = st.file_uploader(
        "Drop an image (jpg/png)", type=["jpg", "jpeg", "png"], key="infer_upload"
    )

    debug = st.checkbox("debug", value=True, key="infer_debug")

    if uploaded is not None:
        if debug:
            st.write("✓ uploaded is not None — entering inference block")

        device = "cuda" if torch.cuda.is_available() else "cpu"
        if debug:
            st.write(f"✓ device = `{device}`")

        try:
            raw_bytes = uploaded.getvalue()
            img = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
        except Exception as e:
            st.error(f"Could not decode uploaded image: {e}")
            st.exception(e)
            st.stop()

        if debug:
            st.write(f"✓ image decoded — {img.size[0]}×{img.size[1]} px, {len(raw_bytes)/1024:.1f} KB")

        st.caption(f"📁 `{uploaded.name}` — {img.size[0]}×{img.size[1]} px, {len(raw_bytes)/1024:.1f} KB")

        # Input preview (no nested columns — flat render so nothing gets swallowed)
        st.markdown("**Input (original / resized to 128×128 the models see)**")
        try:
            x = preprocess(img, image_size=128)
            rgb_small = (x.squeeze().permute(1, 2, 0).numpy() * 255).astype(np.uint8)
        except Exception as e:
            st.error(f"preprocess failed: {e}")
            st.exception(e)
            st.stop()

        st.image([img, rgb_small], width=cell_w, caption=["original", "128×128"])

        if debug:
            st.write(f"✓ tensor shape = {tuple(x.shape)}, rgb_small shape = {rgb_small.shape}")
            st.write(f"✓ MODELS = {MODELS}, SHARES = {SHARES}")

        st.markdown("##### Predictions (one row per model+share)")

        def render_cell(probs, mask, rgb):
            if view == "mask":
                return (mask * 255).astype(np.uint8)
            if view == "overlay":
                return overlay(rgb, mask)
            return heatmap(probs)

        # Single flat row of 6 cells — no nested columns.
        cells = []
        for model_name in MODELS:
            for share in SHARES:
                cells.append((model_name, share))

        if debug:
            st.write(f"✓ rendering {len(cells)} cells: {cells}")

        cols = st.columns(len(cells))
        for col, (model_name, share) in zip(cols, cells):
            with col:
                st.markdown(f"**{PRETTY_MODEL[model_name]}**  \n*{share}%*")
                try:
                    if debug:
                        st.write(f"loading {model_name}_{share}…")
                    m = load_best(model_name, share, device)
                    if m is None:
                        st.warning(f"missing `{model_name}_{share}_best.pth`")
                        continue
                    if debug:
                        st.write("running…")
                    probs, mask = run_inference(m, x, device, threshold=threshold)
                    cell_img = render_cell(probs, mask, rgb_small)
                    st.image(cell_img, width=cell_w)
                    st.caption(
                        f"cov={float(mask.mean())*100:.1f}%   "
                        f"p[{probs.min():.2f},{probs.max():.2f}]"
                    )
                except Exception as e:
                    st.error(f"{model_name} {share}% failed")
                    st.exception(e)

        if debug:
            st.write("✓ render loop complete")
    else:
        st.info("Upload an image to run inference across all six trained models.")