app.py

"""VynFi Financial Sankey Explorer.

Interactive funds-flow viewer for the `financial_reporting/sankey/flows.json`
export produced by `datasynth-data` (`financial_reporting.sankey: true`). Pick an
entity, a period, and a statement and see the **income-statement waterfall**
(Revenue → Gross Profit → Operating Income → Net Income, with operating-expense
sub-bands) or the **cash-flow funds-flow** (Operating / Investing / Financing →
Net change in cash) rendered as a Plotly Sankey.

The flow JSON is the engine's native shape — a list of self-describing
`SankeyFlow` objects:

    {statement_type, fiscal_year, fiscal_period, company_code, currency,
     is_consolidated, period_start, period_end,
     nodes: [{id, label, section, value, is_subtotal, line_code}],
     links: [{source, target, value, label, is_contra}]}

so this Space renders real `datasynth-data` output verbatim. With no dataset
loaded it falls back to a small built-in demo set (constructed with the same
waterfall logic as the Rust builder) so the app is self-contained.
"""
from __future__ import annotations

import json
from decimal import Decimal
from pathlib import Path
from typing import Any

import gradio as gr
import plotly.graph_objects as go

# ── Section palette (matches the engine's statement sections) ────────────────
SECTION_COLOR = {
    "Revenue": "#2563eb",
    "Cost of Sales": "#dc2626",
    "Gross Profit": "#16a34a",
    "Operating Expenses": "#ea580c",
    "Operating Income": "#16a34a",
    "Tax": "#9333ea",
    "Net Income": "#15803d",
    "Operating": "#2563eb",
    "Investing": "#9333ea",
    "Financing": "#ea580c",
    "Net Change": "#15803d",
}
CONTRA_LINK = "rgba(220, 38, 38, 0.35)"   # red-ish: reductions / outflows
NORMAL_LINK = "rgba(37, 99, 235, 0.30)"   # blue-ish: retained / inflows


def _num(v: Any) -> float:
    """Coerce a Sankey value (serde_decimal serialises as a string by default,
    or as a JSON number in CLI mode) to float."""
    if isinstance(v, (int, float)):
        return float(v)
    try:
        return float(Decimal(str(v)))
    except Exception:
        return 0.0


# ── Demo flow builders (mirror the Rust sankey_builder waterfall) ────────────
def _node(nid, label, section, value, is_subtotal=False, code=None):
    # Magnitudes are absolute (matches the Rust builder; sign lives in is_contra).
    return {
        "id": nid, "label": label, "section": section,
        "value": str(abs(value)), "is_subtotal": is_subtotal, "line_code": code,
    }


def _link(src, tgt, value, label, contra=False):
    return {"source": src, "target": tgt, "value": str(abs(value)), "label": label, "is_contra": contra}


def demo_income_statement(company, fy, fp, period_start, period_end, revenue, cogs, opex_subs):
    """Build an IS waterfall flow exactly as the engine does."""
    gross = revenue - cogs
    opex_total = sum(v for _, v in opex_subs)
    oi = gross - opex_total
    tax = round(oi * 0.21)
    ni = oi - tax
    nodes, links = [], []
    rev_id = len(nodes); nodes.append(_node(rev_id, "Revenue", "Revenue", revenue, code="IS-REV"))
    gp_id = len(nodes); nodes.append(_node(gp_id, "Gross Profit", "Gross Profit", gross, True, "IS-GP"))
    if cogs > 0:
        c_id = len(nodes); nodes.append(_node(c_id, "Cost of Goods Sold", "Cost of Sales", cogs, code="IS-COGS"))
        links.append(_link(rev_id, c_id, cogs, "Cost of Goods Sold", True))
    links.append(_link(rev_id, gp_id, gross, "Gross Profit"))
    oi_id = len(nodes); nodes.append(_node(oi_id, "Operating Income", "Operating Income", oi, True, "IS-OI"))
    for label, amount in opex_subs:
        if amount <= 0:
            continue
        s_id = len(nodes); nodes.append(_node(s_id, label, "Operating Expenses", amount))
        links.append(_link(gp_id, s_id, amount, label, True))
    links.append(_link(gp_id, oi_id, oi, "Operating Income"))
    ni_id = len(nodes); nodes.append(_node(ni_id, "Net Income", "Net Income", ni, True, "IS-NI"))
    if tax > 0:
        t_id = len(nodes); nodes.append(_node(t_id, "Income Tax Expense", "Tax", tax, code="IS-TAX"))
        links.append(_link(oi_id, t_id, tax, "Income Tax Expense", True))
    links.append(_link(oi_id, ni_id, ni, "Net Income"))
    return {
        "statement_type": "income_statement", "fiscal_year": fy, "fiscal_period": fp,
        "company_code": company, "currency": "USD", "is_consolidated": False,
        "period_start": period_start, "period_end": period_end, "nodes": nodes, "links": links,
    }


