app.py

import gradio as gr
import pandas as pd
import folium
import numpy as np
import os
import re
import json

BASE        = os.path.dirname(os.path.abspath(__file__)) if "__file__" in dir() else os.getcwd()
STAY_POINTS = os.path.join(BASE, "data", "stay_points_inference_sample.csv")
POI_PATH    = os.path.join(BASE, "data", "poi_inference_sample.csv")
DEMO_PATH   = os.path.join(BASE, "data", "demographics_inference_sample.csv")
COT_PATH    = os.path.join(BASE, "data", "inference_results_sample.json")

SEX_MAP  = {1:"Male", 2:"Female", -8:"Unknown", -7:"Prefer not to answer"}
EDU_MAP  = {1:"Less than HS", 2:"HS Graduate/GED", 3:"Some College/Associate",
            4:"Bachelor's Degree", 5:"Graduate/Professional Degree",
            -1:"N/A", -7:"Prefer not to answer", -8:"Unknown"}
INC_MAP  = {1:"<$10,000", 2:"$10,000–$14,999", 3:"$15,000–$24,999",
            4:"$25,000–$34,999", 5:"$35,000–$49,999", 6:"$50,000–$74,999",
            7:"$75,000–$99,999", 8:"$100,000–$124,999", 9:"$125,000–$149,999",
            10:"$150,000–$199,999", 11:"$200,000+",
            -7:"Prefer not to answer", -8:"Unknown", -9:"Not ascertained"}
RACE_MAP = {1:"White", 2:"Black or African American", 3:"Asian",
            4:"American Indian or Alaska Native",
            5:"Native Hawaiian or Other Pacific Islander",
            6:"Multiple races", 97:"Other",
            -7:"Prefer not to answer", -8:"Unknown"}
ACT_MAP  = {0:"Transportation", 1:"Home", 2:"Work", 3:"School", 4:"ChildCare",
            5:"BuyGoods", 6:"Services", 7:"EatOut", 8:"Errands", 9:"Recreation",
            10:"Exercise", 11:"Visit", 12:"HealthCare", 13:"Religious",
            14:"SomethingElse", 15:"DropOff"}

print("Loading data...")
sp   = pd.read_csv(STAY_POINTS)
poi  = pd.read_csv(POI_PATH)
demo = pd.read_csv(DEMO_PATH)

sp = sp.merge(poi, on="poi_id", how="left")
sp["start_datetime"] = pd.to_datetime(sp["start_datetime"], utc=True)
sp["end_datetime"]   = pd.to_datetime(sp["end_datetime"],   utc=True)
sp["duration_min"]   = ((sp["end_datetime"] - sp["start_datetime"]).dt.total_seconds() / 60).round(1)

def parse_act_types(x):
    try:
        codes = list(map(int, str(x).strip("[]").split()))
        return ", ".join(ACT_MAP.get(c, str(c)) for c in codes)
    except:
        return str(x)

sp["act_label"] = sp["act_types"].apply(parse_act_types)

# Load CoT JSON (optional)
cot_by_agent = {}
if os.path.exists(COT_PATH):
    with open(COT_PATH, "r") as f:
        cot_raw = json.load(f)
    # Support both list and {"inference_results": [...]} formats
    records = cot_raw if isinstance(cot_raw, list) else cot_raw.get("inference_results", [])
    for result in records:
        cot_by_agent[int(result["agent_id"])] = result
    print(f"Loaded CoT for {len(cot_by_agent)} agents.")

sample_agents = sorted(sp["agent_id"].unique().tolist())
print(f"Ready. {len(sample_agents)} agents loaded.")


def get_cot(agent_id):
    result = cot_by_agent.get(int(agent_id), {})
    s1 = result.get("step1_response", "")
    s2 = result.get("step2_response", "")
    s3 = result.get("step3_response", "")
    return s1, s2, s3


# ── Mobility text builders ────────────────────────────────────────────────────

def build_mobility_summary(agent_sp):
    top5 = (agent_sp.groupby("name")["duration_min"]
            .agg(visits="count", avg_dur="mean")
            .sort_values("visits", ascending=False)
            .head(5))

    obs_start = agent_sp["start_datetime"].min().strftime("%Y-%m-%d")
    obs_end   = agent_sp["end_datetime"].max().strftime("%Y-%m-%d")
    days      = (agent_sp["end_datetime"].max() - agent_sp["start_datetime"].min()).days

    # Top activity types
    act_counts = agent_sp["act_label"].value_counts().head(3)
    top_acts = ", ".join(f"{a} ({n})" for a, n in act_counts.items())

    # Time of day
    agent_sp2 = agent_sp.copy()
    agent_sp2["hour"] = agent_sp2["start_datetime"].dt.hour
    def tod(h):
        if 5 <= h < 12:  return "Morning"
        if 12 <= h < 17: return "Afternoon"
        if 17 <= h < 21: return "Evening"
        return "Night"
    agent_sp2["tod"] = agent_sp2["hour"].apply(tod)
    peak_tod = agent_sp2["tod"].value_counts().idxmax()
    agent_sp2["is_weekend"] = agent_sp2["start_datetime"].dt.dayofweek >= 5
    wd_pct = int((~agent_sp2["is_weekend"]).mean() * 100)

    lines = [
        f"Period: {obs_start} ~ {obs_end}  ({days} days)",
        f"Stay points: {len(agent_sp)}  |  Unique locations: {agent_sp['name'].nunique()}",
        f"Weekday/Weekend: {wd_pct}% / {100-wd_pct}%  |  Peak time: {peak_tod}",
        f"Top activities: {top_acts}",
        "",
        "Top Locations:",
    ]
    for i, (name, row) in enumerate(top5.iterrows(), 1):
        lines.append(f"  {i}. {name}  —  {int(row['visits'])} visits, avg {int(row['avg_dur'])} min")

    return "\n".join(lines)


