# app.py
# -*- coding: utf-8 -*-
"""
EarningALZ Information Flow Observatory

A dark sci-fi Streamlit dashboard deployable on Hugging Face Spaces.
It loads Parquet data directly from Hugging Face datasets and visualizes
company-level, cluster-level, and signal-level information propagation.
"""

from __future__ import annotations

import math
import re
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

import numpy as np
import pandas as pd
import networkx as nx
import plotly.graph_objects as go
import streamlit as st
from huggingface_hub import hf_hub_download, list_repo_files


# ============================================================
# Page config and dark sci-fi CSS
# ============================================================

st.set_page_config(
    page_title="EarningALZ Information Flow Observatory",
    page_icon="🛰️",
    layout="wide",
    initial_sidebar_state="expanded",
)

st.markdown(
    """
<style>
:root {
    --bg: #050A12;
    --panel: rgba(12, 20, 36, 0.92);
    --cyan: #00E5FF;
    --blue: #4CC9F0;
    --violet: #9D4EDD;
    --green: #00F5A0;
    --amber: #FFB703;
    --red: #FF5A5F;
    --text: #E5F3FF;
    --muted: #9FB3C8;
}
.stApp {
    background:
        radial-gradient(circle at top left, rgba(0, 229, 255, 0.11), transparent 28%),
        radial-gradient(circle at top right, rgba(157, 78, 221, 0.12), transparent 25%),
        linear-gradient(135deg, #050A12 0%, #08111F 45%, #030712 100%);
    color: var(--text);
}
.block-container { padding-top: 1.2rem; padding-bottom: 2rem; }
[data-testid="stSidebar"] {
    background: linear-gradient(180deg, rgba(8, 17, 31, 0.98), rgba(3, 7, 18, 0.98));
    border-right: 1px solid rgba(0, 229, 255, 0.22);
}
h1, h2, h3 { color: #E5F3FF !important; letter-spacing: 0.02em; }
h1 { text-shadow: 0 0 18px rgba(0, 229, 255, 0.45); }
div[data-testid="stMetric"] {
    background: linear-gradient(180deg, rgba(17, 24, 39, 0.92), rgba(5, 10, 18, 0.90));
    border: 1px solid rgba(0, 229, 255, 0.22);
    border-radius: 16px;
    padding: 14px 16px;
    box-shadow: 0 0 24px rgba(0, 229, 255, 0.08);
}
div[data-testid="stMetricLabel"] { color: #9FB3C8 !important; }
div[data-testid="stMetricValue"] {
    color: #00E5FF !important;
    text-shadow: 0 0 10px rgba(0, 229, 255, 0.38);
}
.dashboard-card {
    background: linear-gradient(180deg, rgba(17,24,39,0.92), rgba(5,10,18,0.88));
    border: 1px solid rgba(0,229,255,0.20);
    border-radius: 18px;
    padding: 16px;
    box-shadow: 0 0 28px rgba(0,229,255,0.08);
}
.small-muted { color: #9FB3C8; font-size: 0.88rem; }
.neon-line {
    height: 1px;
    background: linear-gradient(90deg, transparent, rgba(0,229,255,0.7), transparent);
    margin: 0.65rem 0 1rem 0;
}
</style>
""",
    unsafe_allow_html=True,
)


# ============================================================
# Constants
# ============================================================

DEFAULT_TWO_PART_DATASET = "soysouce/earningALZ_twopart"
DEFAULT_EVIDENCE_DATASET = "soysouce/earningALZ_SBERT_evidence"

TWO_PART_FILES = {
    "outlook": "cleaned_outlook_all.parquet",
    "relationships": "matched_company_relationships.parquet",
    "cross_events": "cross_quarter_events.parquet",
    "same_events": "same_quarter_events.parquet",
    "clusters": "best_company_cluster_assignment.parquet",
}

EVIDENCE_FILES = {
    "chunks": "rag_evidence_chunks_flat_full_gpu_direct.parquet",
}

SIGNAL_COLORS = {
    "demand_outlook": "#00E5FF",
    "margin_outlook": "#00F5A0",
    "supply_outlook": "#FFB703",
    "inventory_outlook": "#9D4EDD",
    "pricing_outlook": "#FF4D9D",
    "capex_outlook": "#4CC9F0",
}

DIRECTION_COLORS = {
    "positive": "#00F5A0",
    "negative": "#FF5A5F",
    "mixed": "#FFB703",
    "neutral": "#9FB3C8",
    "not_mentioned": "#334155",
    "": "#334155",
}

RELATION_COLORS = {
    "upstream": "#FFB703",
    "downstream": "#00E5FF",
    "partner": "#9D4EDD",
    "parent": "#4CC9F0",
    "subsidiary": "#00F5A0",
    "competitor": "#FF4D9D",
    "related": "#94A3B8",
    "customer": "#38BDF8",
    "customer_group": "#38BDF8",
    "supplier_group": "#F59E0B",
    "unknown": "#64748B",
}


# ============================================================
# Utility functions
# ============================================================

def clean_node(value) -> str:
    if pd.isna(value):
        return ""
    text = str(value).strip()
    if not text or text.lower() in {"nan", "none", "null", "0"}:
        return ""
    return text


def quarter_to_index(quarter: str) -> float:
    match = re.match(r"^(\d{4})Q([1-4])$", str(quarter).strip())
    if not match:
        return np.nan
    return int(match.group(1)) * 4 + int(match.group(2))