def demo_cash_flow(company, fy, fp, period_start, period_end, items):
    """Build a CF funds-flow flow exactly as the engine does. `items` =
    list of (label, category, amount) with signed amounts."""
    secs = {"Operating": 0, "Investing": 0, "Financing": 0}
    for _, cat, amt in items:
        secs[cat] += amt
    net = sum(secs.values())
    nodes, links = [], []
    op_id = len(nodes); nodes.append(_node(op_id, "Operating Activities", "Operating", secs["Operating"], True))
    inv_id = len(nodes); nodes.append(_node(inv_id, "Investing Activities", "Investing", secs["Investing"], True))
    fin_id = len(nodes); nodes.append(_node(fin_id, "Financing Activities", "Financing", secs["Financing"], True))
    net_id = len(nodes); nodes.append(_node(net_id, "Net Change in Cash", "Net Change", net, True))
    sec_id = {"Operating": op_id, "Investing": inv_id, "Financing": fin_id}
    for label, cat, amt in items:
        if amt == 0:
            continue
        i_id = len(nodes); nodes.append(_node(i_id, label, cat, amt))
        links.append(_link(i_id, sec_id[cat], amt, label, amt < 0))
    links.append(_link(op_id, net_id, secs["Operating"], "Operating cash flow", secs["Operating"] < 0))
    links.append(_link(inv_id, net_id, secs["Investing"], "Investing cash flow", secs["Investing"] < 0))
    links.append(_link(fin_id, net_id, secs["Financing"], "Financing cash flow", secs["Financing"] < 0))
    return {
        "statement_type": "cash_flow", "fiscal_year": fy, "fiscal_period": fp,
        "company_code": company, "currency": "USD", "is_consolidated": False,
        "period_start": period_start, "period_end": period_end, "nodes": nodes, "links": links,
    }


def _demo_flows():
    flows = []
    opex = [("Salaries & Benefits", 180_000), ("Rent & Facilities", 60_000),
            ("Utilities", 25_000), ("Depreciation & Amortization", 40_000),
            ("Marketing & Advertising", 35_000), ("Other Operating Expenses", 18_000)]
    flows.append(demo_income_statement("1000", 2024, 3, "2024-03-01", "2024-03-31", 1_200_000, 720_000, opex))
    opex6 = [("Salaries & Benefits", 195_000), ("Rent & Facilities", 60_000),
             ("Utilities", 28_000), ("Depreciation & Amortization", 40_000),
             ("Marketing & Advertising", 52_000), ("Other Operating Expenses", 21_000)]
    flows.append(demo_income_statement("1000", 2024, 6, "2024-06-01", "2024-06-30", 1_410_000, 838_000, opex6))
    # Full-year rollup (period 0)
    opex_fy = [(l, v * 12) for l, v in opex]
    flows.append(demo_income_statement("1000", 2024, 0, "2024-01-01", "2024-12-31", 14_900_000, 8_940_000, opex_fy))
    flows.append(demo_cash_flow("1000", 2024, 3, "2024-03-01", "2024-03-31", [
        ("Net Income", "Operating", 110_600), ("Depreciation & Amortization", "Operating", 40_000),
        ("Change in Accounts Receivable", "Operating", -50_000), ("Change in Accounts Payable", "Operating", 20_000),
        ("Change in Inventory", "Operating", -15_000), ("Capital Expenditure", "Investing", -80_000),
        ("Debt Issuance", "Financing", 100_000), ("Dividends Paid", "Financing", -30_000),
    ]))
    # A second entity, so the entity selector is meaningful.
    flows.append(demo_income_statement("2000", 2024, 3, "2024-03-01", "2024-03-31", 640_000, 410_000,
                 [("Salaries & Benefits", 95_000), ("Rent & Facilities", 28_000),
                  ("Utilities", 12_000), ("Other Operating Expenses", 14_000)]))
    return flows


# ── Flow store ───────────────────────────────────────────────────────────────
def load_flows(path_or_none) -> list[dict]:
    if path_or_none:
        try:
            data = json.loads(Path(path_or_none).read_text())
            if isinstance(data, list) and data:
                return data
        except Exception:
            pass
    # Bundled file (e.g. shipped alongside the Space), else built-in demo.
    bundled = Path(__file__).with_name("sample_flows.json")
    if bundled.exists():
        try:
            data = json.loads(bundled.read_text())
            if isinstance(data, list) and data:
                return data
        except Exception:
            pass
    return _demo_flows()


def _period_label(flow: dict) -> str:
    fp = flow.get("fiscal_period", 0)
    fy = flow.get("fiscal_year", 0)
    return f"{fy} Full year" if fp == 0 else f"{fy}-{int(fp):02d}"


def _stmt_label(flow: dict) -> str:
    return "Income Statement" if flow.get("statement_type") == "income_statement" else "Cash Flow"