def build_weekly_checkin(agent_sp, max_days=None):
    agent_sp2 = agent_sp.copy()
    agent_sp2["date"] = agent_sp2["start_datetime"].dt.date
    all_dates = sorted(agent_sp2["date"].unique())
    dates_to_show = all_dates[:max_days] if max_days else all_dates
    total_days = len(all_dates)

    lines = ["WEEKLY CHECK-IN SUMMARY", "======================="]
    for date in dates_to_show:
        grp = agent_sp2[agent_sp2["date"] == date]
        dow   = grp["start_datetime"].iloc[0].strftime("%A")
        label = "Weekend" if grp["start_datetime"].iloc[0].dayofweek >= 5 else "Weekday"
        lines.append(f"\n--- {dow}, {date} ({label}) ---")
        lines.append(f"Total activities: {len(grp)}")
        for _, row in grp.iterrows():
            lines.append(
                f"- {row['start_datetime'].strftime('%H:%M')}-"
                f"{row['end_datetime'].strftime('%H:%M')} "
                f"({int(row['duration_min'])} mins): "
                f"{row['name']} - {row['act_label']}"
            )
    if max_days and total_days > max_days:
        lines.append(f"\n... ({total_days - max_days} more days)")
    return "\n".join(lines)


# ── HTML reasoning chain ──────────────────────────────────────────────────────

# ── Paste this entire block into app.py, replacing the existing CHAIN_CSS, render_chain, and helper functions ──

import re