def sort_quarters(quarters: Iterable[str]) -> List[str]:
    return sorted([q for q in quarters if not pd.isna(quarter_to_index(q))], key=quarter_to_index)


def normalize_bool(series: pd.Series) -> pd.Series:
    if series.dtype == bool:
        return series.fillna(False)
    return series.astype(str).str.lower().isin(["true", "1", "yes", "y"])


def safe_float(value, default=0.0) -> float:
    try:
        if pd.isna(value):
            return default
        return float(value)
    except Exception:
        return default


def fmt_rate(value) -> str:
    if pd.isna(value):
        return "N/A"
    return f"{100 * float(value):.1f}%"


def prefixed(prefix: str, filename: str) -> str:
    prefix = prefix.strip().strip("/")
    return f"{prefix}/{filename}" if prefix else filename


# ============================================================
# Hugging Face loaders
# ============================================================

@st.cache_data(show_spinner=False, ttl=3600)
def list_hf_parquets(repo_id: str, revision: str = "main") -> List[str]:
    files = list_repo_files(repo_id=repo_id, repo_type="dataset", revision=revision)
    return sorted([f for f in files if f.endswith(".parquet")])


def select_hf_file(repo_id: str, filename: str, prefix: str = "", revision: str = "main") -> str:
    files = list_hf_parquets(repo_id, revision)
    stem = Path(filename).stem
    candidates = [
        prefixed(prefix, filename),
        filename,
        prefixed(prefix, f"data/{filename}"),
        prefixed(prefix, f"results/{filename}"),
        prefixed(prefix, f"two_part_network_prediction_analysis_hf/{filename}"),
        prefixed(prefix, f"two_part_network_prediction_analysis_parquet/{filename}"),
        prefixed(prefix, f"cluster_method_comparison_v4_hf/{filename}"),
    ]
    candidates += [f for f in files if Path(f).name == filename]
    candidates += [f for f in files if stem in Path(f).stem]

    seen = set()
    for candidate in candidates:
        if candidate in seen:
            continue
        seen.add(candidate)
        if candidate in files:
            return candidate

    raise FileNotFoundError(f"Could not find {filename} in {repo_id}. Available parquet files: {files[:80]}")


@st.cache_data(show_spinner=False, ttl=3600)
def read_hf_parquet(repo_id: str, filename: str, prefix: str = "", revision: str = "main") -> pd.DataFrame:
    remote_file = select_hf_file(repo_id, filename, prefix, revision)
    local_path = hf_hub_download(repo_id=repo_id, filename=remote_file, repo_type="dataset", revision=revision)
    df = pd.read_parquet(local_path)
    df["_hf_dataset"] = repo_id
    df["_hf_file"] = remote_file
    return df


@st.cache_data(show_spinner=True, ttl=3600)
def load_dashboard_data(two_part_dataset: str, evidence_dataset: str, prefix: str, revision: str) -> Dict[str, pd.DataFrame]:
    data = {}
    data["outlook"] = read_hf_parquet(two_part_dataset, TWO_PART_FILES["outlook"], prefix, revision)
    data["relationships"] = read_hf_parquet(two_part_dataset, TWO_PART_FILES["relationships"], prefix, revision)
    data["cross_events"] = read_hf_parquet(two_part_dataset, TWO_PART_FILES["cross_events"], prefix, revision)
    data["same_events"] = read_hf_parquet(two_part_dataset, TWO_PART_FILES["same_events"], prefix, revision)

    try:
        data["clusters"] = read_hf_parquet(two_part_dataset, TWO_PART_FILES["clusters"], prefix, revision)
    except Exception:
        data["clusters"] = pd.DataFrame()

    try:
        data["evidence_chunks"] = read_hf_parquet(evidence_dataset, EVIDENCE_FILES["chunks"], "", revision)
    except Exception:
        data["evidence_chunks"] = pd.DataFrame()

    return data


# ============================================================
# Data preparation
# ============================================================

@st.cache_data(show_spinner=False)
def prepare_events(cross_events: pd.DataFrame, same_events: pd.DataFrame) -> pd.DataFrame:
    cross = cross_events.copy()
    same = same_events.copy()
    if "analysis_mode" not in cross.columns:
        cross["analysis_mode"] = "cross_quarter"
    if "analysis_mode" not in same.columns:
        same["analysis_mode"] = "same_quarter"

    events = pd.concat([cross, same], ignore_index=True, sort=False)

    for col in ["source_active", "target_active", "direction_match", "exact_match"]:
        events[col] = normalize_bool(events[col]) if col in events.columns else False

    for col in ["source_node", "target_node", "source_quarter", "target_quarter", "signal", "relation_group"]:
        events[col] = events[col].astype(str).str.strip() if col in events.columns else ""

    for col in ["source_direction", "target_direction", "source_label", "target_label"]:
        events[col] = events[col].astype(str).str.strip() if col in events.columns else ""

    events["success"] = events["source_active"] & events["target_active"] & events["direction_match"]
    events["window"] = events["source_quarter"] + "→" + events["target_quarter"]
    events["release_date_gap_days"] = pd.to_numeric(events.get("release_date_gap_days", np.nan), errors="coerce")
    return events


