palladium/app/components/tables.py

"""Streamlit data-editor wrappers for benefit/cost rows."""

from __future__ import annotations

import pandas as pd
import streamlit as st

from app.locale import currency_fmt, fmt_currency, fmt_pct, pct_fmt, _STANDARD_LOCALE


def _years_for_table(fields: list[dict], analysis_years: int) -> list[int]:
    """Years 1..N -- taken from analysis_period_years on the report."""
    return list(range(1, max(int(analysis_years or 3), 1) + 1))


def value_editor(
    table: str,
    fields: list[dict],
    values: list[dict],
    *,
    analysis_years: int,
    key: str,
) -> pd.DataFrame:
    """
    Render an ``st.data_editor`` for benefit or cost values.

    The editor shows one row per field (filtered to ``table``), with year
    columns, an ``initial`` column for costs, a risk_adjustment column, and
    a notes column. Returns the edited DataFrame; the caller is responsible
    for converting it back to value-row dicts and PUTting to Athena.

    Currency columns use the locale configured via PALLADIUM_CURRENCY_SYMBOL /
    PALLADIUM_THOUSANDS_SEP / PALLADIUM_DECIMAL_SEP in .env.
    The risk_adj column is stored as a 0-1 fraction and displayed as a
    percentage (e.g. 0.20 -> "20.00%").
    """
    fields = [
        f
        for f in fields
        if f.get("table") == table
        # Companion "<key>_initial" fields are edited via the Initial column
        # on their parent cost row, not as standalone rows.
        and not str(f.get("field_key", "")).endswith("_initial")
    ]
    fields.sort(key=lambda f: int(f.get("sort_order") or 0))

    by_key = {v.get("field_key"): v for v in values}
    years = _years_for_table(fields, analysis_years)

    rows: list[dict] = []
    for f in fields:
        v = by_key.get(f["field_key"], {}) or {}
        yv = v.get("year_values") or {}
        risk_raw = float(v.get("risk_adjustment") or 0.0)
        row = {
            "field_key": f["field_key"],
            "label": f.get("label", f["field_key"]),
            "category": f.get("category", "") or "",
        }
        if table == "costs":
            if _STANDARD_LOCALE:
                row["Initial"] = float(v.get("initial") or 0.0)
            else:
                row["Initial"] = fmt_currency(float(v.get("initial") or 0.0))
        for y in years:
            raw = float(yv.get(str(y)) or 0.0)
            if _STANDARD_LOCALE:
                row[f"Year {y}"] = raw
            else:
                row[f"Year {y}"] = fmt_currency(raw)
        # Risk adj: store as fraction for standard locales (NumberColumn handles
        # display), or pre-format as "20.00%" string for non-standard locales.
        if _STANDARD_LOCALE:
            row["risk_adj"] = risk_raw
        else:
            row["risk_adj"] = fmt_pct(risk_raw)
        row["notes"] = v.get("notes", "") or ""
        rows.append(row)

    df = pd.DataFrame(rows)

    _cur_fmt = currency_fmt()
    _pct_fmt_str = pct_fmt()

    column_config: dict = {
        "field_key": st.column_config.TextColumn("Key", disabled=True, width="small"),
        "label": st.column_config.TextColumn("Field", disabled=True),
        "category": st.column_config.TextColumn("Category", disabled=True, width="small"),
        "notes": st.column_config.TextColumn("Notes", width="medium"),
    }

    if _STANDARD_LOCALE:
        column_config["risk_adj"] = st.column_config.NumberColumn(
            "Risk Adj.",
            min_value=0.0,
            max_value=1.0,
            step=0.05,
            format=_pct_fmt_str,
            help="Enter as a decimal fraction (e.g. 0.20 = 20%)",
        )
        if table == "costs":
            column_config["Initial"] = st.column_config.NumberColumn(
                "Initial", format=_cur_fmt
            )
        for y in years:
            column_config[f"Year {y}"] = st.column_config.NumberColumn(
                f"Year {y}", format=_cur_fmt
            )
    else:
        # Non-standard locale: display as pre-formatted strings (read-only display;
        # user edits the raw number and we re-format on save).
        column_config["risk_adj"] = st.column_config.TextColumn(
            "Risk Adj.", help="Displayed as percentage; stored as 0-1 fraction"
        )
        if table == "costs":
            column_config["Initial"] = st.column_config.TextColumn("Initial")
        for y in years:
            column_config[f"Year {y}"] = st.column_config.TextColumn(f"Year {y}")

    edited = st.data_editor(
        df,
        column_config=column_config,
        width="stretch",
        num_rows="fixed",
        hide_index=True,
        key=key,
    )
    return edited


def df_to_values(df: pd.DataFrame, table: str, analysis_years: int) -> list[dict]:
    """Convert an edited DataFrame back to wire-format value rows."""
    out: list[dict] = []
    years = list(range(1, max(int(analysis_years or 3), 1) + 1))
    for _, row in df.iterrows():
        item: dict = {"field_key": row["field_key"], "table": table}
        yv = {}
        for y in years:
            col = f"Year {y}"
            if col in df.columns:
                yv[str(y)] = float(row[col] or 0)
        if yv:
            item["year_values"] = yv
        if table == "costs" and "Initial" in df.columns:
            item["initial"] = float(row["Initial"] or 0)
        ra = row.get("risk_adj")
        if ra is not None and not pd.isna(ra):
            item["risk_adjustment"] = float(ra)
        notes = row.get("notes")
        if isinstance(notes, str) and notes.strip():
            item["notes"] = notes.strip()
        out.append(item)
    return out