CHAIN_CSS = """
<style>
@import url('https://fonts.googleapis.com/css2?family=DM+Mono:wght@400;500&family=DM+Sans:wght@300;400;500;600&display=swap');

.hct-root {
  font-family: 'DM Sans', sans-serif;
  display: flex;
  flex-direction: column;
  gap: 0;
  padding: 4px 0 8px;
}

/* ── Stage shell ── */
.hct-stage {
  border-radius: 12px;
  overflow: hidden;
  transition: opacity 0.3s, filter 0.3s;
}
.hct-stage.dim  { opacity: 0.28; filter: grayscale(0.6); pointer-events: none; }
.hct-stage.active { opacity: 1; }

/* ── Stage header strip ── */
.hct-head {
  display: flex;
  align-items: center;
  gap: 10px;
  padding: 9px 14px;
}
.hct-num {
  font-family: 'DM Mono', monospace;
  font-size: 10px;
  font-weight: 500;
  letter-spacing: 0.12em;
  padding: 2px 7px;
  border-radius: 4px;
}
.hct-title {
  font-size: 11px;
  font-weight: 600;
  letter-spacing: 0.04em;
  text-transform: uppercase;
}

/* Stage 1 colors */
.hct-s1 { background: #f4f6fb; border: 1.5px solid #d4daf0; }
.hct-s1 .hct-head { background: #eaecf7; border-bottom: 1px solid #d4daf0; }
.hct-s1 .hct-num { background: #dde2f3; color: #3a4a80; }
.hct-s1 .hct-title { color: #3a4a80; }

/* Stage 2 colors */
.hct-s2 { background: #fdf8f2; border: 1.5px solid #e8d5b8; }
.hct-s2 .hct-head { background: #f7ede0; border-bottom: 1px solid #e8d5b8; }
.hct-s2 .hct-num { background: #f0dcbf; color: #7a4a10; }
.hct-s2 .hct-title { color: #7a4a10; }

/* Stage 3 colors */
.hct-s3 { background: #fff6f5; border: 2px solid #d4453a; }
.hct-s3 .hct-head { background: #fce8e7; border-bottom: 1px solid #d4453a; }
.hct-s3 .hct-num { background: #d4453a; color: #fff; }
.hct-s3 .hct-title { color: #b0302a; }

/* ── Body ── */
.hct-body { padding: 12px 14px; }

/* ── Arrow connector ── */
.hct-arrow {
  display: flex;
  align-items: center;
  gap: 8px;
  padding: 5px 18px;
  transition: opacity 0.3s;
}
.hct-arrow-line { flex: 1; height: 1px; background: #d8d4ce; }
.hct-arrow-label {
  font-family: 'DM Mono', monospace;
  font-size: 9px;
  color: #b0a898;
  letter-spacing: 0.08em;
  text-transform: uppercase;
  white-space: nowrap;
  background: white;
  padding: 2px 8px;
  border: 1px solid #e0dbd4;
  border-radius: 20px;
}

/* ── Stage 1: Location table ── */
.hct-loc-table {
  width: 100%;
  border-collapse: collapse;
  font-size: 11.5px;
  margin-bottom: 10px;
}
.hct-loc-table th {
  font-family: 'DM Mono', monospace;
  font-size: 9px;
  font-weight: 500;
  letter-spacing: 0.1em;
  text-transform: uppercase;
  color: #8090b0;
  border-bottom: 1px solid #d4daf0;
  padding: 3px 6px 5px;
  text-align: left;
}
.hct-loc-table th:not(:first-child) { text-align: right; }
.hct-loc-table td {
  padding: 5px 6px;
  color: #2a3050;
  border-bottom: 1px solid #eaecf5;
  line-height: 1.3;
}
.hct-loc-table td:not(:first-child) { text-align: right; font-family: 'DM Mono', monospace; font-size: 11px; color: #5060a0; }
.hct-loc-table tr:last-child td { border-bottom: none; }
.hct-loc-name { font-weight: 500; max-width: 160px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; display: block; }
.hct-visit-bar-wrap { display: flex; align-items: center; gap: 6px; justify-content: flex-end; }
.hct-visit-bar { height: 4px; border-radius: 2px; background: #6878c8; opacity: 0.55; }

/* ── Stage 1: Temporal panel ── */
.hct-temporal {
  display: grid;
  grid-template-columns: 1fr 1fr;
  gap: 8px;
}
.hct-temp-block {
  background: #eef0fa;
  border-radius: 8px;
  padding: 8px 10px;
}
.hct-temp-label {
  font-family: 'DM Mono', monospace;
  font-size: 9px;
  font-weight: 500;
  letter-spacing: 0.1em;
  text-transform: uppercase;
  color: #7080b0;
  margin-bottom: 6px;
}
.hct-seg-row { display: flex; height: 10px; border-radius: 5px; overflow: hidden; margin-bottom: 5px; }
.hct-seg { display: flex; align-items: center; justify-content: center; font-size: 0; transition: width 0.5s; }
.seg-morning  { background: #fbbf24; }
.seg-afternoon{ background: #f97316; }
.seg-evening  { background: #8b5cf6; }
.seg-night    { background: #1e3a5f; }
.seg-weekday  { background: #6878c8; }
.seg-weekend  { background: #e8c080; }
.hct-legend { display: flex; flex-wrap: wrap; gap: 4px 10px; }
.hct-leg-item { display: flex; align-items: center; gap: 4px; font-size: 10px; color: #5a6080; }
.hct-leg-dot  { width: 8px; height: 8px; border-radius: 2px; flex-shrink: 0; }
.hct-dist-line {
  margin-top: 8px;
  font-size: 11px;
  color: #6070a0;
  font-family: 'DM Mono', monospace;
  padding: 5px 8px;
  background: #eef0fa;
  border-radius: 6px;
  display: flex;
  align-items: center;
  gap: 6px;
}

/* ── Stage 2: 2×2 grid ── */
.hct-dim-grid {
  display: grid;
  grid-template-columns: 1fr 1fr;
  gap: 8px;
}
.hct-dim-card {
  background: #fff;
  border: 1px solid #e8d5b8;
  border-radius: 8px;
  padding: 9px 11px;
}
.hct-dim-head {
  display: flex;
  align-items: center;
  gap: 6px;
  margin-bottom: 5px;
}
.hct-dim-icon { font-size: 13px; line-height: 1; }
.hct-dim-name {
  font-family: 'DM Mono', monospace;
  font-size: 9px;
  font-weight: 500;
  letter-spacing: 0.1em;
  text-transform: uppercase;
  color: #a07040;
}
.hct-dim-text {
  font-size: 11px;
  color: #3a2a10;
  line-height: 1.55;
}
.hct-dim-empty { color: #ccc; font-style: italic; }

/* ── Stage 3: prediction ── */
.hct-pred-row {
  display: flex;
  align-items: flex-start;
  gap: 16px;
  margin-bottom: 10px;
}
.hct-pred-badge {
  background: #d4453a;
  color: white;
  border-radius: 8px;
  padding: 8px 14px;
  text-align: center;
  flex-shrink: 0;
}
.hct-pred-val { font-size: 18px; font-weight: 600; line-height: 1.2; white-space: nowrap; }
.hct-pred-sub { font-family: 'DM Mono', monospace; font-size: 9px; opacity: 0.8; letter-spacing: 0.08em; text-transform: uppercase; margin-top: 2px; }
.hct-conf-col { flex: 1; padding-top: 4px; }
.hct-conf-label { font-family: 'DM Mono', monospace; font-size: 9px; color: #a04040; letter-spacing: 0.08em; text-transform: uppercase; margin-bottom: 4px; }
.hct-conf-track { height: 6px; background: #f0d0cf; border-radius: 3px; overflow: hidden; margin-bottom: 6px; }
.hct-conf-fill  { height: 100%; background: linear-gradient(90deg, #e74c3c, #8b0000); border-radius: 3px; }
.hct-reasoning {
  font-size: 11.5px;
  color: #4a2020;
  line-height: 1.6;
  border-left: 3px solid #e8b0ae;
  padding-left: 10px;
}

/* ── Idle / loading states ── */
.hct-idle { font-size: 12px; color: #b0bac8; padding: 6px 0; font-style: italic; }
.hct-loading {
  font-size: 12px; padding: 6px 0;
  display: flex; align-items: center; gap: 8px;
}
.hct-dot { width: 6px; height: 6px; border-radius: 50%; display: inline-block; animation: hct-pulse 1.2s ease-in-out infinite; }
.hct-dot:nth-child(2) { animation-delay: 0.2s; }
.hct-dot:nth-child(3) { animation-delay: 0.4s; }
@keyframes hct-pulse {
  0%,100% { opacity: 0.2; transform: scale(0.8); }
  50% { opacity: 1; transform: scale(1.1); }
}
.hct-s1 .hct-dot { background: #6878c8; }
.hct-s2 .hct-dot { background: #c08040; }
.hct-s3 .hct-dot { background: #d4453a; }
</style>
"""