@st.cache_data(show_spinner=False)
def prepare_relationships(relationships: pd.DataFrame) -> pd.DataFrame:
    rel = relationships.copy()
    if "source_company_node" not in rel.columns and "source_node" in rel.columns:
        rel["source_company_node"] = rel["source_node"]
    if "target_company_node" not in rel.columns and "target_node" in rel.columns:
        rel["target_company_node"] = rel["target_node"]

    rel["source_company_node"] = rel["source_company_node"].map(clean_node)
    rel["target_company_node"] = rel["target_company_node"].map(clean_node)
    if "relation_group_clean" in rel.columns:
        rel["relation_group_view"] = rel["relation_group_clean"].astype(str).str.strip().str.lower()
    elif "relation_group" in rel.columns:
        rel["relation_group_view"] = rel["relation_group"].astype(str).str.strip().str.lower()
    else:
        rel["relation_group_view"] = "unknown"

    rel = rel[
        rel["source_company_node"].ne("")
        & rel["target_company_node"].ne("")
        & rel["source_company_node"].ne(rel["target_company_node"])
    ].copy()
    return rel


@st.cache_data(show_spinner=False)
def build_company_table(outlook: pd.DataFrame, events: pd.DataFrame) -> pd.DataFrame:
    frames = []
    if {"company_node", "ticker", "current_company"}.issubset(outlook.columns):
        frames.append(outlook[["company_node", "ticker", "current_company"]].rename(columns={"current_company": "company"}))

    for side in ["source", "target"]:
        node_col = f"{side}_node"
        ticker_col = f"{side}_ticker"
        company_col = f"{side}_company"
        if node_col in events.columns:
            cols = [node_col]
            if ticker_col in events.columns:
                cols.append(ticker_col)
            if company_col in events.columns:
                cols.append(company_col)
            x = events[cols].rename(columns={node_col: "company_node", ticker_col: "ticker", company_col: "company"})
            frames.append(x)

    if not frames:
        nodes = sorted(set(events["source_node"]) | set(events["target_node"]))
        return pd.DataFrame({"company_node": nodes, "ticker": nodes, "company": nodes, "display_name": nodes})

    company = pd.concat(frames, ignore_index=True, sort=False).fillna("")
    if "ticker" not in company.columns:
        company["ticker"] = ""
    if "company" not in company.columns:
        company["company"] = ""
    company["company_node"] = company["company_node"].map(clean_node)
    company = company[company["company_node"].ne("")].copy()

    out = company.groupby("company_node", as_index=False).agg(
        ticker=("ticker", lambda x: next((str(v) for v in x if clean_node(v)), "")),
        company=("company", lambda x: next((str(v) for v in x if clean_node(v)), "")),
    )
    out["display_name"] = np.where(out["ticker"].astype(str).str.len() > 0, out["ticker"], out["company_node"].str.replace("COMPANY::", "", regex=False))
    return out


@st.cache_data(show_spinner=False)
def build_clusters(relationships: pd.DataFrame, clusters: pd.DataFrame, company_table: pd.DataFrame) -> pd.DataFrame:
    if not clusters.empty and "company_node" in clusters.columns:
        out = clusters.copy()
        if "cluster_id" not in out.columns:
            id_cols = [c for c in out.columns if "cluster" in c.lower() and "id" in c.lower()]
            out["cluster_id"] = out[id_cols[0]] if id_cols else 0
        if "cluster_theme_label" not in out.columns:
            label_cols = [c for c in out.columns if "theme" in c.lower() or "label" in c.lower()]
            out["cluster_theme_label"] = out[label_cols[0]] if label_cols else "Network community " + out["cluster_id"].astype(str)
        return out[["company_node", "cluster_id", "cluster_theme_label"]].drop_duplicates()

    graph = nx.Graph()
    for _, row in relationships.iterrows():
        s, t = row["source_company_node"], row["target_company_node"]
        if s and t and s != t:
            graph.add_edge(s, t, weight=graph.get_edge_data(s, t, default={}).get("weight", 0) + 1)

    if graph.number_of_nodes() == 0:
        return pd.DataFrame({"company_node": company_table["company_node"], "cluster_id": 0, "cluster_theme_label": "All companies"})

    communities = list(nx.algorithms.community.greedy_modularity_communities(graph, weight="weight"))
    rows = []
    for cid, nodes in enumerate(communities):
        for node in nodes:
            rows.append({"company_node": node, "cluster_id": cid, "cluster_theme_label": f"Network community {cid}"})
    out = pd.DataFrame(rows)
    missing = sorted(set(company_table["company_node"]) - set(out["company_node"]))
    if missing:
        out = pd.concat([out, pd.DataFrame({"company_node": missing, "cluster_id": -1, "cluster_theme_label": "Unassigned"})], ignore_index=True)
    return out