# ── Rendering ─────────────────────────────────────────────────────────────────
def render(flow: dict) -> go.Figure:
    nodes = flow["nodes"]
    labels = [f"{n['label']}  ({_num(n['value']):,.0f})" for n in nodes]
    node_colors = [SECTION_COLOR.get(n.get("section", ""), "#6b7280") for n in nodes]
    links = flow["links"]
    fig = go.Figure(go.Sankey(
        arrangement="snap",
        node=dict(label=labels, color=node_colors, pad=18, thickness=20,
                  line=dict(color="rgba(0,0,0,0.25)", width=0.5)),
        link=dict(
            source=[l["source"] for l in links],
            target=[l["target"] for l in links],
            value=[max(_num(l["value"]), 1e-9) for l in links],
            label=[l.get("label", "") for l in links],
            color=[CONTRA_LINK if l.get("is_contra") else NORMAL_LINK for l in links],
        ),
    ))
    cur = flow.get("currency", "USD")
    title = (f"{flow.get('company_code', '')} · {_stmt_label(flow)} · {_period_label(flow)} "
             f"({cur}){' · consolidated' if flow.get('is_consolidated') else ''}")
    fig.update_layout(title_text=title, font_size=12, height=560,
                      margin=dict(l=10, r=10, t=50, b=10))
    return fig


def _find(flows, company, period, stmt):
    for f in flows:
        if (f.get("company_code") == company and _period_label(f) == period
                and _stmt_label(f) == stmt):
            return f
    return None


# ── App ───────────────────────────────────────────────────────────────────────
def build_app():
    with gr.Blocks(title="Financial Sankey Explorer") as demo:
        # State MUST be created inside the Blocks context so it is registered
        # (a State created outside has an unregistered id → KeyError on every
        # callback that reads it, including the on-load).
        state = gr.State(_demo_flows())
        gr.Markdown(
            "# 💸 Financial Sankey Explorer\n"
            "Multi-step funds-flow view of synthetic financial statements from "
            "**`datasynth-data`** (`financial_reporting.sankey`). Income-statement "
            "waterfall and cash-flow funds-flow, per entity / period, with operating-"
            "expense sub-bands. Red bands are reductions / outflows; blue are retained "
            "amounts / inflows.\n\n"
            "Upload a `financial_reporting/sankey/flows.json` to explore your own run, "
            "or browse the built-in demo."
        )
        with gr.Row():
            upload = gr.File(label="flows.json (optional)", file_types=[".json"], scale=2)
            entity = gr.Dropdown(label="Entity", scale=1)
            period = gr.Dropdown(label="Period", scale=1)
            stmt = gr.Dropdown(label="Statement", scale=1)
        plot = gr.Plot(label="Funds flow")

        def _refresh(flows, company=None, period_sel=None, stmt_sel=None):
            companies = sorted({f["company_code"] for f in flows})
            company = company if company in companies else (companies[0] if companies else None)
            periods = sorted({_period_label(f) for f in flows if f["company_code"] == company},
                             key=lambda s: (("Full year" in s), s))
            period_sel = period_sel if period_sel in periods else (periods[0] if periods else None)
            stmts = sorted({_stmt_label(f) for f in flows
                            if f["company_code"] == company and _period_label(f) == period_sel})
            stmt_sel = stmt_sel if stmt_sel in stmts else (stmts[0] if stmts else None)
            flow = _find(flows, company, period_sel, stmt_sel)
            fig = render(flow) if flow else go.Figure()
            return (gr.update(choices=companies, value=company),
                    gr.update(choices=periods, value=period_sel),
                    gr.update(choices=stmts, value=stmt_sel), fig)

        def _on_load(file):
            flows = load_flows(file.name if file else None)
            e, p, s, fig = _refresh(flows)
            return flows, e, p, s, fig

        def _on_entity(flows, company):
            return _refresh(flows, company)[1:]

        def _on_period(flows, company, period_sel):
            return _refresh(flows, company, period_sel)[2:]

        def _on_stmt(flows, company, period_sel, stmt_sel):
            flow = _find(flows, company, period_sel, stmt_sel)
            return render(flow) if flow else go.Figure()

        upload.change(_on_load, inputs=upload, outputs=[state, entity, period, stmt, plot])
        entity.change(_on_entity, inputs=[state, entity], outputs=[period, stmt, plot])
        period.change(_on_period, inputs=[state, entity, period], outputs=[stmt, plot])
        stmt.change(_on_stmt, inputs=[state, entity, period, stmt], outputs=plot)

        demo.load(lambda flows: _refresh(flows), inputs=state, outputs=[entity, period, stmt, plot])
    return demo


# Module-level `demo` so the HF Spaces launcher auto-detects it.
demo = build_app()

if __name__ == "__main__":
    demo.launch()