def _loading(msg):
    return f'<div class="hct-loading"><span class="hct-dot"></span><span class="hct-dot"></span><span class="hct-dot"></span><span style="color:#8090a0;font-size:12px">{msg}</span></div>'


def _parse_s1(text):
    """Returns (locations, tod, wk, dist)"""
    locations = []
    dur_map = {}
    tod = {}
    wk = {}
    dist = None

    for line in text.splitlines():
        s = line.strip()
        # Location inventory: "- Name: N visits, ..."
        m = re.match(r'-\s+(.+?):\s+(\d+)\s+visit', s, re.IGNORECASE)
        if m:
            locations.append((m.group(1).strip(), int(m.group(2))))
        # Duration: "- LocationName: Average duration of X minutes"
        m2 = re.match(r'-?\s*(.+?):\s+Average duration of ([\d.]+)\s+min', s, re.IGNORECASE)
        if m2:
            dur_map[m2.group(1).strip()] = float(m2.group(2))
        # TOD: "65% morning, 23% afternoon, 6% evening, 5% night"
        if not tod:
            m3 = re.search(r'(\d+)%\s*morning.*?(\d+)%\s*afternoon.*?(\d+)%\s*evening.*?(\d+)%\s*night', s, re.IGNORECASE)
            if m3:
                tod = {'Morning': int(m3.group(1)), 'Afternoon': int(m3.group(2)),
                       'Evening': int(m3.group(3)), 'Night': int(m3.group(4))}
        # Weekday/weekend
        if not wk:
            m4 = re.search(r'(\d+)%\s*weekday.*?(\d+)%\s*weekend', s, re.IGNORECASE)
            if m4:
                wk = {'Weekday': int(m4.group(1)), 'Weekend': int(m4.group(2))}
        # Distance
        if not dist:
            m5 = re.search(r'average distance of approximately ([\d.]+)\s*miles', s, re.IGNORECASE)
            if m5:
                dist = float(m5.group(1))

    result_locs = [(n, v, dur_map.get(n)) for n, v in locations[:7]]
    return result_locs, tod, wk, dist


def _parse_s2(text):
    """Returns dict: ROUTINE, ECONOMIC, SOCIAL, URBAN, STABILITY → short summary string"""
    DIMS = {
        'ROUTINE':    ['ROUTINE', 'SCHEDULE'],
        'ECONOMIC':   ['ECONOMIC', 'SPENDING'],
        'SOCIAL':     ['SOCIAL', 'LIFESTYLE'],
        'URBAN':      ['URBAN', 'COMMUNITY'],
        'STABILITY':  ['STABILITY', 'REGULARITY', 'CONSISTENCY'],
    }
    sections = {}
    current_key = None
    current_lines = []

    for line in text.splitlines():
        s = line.strip()
        # Format A: "1. TITLE ANALYSIS:" or "2. ECONOMIC BEHAVIOR PATTERNS:"
        mA = re.match(r'^\d+\.\s+([A-Z][A-Z\s&]+?)(?:\s+ANALYSIS|\s+PATTERNS|\s+INDICATORS|\s+CHARACTERISTICS|\s+STABILITY)?:\s*$', s, re.IGNORECASE)
        # Format B: "STEP 1: ROUTINE & SCHEDULE ANALYSIS"
        mB = re.match(r'^STEP\s+\d+:\s+([A-Z][A-Z\s&]+?)(?:\s+ANALYSIS|\s+PATTERNS|\s+INDICATORS|\s+CHARACTERISTICS|\s+STABILITY)?\s*$', s, re.IGNORECASE)
        mm = mA or mB
        if mm:
            if current_key and current_lines:
                sections[current_key] = ' '.join(current_lines)
            current_key = mm.group(1).upper().strip()
            current_lines = []
        elif current_key and s:
            if re.match(r'^\d+\.\d+', s):
                sub = re.sub(r'^\d+\.\d+[^:]*:\s*', '', s)
                if sub:
                    current_lines.append(sub)
            elif s.startswith('-'):
                current_lines.append(s.lstrip('-').strip())
            elif not re.match(r'^\d+\.', s):
                current_lines.append(s)

    if current_key and current_lines:
        sections[current_key] = ' '.join(current_lines)

    result = {}
    for dim, keywords in DIMS.items():
        for k, txt in sections.items():
            if any(kw in k for kw in keywords) and txt:
                sents = re.split(r'(?<=[.!?])\s+', txt.strip())
                summary = ' '.join(sents[:2])
                if len(summary) > 160:
                    summary = summary[:157] + '…'
                result[dim] = summary
                break
    return result