@st.cache_data(show_spinner=False)
def build_leader_follower(events: pd.DataFrame) -> pd.DataFrame:
    active = events[events["source_active"]].copy()
    if active.empty:
        return pd.DataFrame()

    source = active.groupby("source_node", as_index=False).agg(
        outgoing_exposures=("signal", "count"),
        outgoing_successes=("success", "sum"),
        outgoing_direction_match_rate=("direction_match", "mean"),
        distinct_targets=("target_node", "nunique"),
        avg_gap_days=("release_date_gap_days", "mean"),
    ).rename(columns={"source_node": "company_node"})

    target = active.groupby("target_node", as_index=False).agg(
        incoming_exposures=("signal", "count"),
        incoming_successes=("success", "sum"),
        incoming_direction_match_rate=("direction_match", "mean"),
        distinct_sources=("source_node", "nunique"),
    ).rename(columns={"target_node": "company_node"})

    out = source.merge(target, on="company_node", how="outer").fillna(0)
    out["leader_score"] = np.log1p(out["outgoing_exposures"]) * out["outgoing_direction_match_rate"]
    out["follower_score"] = np.log1p(out["incoming_exposures"]) * out["incoming_direction_match_rate"]
    out["leader_minus_follower"] = out["leader_score"] - out["follower_score"]
    return out.sort_values("leader_score", ascending=False)


# ============================================================
# Filtering
# ============================================================

def filter_events(events: pd.DataFrame, mode: str, signal: str, directions: List[str], relations: List[str], quarters: List[str], only_successful: bool) -> pd.DataFrame:
    data = events.copy()
    if mode != "Both":
        data = data[data["analysis_mode"].eq(mode)]
    if signal != "All":
        data = data[data["signal"].eq(signal)]
    if directions and "All" not in directions:
        data = data[data["source_direction"].isin(directions)]
    if relations and "All" not in relations:
        data = data[data["relation_group"].isin(relations)]
    if quarters:
        qset = set(quarters)
        data = data[data["source_quarter"].isin(qset) | data["target_quarter"].isin(qset)]
    if only_successful:
        data = data[data["success"]]
    return data.copy()


def ego_nodes_from_events(events: pd.DataFrame, center_node: str, depth: int, max_nodes: int = 250) -> set[str]:
    graph = nx.Graph()
    for _, row in events.iterrows():
        s = clean_node(row.get("source_node", ""))
        t = clean_node(row.get("target_node", ""))
        if s and t and s != t:
            graph.add_edge(s, t)
    if center_node not in graph:
        return {center_node}
    distances = nx.single_source_shortest_path_length(graph, center_node, cutoff=depth)
    ordered = sorted(distances.keys(), key=lambda n: (distances[n], n))
    return set(ordered[:max_nodes])


def nodes_from_cluster(cluster_df: pd.DataFrame, cluster_label: str, max_nodes: int = 350) -> set[str]:
    if cluster_label == "All":
        return set(cluster_df["company_node"].head(max_nodes))
    data = cluster_df[cluster_df["cluster_theme_label"].eq(cluster_label)]
    return set(data["company_node"].head(max_nodes))


# ============================================================
# Network construction and plotting
# ============================================================

def build_edge_summary(events: pd.DataFrame) -> pd.DataFrame:
    if events.empty:
        return pd.DataFrame()
    group_cols = ["source_node", "target_node", "signal", "relation_group", "source_direction"]
    summary = events.groupby(group_cols, dropna=False).agg(
        event_count=("signal", "count"),
        success_count=("success", "sum"),
        target_active_rate=("target_active", "mean"),
        direction_match_rate=("direction_match", "mean"),
        avg_gap_days=("release_date_gap_days", "mean"),
        first_source_quarter=("source_quarter", "min"),
        last_target_quarter=("target_quarter", "max"),
    ).reset_index()
    summary["success_rate"] = summary["success_count"] / summary["event_count"].replace(0, np.nan)
    return summary.sort_values(["success_count", "event_count"], ascending=[False, False])


@st.cache_data(show_spinner=False)
def compute_layout(edge_summary_records: list, seed: int = 42) -> Dict[str, Tuple[float, float]]:
    edge_summary = pd.DataFrame(edge_summary_records)
    graph = nx.Graph()
    for _, row in edge_summary.iterrows():
        graph.add_edge(row["source_node"], row["target_node"], weight=max(1.0, safe_float(row.get("event_count", 1))))
    if graph.number_of_nodes() == 0:
        return {}
    iterations = 80 if graph.number_of_nodes() <= 600 else 35
    pos = nx.spring_layout(graph, seed=seed, iterations=iterations, weight="weight")
    return {node: (float(x), float(y)) for node, (x, y) in pos.items()}


def build_network_figure(edge_summary: pd.DataFrame, company_table: pd.DataFrame, cluster_table: pd.DataFrame, leader_table: pd.DataFrame, animate: bool, selected_signal: str, max_edges: int = 600) -> go.Figure:
    if edge_summary.empty:
        fig = go.Figure()
        fig.update_layout(template="plotly_dark", paper_bgcolor="#050A12", plot_bgcolor="#050A12", title="No network events match the current filters")
        return fig

    edge_summary = edge_summary.head(max_edges).copy()
    layout = compute_layout(edge_summary.to_dict(orient="records"))

    nodes = sorted(set(edge_summary["source_node"]) | set(edge_summary["target_node"]))
    company_map = company_table.set_index("company_node").to_dict(orient="index") if not company_table.empty else {}
    cluster_map = cluster_table.set_index("company_node").to_dict(orient="index") if not cluster_table.empty else {}
    leader_map = leader_table.set_index("company_node").to_dict(orient="index") if not leader_table.empty else {}

    traces = []
    for relation, group in edge_summary.groupby("relation_group", dropna=False):
        x_lines, y_lines = [], []
        for _, row in group.iterrows():
            s, t = row["source_node"], row["target_node"]
            if s not in layout or t not in layout:
                continue
            x0, y0 = layout[s]
            x1, y1 = layout[t]
            x_lines += [x0, x1, None]
            y_lines += [y0, y1, None]
        traces.append(go.Scatter(x=x_lines, y=y_lines, mode="lines", line=dict(width=1.2, color=RELATION_COLORS.get(str(relation), "#64748B")), opacity=0.36, hoverinfo="skip", name=f"edge: {relation}"))

    arrow_x, arrow_y, arrow_text, arrow_color, arrow_size = [], [], [], [], []
    for _, row in edge_summary.iterrows():
        s, t = row["source_node"], row["target_node"]
        if s not in layout or t not in layout:
            continue
        x0, y0 = layout[s]
        x1, y1 = layout[t]
        arrow_x.append(0.68 * x0 + 0.32 * x1)
        arrow_y.append(0.68 * y0 + 0.32 * y1)
        signal = row.get("signal", selected_signal)
        direction = row.get("source_direction", "")
        arrow_color.append(SIGNAL_COLORS.get(signal, DIRECTION_COLORS.get(direction, "#00E5FF")))
        arrow_size.append(8 + 10 * min(1.0, safe_float(row.get("success_rate", 0))))
        arrow_text.append(
            f"<b>{s}</b> → <b>{t}</b><br>signal: {signal}<br>direction: {direction}<br>relation: {row.get('relation_group', '')}<br>events: {int(row.get('event_count', 0))}<br>success rate: {fmt_rate(row.get('success_rate', np.nan))}<br>avg gap days: {safe_float(row.get('avg_gap_days', np.nan), np.nan):.1f}"
        )

    traces.append(go.Scatter(x=arrow_x, y=arrow_y, mode="markers", marker=dict(size=arrow_size, color=arrow_color, symbol="triangle-right", line=dict(width=0.6, color="#E5F3FF"), opacity=0.88), text=arrow_text, hoverinfo="text", name="flow direction"))

    node_x, node_y, node_text, node_size, node_color = [], [], [], [], []
    palette = ["#00E5FF", "#9D4EDD", "#00F5A0", "#FFB703", "#FF4D9D", "#4CC9F0", "#F97316", "#A3E635"]
    for node in nodes:
        if node not in layout:
            continue
        x, y = layout[node]
        node_x.append(x)
        node_y.append(y)
        company_info = company_map.get(node, {})
        cluster_info = cluster_map.get(node, {})
        leader_info = leader_map.get(node, {})
        display_name = company_info.get("display_name", node.replace("COMPANY::", ""))
        cluster_label = cluster_info.get("cluster_theme_label", "Unknown community")
        visible_degree = int((edge_summary["source_node"].eq(node)).sum() + (edge_summary["target_node"].eq(node)).sum())
        leader_score = safe_float(leader_info.get("leader_score", 0))
        follower_score = safe_float(leader_info.get("follower_score", 0))
        node_text.append(f"<b>{display_name}</b><br>node: {node}<br>cluster: {cluster_label}<br>visible degree: {visible_degree}<br>leader score: {leader_score:.3f}<br>follower score: {follower_score:.3f}")
        node_size.append(12 + 4 * math.log1p(visible_degree) + 2 * math.log1p(max(leader_score, follower_score)))
        cid = int(safe_float(cluster_info.get("cluster_id", 0), 0))
        node_color.append(palette[cid % len(palette)])

    traces.append(go.Scatter(x=node_x, y=node_y, mode="markers", marker=dict(size=node_size, color=node_color, opacity=0.94, line=dict(width=1.4, color="#E5F3FF")), text=node_text, hoverinfo="text", name="companies"))

    frames = []
    if animate:
        animated_edges = edge_summary.head(min(120, len(edge_summary))).copy()
        steps = 22
        for step in range(steps):
            tval = step / max(1, steps - 1)
            px, py, ptext, pcolor, psize = [], [], [], [], []
            for _, row in animated_edges.iterrows():
                s, t = row["source_node"], row["target_node"]
                if s not in layout or t not in layout:
                    continue
                x0, y0 = layout[s]
                x1, y1 = layout[t]
                tt = (tval + 0.07 * (hash(s + t) % 10)) % 1.0
                px.append((1 - tt) * x0 + tt * x1)
                py.append((1 - tt) * y0 + tt * y1)
                pcolor.append(SIGNAL_COLORS.get(row.get("signal", selected_signal), "#00E5FF"))
                psize.append(8 + 12 * min(1.0, safe_float(row.get("success_rate", 0))))
                ptext.append(f"Pulse: {s} → {t}<br>signal: {row.get('signal', '')}<br>relation: {row.get('relation_group', '')}<br>window: {row.get('first_source_quarter', '')} → {row.get('last_target_quarter', '')}")
            frames.append(go.Frame(data=[go.Scatter(x=px, y=py, mode="markers", marker=dict(size=psize, color=pcolor, opacity=0.96, symbol="circle", line=dict(width=1.0, color="#FFFFFF")), text=ptext, hoverinfo="text", name="diffusion pulse")], name=f"frame_{step}"))
        if frames:
            traces.append(frames[0].data[0])

    updatemenus = []
    if animate and frames:
        updatemenus = [dict(type="buttons", showactive=False, x=0.02, y=1.06, xanchor="left", yanchor="top", buttons=[
            dict(label="▶ Play diffusion", method="animate", args=[None, {"frame": {"duration": 150, "redraw": True}, "fromcurrent": True, "transition": {"duration": 40}}]),
            dict(label="⏸ Pause", method="animate", args=[[None], {"frame": {"duration": 0, "redraw": False}, "mode": "immediate", "transition": {"duration": 0}}]),
        ])]

    fig = go.Figure(data=traces, frames=frames)
    fig.update_layout(template="plotly_dark", paper_bgcolor="#050A12", plot_bgcolor="#050A12", height=760, margin=dict(l=10, r=10, t=50, b=10), title=dict(text="Information Flow Network", font=dict(size=22, color="#E5F3FF")), showlegend=True, legend=dict(orientation="h", yanchor="bottom", y=1.01, xanchor="right", x=1, bgcolor="rgba(5,10,18,0.35)", font=dict(color="#CFE9FF", size=10)), xaxis=dict(visible=False), yaxis=dict(visible=False), updatemenus=updatemenus)
    return fig