def _parse_s3(text):
    pred, conf, reasoning = '', 0, ''
    in_r = False
    r_lines = []
    for line in text.splitlines():
        s = line.strip()
        if s.startswith('INCOME_PREDICTION:'):
            pred = s.replace('INCOME_PREDICTION:', '').strip()
        elif s.startswith('INCOME_CONFIDENCE:'):
            try:
                conf = int(re.search(r'\d+', s).group())
            except:
                conf = 0
        elif s.startswith('INCOME_REASONING:'):
            in_r = True
            r_lines.append(s.replace('INCOME_REASONING:', '').strip())
        elif in_r:
            if re.match(r'^2\.', s) or s.startswith('INCOME_'):
                break
            if s:
                r_lines.append(s)
    reasoning = ' '.join(r_lines).strip()
    sents = re.split(r'(?<=[.!?])\s+', reasoning)
    reasoning = ' '.join(sents[:3])
    if len(reasoning) > 280:
        reasoning = reasoning[:277] + '…'
    return pred, conf, reasoning


def _s1_body(text, active):
    if not active:
        return '<div class="hct-idle">Press ▶ to start</div>'
    if not text:
        return _loading('Extracting features')
    locs, tod, wk, dist = _parse_s1(text)

    # Location table
    max_v = max((v for _, v, _ in locs), default=1)
    rows = ''
    for name, visits, dur in locs:
        bar_w = int(60 * visits / max_v)
        dur_str = f'{int(dur)}m' if dur else '—'
        rows += (
            f'<tr>'
            f'<td><span class="hct-loc-name" title="{name}">{name}</span></td>'
            f'<td><div class="hct-visit-bar-wrap">'
            f'<div class="hct-visit-bar" style="width:{bar_w}px"></div>'
            f'{visits}</div></td>'
            f'<td>{dur_str}</td>'
            f'</tr>'
        )
    table = (
        f'<table class="hct-loc-table">'
        f'<thead><tr><th>Location</th><th>Visits</th><th>Avg Stay</th></tr></thead>'
        f'<tbody>{rows}</tbody>'
        f'</table>'
    ) if rows else ''

    # Temporal panels
    def seg_bar(data, seg_classes):
        total = sum(data.values()) or 1
        segs = ''.join(
            f'<div class="hct-seg {cls}" style="width:{int(100*v/total)}%"></div>'
            for (label, v), cls in zip(data.items(), seg_classes)
        )
        legend = ''.join(
            f'<div class="hct-leg-item"><div class="hct-leg-dot {cls}"></div>{label} {v}%</div>'
            for (label, v), cls in zip(data.items(), seg_classes)
        )
        return f'<div class="hct-seg-row">{segs}</div><div class="hct-legend">{legend}</div>'

    tod_panel = ''
    if tod:
        tod_panel = (
            f'<div class="hct-temp-block">'
            f'<div class="hct-temp-label">Time of Day</div>'
            f'{seg_bar(tod, ["seg-morning","seg-afternoon","seg-evening","seg-night"])}'
            f'</div>'
        )
    wk_panel = ''
    if wk:
        wk_panel = (
            f'<div class="hct-temp-block">'
            f'<div class="hct-temp-label">Weekday / Weekend</div>'
            f'{seg_bar(wk, ["seg-weekday","seg-weekend"])}'
            f'</div>'
        )
    temporal = f'<div class="hct-temporal">{tod_panel}{wk_panel}</div>' if (tod_panel or wk_panel) else ''

    dist_line = ''
    if dist:
        dist_line = f'<div class="hct-dist-line">📍 Avg trip distance &nbsp;{dist} mi</div>'

    return table + temporal + dist_line


def _s2_body(text, active):
    if not active:
        return '<div class="hct-idle">Waiting…</div>'
    if not text:
        return _loading('Analyzing behavior')
    dims = _parse_s2(text)

    DIM_META = [
        ('ROUTINE',   '🕐', 'Schedule'),
        ('ECONOMIC',  '💰', 'Economic'),
        ('SOCIAL',    '👥', 'Social'),
        ('STABILITY', '🔄', 'Stability'),
    ]
    # fallback to URBAN if STABILITY missing
    if 'STABILITY' not in dims and 'URBAN' in dims:
        dims['STABILITY'] = dims['URBAN']

    cards = ''
    for key, icon, label in DIM_META:
        txt = dims.get(key, '')
        content = f'<div class="hct-dim-text">{txt}</div>' if txt else '<div class="hct-dim-text hct-dim-empty">—</div>'
        cards += (
            f'<div class="hct-dim-card">'
            f'<div class="hct-dim-head">'
            f'<span class="hct-dim-icon">{icon}</span>'
            f'<span class="hct-dim-name">{label}</span>'
            f'</div>'
            f'{content}'
            f'</div>'
        )
    return f'<div class="hct-dim-grid">{cards}</div>'