# ============================================================
# Charts
# ============================================================

def make_timeline_chart(events: pd.DataFrame) -> go.Figure:
    active = events[events["source_active"]].copy()
    if active.empty:
        return go.Figure()
    summary = active.groupby(["source_quarter", "signal"], as_index=False).agg(event_count=("signal", "count"), success_rate=("success", "mean"))
    summary["quarter_index"] = summary["source_quarter"].map(quarter_to_index)
    summary = summary.sort_values("quarter_index")
    fig = go.Figure()
    for signal, group in summary.groupby("signal"):
        fig.add_trace(go.Scatter(x=group["source_quarter"], y=group["success_rate"], mode="lines+markers", name=signal, line=dict(color=SIGNAL_COLORS.get(signal, "#00E5FF"), width=2.5), marker=dict(size=np.clip(np.log1p(group["event_count"]) * 2.5, 6, 20)), customdata=group["event_count"], hovertemplate="quarter=%{x}<br>success rate=%{y:.3f}<br>events=%{customdata}<extra></extra>"))
    fig.update_layout(template="plotly_dark", paper_bgcolor="#050A12", plot_bgcolor="#050A12", height=360, title="Signal Propagation Timeline", yaxis_title="success rate", xaxis_title="source quarter", legend=dict(orientation="h", y=1.08))
    return fig


def make_relation_bar(events: pd.DataFrame) -> go.Figure:
    active = events[events["source_active"]].copy()
    if active.empty:
        return go.Figure()
    summary = active.groupby("relation_group", as_index=False).agg(event_count=("signal", "count"), success_rate=("success", "mean")).sort_values("success_rate", ascending=True).tail(20)
    fig = go.Figure(go.Bar(x=summary["success_rate"], y=summary["relation_group"], orientation="h", marker=dict(color="#00E5FF", line=dict(color="#E5F3FF", width=0.5)), text=summary["event_count"], hovertemplate="relation=%{y}<br>success rate=%{x:.3f}<br>events=%{text}<extra></extra>"))
    fig.update_layout(template="plotly_dark", paper_bgcolor="#050A12", plot_bgcolor="#050A12", height=360, title="Propagation Strength by Relationship Type", xaxis_title="success rate", yaxis_title="")
    return fig


def find_evidence_rows(evidence_chunks: pd.DataFrame, chunk_ids: str, max_rows: int = 12) -> pd.DataFrame:
    if evidence_chunks.empty or not chunk_ids:
        return pd.DataFrame()
    ids = [x.strip() for x in re.split(r"[|,;]", str(chunk_ids)) if x.strip()]
    if not ids:
        return pd.DataFrame()
    possible_cols = [c for c in ["chunk_id", "chunk_uid", "evidence_id"] if c in evidence_chunks.columns]
    if not possible_cols:
        return pd.DataFrame()
    mask = pd.Series(False, index=evidence_chunks.index)
    for col in possible_cols:
        mask |= evidence_chunks[col].astype(str).isin(ids)
    cols = [c for c in ["doc_id", "ticker", "current_company", "quarter", "chunk_id", "chunk_uid", "query_group", "rank", "hybrid_score", "text", "chunk_text"] if c in evidence_chunks.columns]
    return evidence_chunks.loc[mask, cols].head(max_rows).copy()


# ============================================================
# App layout
# ============================================================

st.markdown("""
# 🛰️ EarningALZ Information Flow Observatory
<div class="small-muted">Dark sci-fi dashboard for exploring earnings-call signal propagation across company networks.</div>
<div class="neon-line"></div>
""", unsafe_allow_html=True)