def _s3_body(text, active):
    if not active:
        return '<div class="hct-idle">Waiting…</div>'
    if not text:
        return _loading('Inferring demographics')
    pred, conf, reasoning = _parse_s3(text)
    conf_pct = int(conf / 5 * 100)
    return (
        f'<div class="hct-pred-row">'
        f'<div class="hct-pred-badge">'
        f'<div class="hct-pred-val">{pred or "—"}</div>'
        f'<div class="hct-pred-sub">Income</div>'
        f'</div>'
        f'<div class="hct-conf-col">'
        f'<div class="hct-conf-label">Confidence &nbsp;{conf}/5</div>'
        f'<div class="hct-conf-track"><div class="hct-conf-fill" style="width:{conf_pct}%"></div></div>'
        f'</div>'
        f'</div>'
        f'<div class="hct-reasoning">{reasoning}</div>'
    )


def render_chain(s1_text, s2_text, s3_text, status="idle"):
    s1_on = status in ("running1", "running2", "running3", "done")
    s2_on = status in ("running2", "running3", "done")
    s3_on = status in ("running3", "done")

    # For "running" states the text may be empty → show loading dots
    s1_body = _s1_body(s1_text if s1_on else '', s1_on)
    s2_body = _s2_body(s2_text if s2_on else '', s2_on)
    s3_body = _s3_body(s3_text if s3_on else '', s3_on)

    def stage(cls, num, title, body, on):
        dim_cls = 'active' if on else 'dim'
        return (
            f'<div class="hct-stage hct-{cls} {dim_cls}">'
            f'<div class="hct-head">'
            f'<span class="hct-num">{num}</span>'
            f'<span class="hct-title">{title}</span>'
            f'</div>'
            f'<div class="hct-body">{body}</div>'
            f'</div>'
        )

    def arrow(label, on):
        op = '1' if on else '0.2'
        return (
            f'<div class="hct-arrow" style="opacity:{op}">'
            f'<div class="hct-arrow-line"></div>'
            f'<div class="hct-arrow-label">{label}</div>'
            f'<div class="hct-arrow-line"></div>'
            f'</div>'
        )

    html = CHAIN_CSS + '<div class="hct-root">'
    html += stage('s1', 'Stage 01', 'Feature Extraction', s1_body, s1_on)
    html += arrow('behavioral abstraction', s2_on)
    html += stage('s2', 'Stage 02', 'Behavioral Analysis', s2_body, s2_on)
    html += arrow('demographic inference', s3_on)
    html += stage('s3', 'Stage 03', 'Demographic Inference', s3_body, s3_on)
    html += '</div>'
    return html


def build_map(agent_sp):
    agent_sp = agent_sp.reset_index(drop=True).copy()
    agent_sp["latitude"]  += np.random.uniform(-0.0003, 0.0003, len(agent_sp))
    agent_sp["longitude"] += np.random.uniform(-0.0003, 0.0003, len(agent_sp))

    lat = agent_sp["latitude"].mean()
    lon = agent_sp["longitude"].mean()
    m = folium.Map(location=[lat, lon], zoom_start=12, tiles="CartoDB positron")

    coords = list(zip(agent_sp["latitude"], agent_sp["longitude"]))
    if len(coords) > 1:
        folium.PolyLine(coords, color="#cc000055", weight=1.5, opacity=0.4).add_to(m)

    n = len(agent_sp)
    for i, row in agent_sp.iterrows():
        ratio = i / max(n - 1, 1)
        r = int(255 - ratio * (255 - 139))
        g = int(204 * (1 - ratio) ** 2)
        b = 0
        color = f"#{r:02x}{g:02x}{b:02x}"
        folium.CircleMarker(
            location=[row["latitude"], row["longitude"]],
            radius=7, color=color, fill=True, fill_color=color, fill_opacity=0.9,
            popup=folium.Popup(
                f"<b>#{i+1} {row['name']}</b><br>"
                f"{row['start_datetime'].strftime('%a %m/%d %H:%M')}<br>"
                f"{int(row['duration_min'])} min<br>{row['act_label']}",
                max_width=220
            )
        ).add_to(m)

    legend_html = """
    <div style="
        position:fixed; bottom:18px; left:18px; z-index:9999;
        background:rgba(255,255,255,0.92); border-radius:8px;
        padding:8px 12px; font-size:11px; font-family:sans-serif;
        box-shadow:0 1px 5px rgba(0,0,0,0.2); line-height:1.8;
    ">
        <div style="font-weight:600;margin-bottom:4px;">Stay Point Legend</div>
        <div style="display:flex;align-items:center;gap:6px;">
            <svg width="60" height="10">
              <defs><linearGradient id="lg" x1="0" x2="1" y1="0" y2="0">
                <stop offset="0%" stop-color="#ffcc00"/>
                <stop offset="100%" stop-color="#8b0000"/>
              </linearGradient></defs>
              <rect width="60" height="10" rx="4" fill="url(#lg)"/>
            </svg>
            <span>Earlier &rarr; Later</span>
        </div>
        <div style="display:flex;align-items:center;gap:6px;margin-top:2px;">
            <svg width="14" height="14"><circle cx="7" cy="7" r="5" fill="#cc4444" opacity="0.5"/></svg>
            <span>Movement path</span>
        </div>
        <div style="color:#999;font-size:10px;margin-top:2px;">Click dot for details</div>
    </div>
    """
    m.get_root().html.add_child(folium.Element(legend_html))
    m.get_root().width = "100%"
    m.get_root().height = "420px"
    return m._repr_html_()