with st.sidebar:
    st.header("Data Source")
    two_part_dataset = st.text_input("Two-part HF dataset", value=DEFAULT_TWO_PART_DATASET)
    evidence_dataset = st.text_input("Evidence HF dataset", value=DEFAULT_EVIDENCE_DATASET)
    two_part_prefix = st.text_input("Optional two-part prefix", value="")
    revision = st.text_input("HF revision", value="main")

    with st.spinner("Loading Hugging Face parquet data..."):
        raw = load_dashboard_data(two_part_dataset, evidence_dataset, two_part_prefix, revision)

    events = prepare_events(raw["cross_events"], raw["same_events"])
    relationships = prepare_relationships(raw["relationships"])
    company_table = build_company_table(raw["outlook"], events)
    cluster_table = build_clusters(relationships, raw["clusters"], company_table)
    leader_table = build_leader_follower(events)

    events = events.merge(cluster_table.rename(columns={"company_node": "source_node", "cluster_id": "source_cluster_id", "cluster_theme_label": "source_cluster_theme_label"}), on="source_node", how="left")
    events = events.merge(cluster_table.rename(columns={"company_node": "target_node", "cluster_id": "target_cluster_id", "cluster_theme_label": "target_cluster_theme_label"}), on="target_node", how="left")

    st.markdown("---")
    st.header("Filters")
    mode_options = ["Both"] + sorted(events["analysis_mode"].dropna().unique().tolist())
    selected_mode = st.selectbox("Propagation mode", mode_options, index=0)
    signal_options = ["All"] + sorted(events["signal"].dropna().unique().tolist())
    selected_signal = st.selectbox("Signal", signal_options, index=signal_options.index("demand_outlook") if "demand_outlook" in signal_options else 0)
    direction_options = ["All"] + sorted([x for x in events["source_direction"].dropna().unique().tolist() if x])
    selected_directions = st.multiselect("Source direction", direction_options, default=["All"])
    relation_options = ["All"] + sorted([x for x in events["relation_group"].dropna().unique().tolist() if x])
    selected_relations = st.multiselect("Relationship type", relation_options, default=["All"])
    quarter_options = sort_quarters(set(events["source_quarter"]) | set(events["target_quarter"]))
    if quarter_options:
        start_q, end_q = st.select_slider("Quarter range", options=quarter_options, value=(quarter_options[0], quarter_options[-1]))
        selected_quarters = [q for q in quarter_options if quarter_to_index(start_q) <= quarter_to_index(q) <= quarter_to_index(end_q)]
    else:
        selected_quarters = []
    only_successful = st.checkbox("Show only successful same-direction events", value=False)

    st.markdown("---")
    st.header("Focus")
    focus_mode = st.radio("Focus mode", ["Whole network", "Company", "Cluster"], horizontal=False)
    selected_company_node = None
    selected_cluster = "All"
    hop_depth = 1
    if focus_mode == "Company":
        company_options = company_table.sort_values("display_name")
        name_to_node = dict(zip(company_options["display_name"], company_options["company_node"]))
        company_name = st.selectbox("Company", list(name_to_node.keys()))
        selected_company_node = name_to_node[company_name]
        hop_depth = st.slider("Hop depth", 1, 3, 2)
    elif focus_mode == "Cluster":
        cluster_options = ["All"] + sorted(cluster_table["cluster_theme_label"].dropna().unique().tolist())
        selected_cluster = st.selectbox("Cluster / community", cluster_options)

    st.markdown("---")
    st.header("Display")
    max_edges = st.slider("Maximum visible edges", 50, 1200, 500, 50)
    animate_flow = st.checkbox("Enable animated diffusion pulses", value=True)
    show_evidence = st.checkbox("Enable evidence inspector", value=True)

base_filtered = filter_events(events, selected_mode, selected_signal, selected_directions, selected_relations, selected_quarters, only_successful)

if focus_mode == "Company" and selected_company_node:
    focus_nodes = ego_nodes_from_events(base_filtered, selected_company_node, hop_depth)
    filtered_events = base_filtered[base_filtered["source_node"].isin(focus_nodes) & base_filtered["target_node"].isin(focus_nodes)].copy()
elif focus_mode == "Cluster":
    focus_nodes = nodes_from_cluster(cluster_table, selected_cluster)
    filtered_events = base_filtered[base_filtered["source_node"].isin(focus_nodes) | base_filtered["target_node"].isin(focus_nodes)].copy()
else:
    filtered_events = base_filtered.copy()

edge_summary = build_edge_summary(filtered_events)
active_events = filtered_events[filtered_events["source_active"]].copy()
success_rate = float(active_events["success"].mean()) if not active_events.empty else np.nan
target_active_rate = float(active_events["target_active"].mean()) if not active_events.empty else np.nan
visible_nodes = set(edge_summary["source_node"]) | set(edge_summary["target_node"]) if not edge_summary.empty else set()

k1, k2, k3, k4, k5 = st.columns(5)
k1.metric("Visible companies", f"{len(visible_nodes):,}")
k2.metric("Visible edges", f"{len(edge_summary):,}")
k3.metric("Source-active events", f"{len(active_events):,}")
k4.metric("Target active rate", fmt_rate(target_active_rate))
k5.metric("Success rate", fmt_rate(success_rate))
st.markdown("<div class='neon-line'></div>", unsafe_allow_html=True)

tab_network, tab_timeline, tab_leaders, tab_events, tab_evidence = st.tabs(["🌐 Network Flow", "⏱️ Timeline", "🧭 Leaders", "📡 Events", "🔎 Evidence"])

with tab_network:
    left, right = st.columns([3.2, 1.0])
    with left:
        fig = build_network_figure(edge_summary, company_table, cluster_table, leader_table, animate_flow, selected_signal, max_edges=max_edges)
        st.plotly_chart(fig, use_container_width=True)
    with right:
        st.markdown("### Top visible edges")
        if not edge_summary.empty:
            cols = ["source_node", "target_node", "signal", "relation_group", "source_direction", "event_count", "success_rate", "avg_gap_days"]
            st.dataframe(edge_summary[[c for c in cols if c in edge_summary.columns]].head(25), use_container_width=True, height=610)
        else:
            st.info("No edges match the current filters.")