def build_demo_text(row):
    age = int(row["age"]) if row["age"] > 0 else "Unknown"
    return (
        f"Age: {age}  |  "
        f"Sex: {SEX_MAP.get(int(row['sex']), row['sex'])}  |  "
        f"Race: {RACE_MAP.get(int(row['race']), row['race'])}  |  "
        f"Education: {EDU_MAP.get(int(row['education']), row['education'])}  |  "
        f"Income: {INC_MAP.get(int(row['hh_income']), row['hh_income'])}"
    )


# ── Callbacks ─────────────────────────────────────────────────────────────────

def on_select(agent_id):
    agent_id   = int(agent_id)
    agent_sp   = sp[sp["agent_id"] == agent_id].sort_values("start_datetime")
    agent_demo = demo[demo["agent_id"] == agent_id].iloc[0]

    map_html  = build_map(agent_sp)
    demo_text = build_demo_text(agent_demo)
    raw_text  = build_mobility_summary(agent_sp) + "\n\n" + build_weekly_checkin(agent_sp)
    chain_html = render_chain("", "", "", status="idle")

    return map_html, raw_text, demo_text, chain_html


def run_step(agent_id, step):
    """Reveal one stage per click. step: 0->1->2->done(-1)"""
    agent_id = int(agent_id)
    s1, s2, s3 = get_cot(agent_id)
    next_step = step + 1
    if next_step == 1:
        html = render_chain(s1, "", "", status="running2")
        label = "▶ Stage 2: Behavioral Analysis"
        return html, 1, gr.update(value=label)
    elif next_step == 2:
        html = render_chain(s1, s2, "", status="running3")
        label = "▶ Stage 3: Demographic Inference"
        return html, 2, gr.update(value=label)
    else:
        html = render_chain(s1, s2, s3, status="done")
        return html, -1, gr.update(value="↺ Reset")


def handle_btn(agent_id, step):
    if step == -1:
        html = render_chain("", "", "", status="idle")
        return html, 0, gr.update(value="▶ Stage 1: Feature Extraction")
    return run_step(agent_id, step)


def on_select_reset(agent_id):
    agent_id_int = int(agent_id)
    agent_sp   = sp[sp["agent_id"] == agent_id_int].sort_values("start_datetime")
    agent_demo = demo[demo["agent_id"] == agent_id_int].iloc[0]
    map_html   = build_map(agent_sp)
    demo_text  = build_demo_text(agent_demo)
    cot_entry  = cot_by_agent.get(agent_id_int, {})
    summary    = build_mobility_summary(agent_sp)
    raw_full   = cot_entry.get("weekly_checkin") or build_weekly_checkin(agent_sp)
    sep = "\n\n--- "
    parts = raw_full.split(sep)
    extra = len(parts) - 1
    raw_text = parts[0] + (sep.join([""] + parts[1:2]) + ("\n\n... ({} more days)".format(extra - 1) if extra > 1 else "")) if extra > 0 else raw_full
    chain_html = render_chain("", "", "", status="idle")
    return map_html, summary, raw_text, demo_text, chain_html, 0, gr.update(value="▶ Stage 1: Feature Extraction")



SHOWCASE_AGENTS = sample_agents[:6]