with tab_timeline:
    c1, c2 = st.columns([1.3, 1.0])
    with c1:
        st.plotly_chart(make_timeline_chart(filtered_events), use_container_width=True)
    with c2:
        st.plotly_chart(make_relation_bar(filtered_events), use_container_width=True)
    st.markdown("### Quarter-window summary")
    if not active_events.empty:
        summary = active_events.groupby(["window", "signal"], as_index=False).agg(event_count=("signal", "count"), success_rate=("success", "mean"), target_active_rate=("target_active", "mean"), direction_match_rate=("direction_match", "mean"))
        st.dataframe(summary, use_container_width=True, height=420)
    else:
        st.info("No active events match the current filters.")

with tab_leaders:
    merged = leader_table.merge(company_table, on="company_node", how="left").merge(cluster_table, on="company_node", how="left")
    if focus_mode == "Cluster" and selected_cluster != "All":
        merged = merged[merged["cluster_theme_label"].eq(selected_cluster)]
    elif focus_mode == "Company" and selected_company_node:
        focus_nodes = ego_nodes_from_events(base_filtered, selected_company_node, hop_depth)
        merged = merged[merged["company_node"].isin(focus_nodes)]
    l1, l2 = st.columns(2)
    with l1:
        st.markdown("### Top information leaders")
        cols = ["display_name", "company_node", "cluster_theme_label", "leader_score", "outgoing_exposures", "outgoing_successes", "outgoing_direction_match_rate", "distinct_targets"]
        st.dataframe(merged.sort_values("leader_score", ascending=False)[[c for c in cols if c in merged.columns]].head(30), use_container_width=True, height=560)
    with l2:
        st.markdown("### Top followers")
        cols = ["display_name", "company_node", "cluster_theme_label", "follower_score", "incoming_exposures", "incoming_successes", "incoming_direction_match_rate", "distinct_sources"]
        st.dataframe(merged.sort_values("follower_score", ascending=False)[[c for c in cols if c in merged.columns]].head(30), use_container_width=True, height=560)

with tab_events:
    st.markdown("### Filtered propagation events")
    event_cols = ["analysis_mode", "source_quarter", "target_quarter", "source_node", "target_node", "signal", "source_direction", "target_direction", "source_label", "target_label", "relation_group", "success", "release_date_gap_days", "source_doc_ids", "target_doc_ids", "relationship_doc_id", "source_signal_raw_values", "target_signal_raw_values", "source_signal_mapping_values", "target_signal_mapping_values"]
    st.dataframe(filtered_events[[c for c in event_cols if c in filtered_events.columns]].head(2000), use_container_width=True, height=650)

with tab_evidence:
    if not show_evidence:
        st.info("Evidence inspector is disabled in the sidebar.")
    elif filtered_events.empty:
        st.info("No events available for evidence inspection.")
    else:
        st.markdown("### Evidence inspector")
        selectable = filtered_events.copy().reset_index(drop=True)
        selectable["event_label"] = selectable.index.astype(str) + " | " + selectable["source_node"].astype(str).str.replace("COMPANY::", "", regex=False) + " → " + selectable["target_node"].astype(str).str.replace("COMPANY::", "", regex=False) + " | " + selectable["signal"].astype(str) + " | " + selectable["source_quarter"].astype(str) + "→" + selectable["target_quarter"].astype(str)
        selected_event_label = st.selectbox("Select event", selectable["event_label"].head(5000).tolist())
        selected_index = int(selected_event_label.split(" | ")[0])
        row = selectable.iloc[selected_index]
        c1, c2, c3 = st.columns(3)
        c1.metric("Source", row.get("source_node", ""))
        c2.metric("Target", row.get("target_node", ""))
        c3.metric("Signal", row.get("signal", ""))
        evidence_fields = ["source_doc_ids", "target_doc_ids", "relationship_doc_id", "source_outlook_row_ids", "target_outlook_row_ids", "relationship_row_id", "source_outlook_chunk_ids", "target_outlook_chunk_ids", "relationship_chunk_id", "source_signal_raw_values", "target_signal_raw_values", "source_signal_mapping_values", "target_signal_mapping_values"]
        evidence_fields = [c for c in evidence_fields if c in selectable.columns]
        st.markdown("#### Evidence-chain fields")
        st.json({field: str(row.get(field, "")) for field in evidence_fields})
        chunk_parts = []
        for field in ["source_outlook_chunk_ids", "target_outlook_chunk_ids", "relationship_chunk_id"]:
            if field in selectable.columns:
                chunk_parts.append(str(row.get(field, "")))
        evidence_rows = find_evidence_rows(raw["evidence_chunks"], "|".join(chunk_parts))
        if not evidence_rows.empty:
            st.markdown("#### Retrieved evidence chunks")
            st.dataframe(evidence_rows, use_container_width=True, height=520)
        else:
            st.info("No matching evidence chunk rows were found. This may happen if chunk IDs are not available in the selected event data.")

st.markdown("<div class='neon-line'></div>", unsafe_allow_html=True)
st.caption("This dashboard visualizes transcript-derived signal alignment and candidate information propagation, not confirmed causal transmission.")