def build_agent_cards(selected_id):
    selected_id = int(selected_id)
    parts = []
    parts.append("<div style='display:grid;grid-template-columns:repeat(3,1fr);gap:10px;padding:4px 0;'>")
    for aid in SHOWCASE_AGENTS:
        row = demo[demo["agent_id"] == aid].iloc[0]
        age = int(row["age"]) if row["age"] > 0 else "?"
        sex = SEX_MAP.get(int(row["sex"]), "?")
        edu = EDU_MAP.get(int(row["education"]), "?")
        inc = INC_MAP.get(int(row["hh_income"]), "?")
        is_sel = (aid == selected_id)
        sel_style = "border:2px solid #c0392b;background:#fff8f8;box-shadow:0 2px 8px rgba(192,57,43,0.15);"
        nor_style = "border:1.5px solid #ddd;background:#fafafa;box-shadow:0 1px 3px rgba(0,0,0,0.06);"
        style = sel_style if is_sel else nor_style
        dot = "<span style='display:inline-block;width:8px;height:8px;border-radius:50%;background:#c0392b;margin-right:5px;'></span>" if is_sel else ""
        js = "var t=document.querySelector('#agent_hidden_input textarea');t.value='AID';t.dispatchEvent(new Event('input',{bubbles:true}));".replace("AID", str(aid))
        parts.append(
            "<div onclick=\"" + js + "\" style=\"cursor:pointer;border-radius:10px;padding:10px 13px;transition:all 0.2s;" + style + "\">"
            "<div style='font-size:11px;font-weight:700;color:#c0392b;margin-bottom:3px;font-family:monospace;'>" + dot + "Agent #" + str(aid) + "</div>"
            "<div style='font-size:11px;color:#333;line-height:1.6;'>"
            "<b>Age:</b> " + str(age) + " &nbsp; <b>Sex:</b> " + sex + "<br>"
            "<b>Edu:</b> " + edu + "<br>"
            "<b>Income:</b> " + inc + "</div></div>"
        )
    parts.append("</div>")
    return "".join(parts)


# ── UI ────────────────────────────────────────────────────────────────────────

with gr.Blocks(title="HiCoTraj Demo") as app:
    gr.Markdown("## HiCoTraj — Trajectory Visualization & Hierarchical CoT Demo")
    gr.Markdown("*Zero-Shot Demographic Reasoning via Hierarchical Chain-of-Thought Prompting from Trajectory*  ·  ACM SIGSPATIAL GeoGenAgent 2025")
    gr.Markdown("""
**Dataset:** NUMOSIM — a synthetic mobility dataset with realistic activity patterns across 6,000 agents.
> Stanford C, Adari S, Liao X, et al. *NUMoSim: A Synthetic Mobility Dataset with Anomaly Detection Benchmarks.* ACM SIGSPATIAL Workshop on Geospatial Anomaly Detection, 2024.
""")

    gr.Markdown("### Select Agent")
    agent_cards = gr.HTML(value=build_agent_cards(SHOWCASE_AGENTS[0]))
    agent_hidden = gr.Textbox(
        value=str(SHOWCASE_AGENTS[0]),
        visible=True,
        elem_id="agent_hidden_input",
        elem_classes=["hidden-input"]
    )
    gr.HTML("<style>.hidden-input { display:none !important; }</style>")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Trajectory Map")
            map_out = gr.HTML()
            gr.Markdown("### NUMOSIM Raw Data")
            with gr.Tabs():
                with gr.Tab("Summary"):
                    summary_out = gr.Textbox(lines=10, interactive=False, label="", show_label=False)
                with gr.Tab("Raw Data"):
                    raw_out = gr.Textbox(lines=10, interactive=False, label="", show_label=False)
                    show_all_btn = gr.Button("Show All Days", size="sm", variant="secondary")

        with gr.Column(scale=1):
            gr.Markdown("### Hierarchical Chain-of-Thought Reasoning")
            step_state = gr.State(value=0)
            run_btn    = gr.Button("▶ Stage 1: Feature Extraction", variant="primary")
            chain_out  = gr.HTML(value=render_chain("", "", "", status="idle"))

    def on_agent_click(agent_id):
        cards_html = build_agent_cards(agent_id)
        map_html, summary, raw_text, _demo_text, chain_html, step, btn = on_select_reset(agent_id)
        return cards_html, map_html, summary, raw_text, chain_html, step, btn

    agent_hidden.change(
        fn=on_agent_click, inputs=agent_hidden,
        outputs=[agent_cards, map_out, summary_out, raw_out, chain_out, step_state, run_btn]
    )
    def on_load(agent_id):
        map_html, summary, raw_text, _demo_text, chain_html, step, btn = on_select_reset(agent_id)
        return map_html, summary, raw_text, chain_html, step, btn

    app.load(
        fn=on_load, inputs=agent_hidden,
        outputs=[map_out, summary_out, raw_out, chain_out, step_state, run_btn]
    )
    run_btn.click(
        fn=handle_btn, inputs=[agent_hidden, step_state],
        outputs=[chain_out, step_state, run_btn]
    )

    def toggle_raw(agent_id, current_text):
        agent_id_int = int(agent_id)
        cot_entry = cot_by_agent.get(agent_id_int, {})
        agent_sp  = sp[sp["agent_id"] == agent_id_int].sort_values("start_datetime")
        raw_full  = cot_entry.get("weekly_checkin") or build_weekly_checkin(agent_sp)
        if "more days" in current_text:
            return raw_full, gr.update(value="Show Less")
        else:
            sep = "\n\n--- "
            parts = raw_full.split(sep)
            extra = len(parts) - 1
            short = parts[0] + (sep.join([""] + parts[1:2]) + ("\n\n... ({} more days)".format(extra - 1) if extra > 1 else "")) if extra > 0 else raw_full
            return short, gr.update(value="Show All Days")

    show_all_btn.click(
        fn=toggle_raw, inputs=[agent_hidden, raw_out],
        outputs=[raw_out, show_all_btn]
    )

app.launch(show_error=True, theme=gr.themes.Soft())