cli.py

"""PolicyEngine Macro CLI. Human-readable tables by default; --json for machine output."""

from __future__ import annotations

import json
import sys

import click

from policyengine_macro import core
from policyengine_macro import capabilities
from policyengine_macro import reporting


def _emit_json(obj) -> None:
    click.echo(json.dumps(obj, indent=2))


def _load_payload_file(path: str | None, option: str) -> dict | None:
    """Load a pre-computed macro payload JSON for a two-process pipeline."""
    if path is None:
        return None
    try:
        with open(path) as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError) as e:
        raise click.ClickException(
            f"{option} {path}: not readable JSON ({e}); pass the unmodified "
            "--json output of the corresponding shock command"
        ) from e


def _table(rows: list[dict], columns: list[str]) -> str:
    widths = {c: max(len(c), *(len(str(r.get(c, ""))) for r in rows)) for c in columns}
    head = "  ".join(c.ljust(widths[c]) for c in columns)
    sep = "  ".join("-" * widths[c] for c in columns)
    body = "\n".join(
        "  ".join(str(r.get(c, "")).ljust(widths[c]) for c in columns) for r in rows
    )
    return f"{head}\n{sep}\n{body}"


@click.group()
def main() -> None:
    """PolicyEngine Macro: unified CLI over the OBR emulator and the UK SVAR model."""


@main.command("model-status")
@click.argument("model_id", required=False)
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def model_status(model_id, as_json):
    """Show supported uses, access, and limitations for one or all models."""
    try:
        rows = ([capabilities.get_status(model_id)] if model_id
                else capabilities.list_capabilities())
    except ValueError as e:
        raise click.ClickException(str(e)) from e
    if as_json:
        _emit_json(rows[0] if model_id else rows)
        return
    click.echo(_table([
        {
            "model": row["model_id"],
            "country": ",".join(row["geography"]),
            "status": row["status"],
            "access": "; ".join(row["access"]),
        }
        for row in rows
    ], ["model", "country", "status", "access"]))


@main.command()
@click.option("--country", type=click.Choice(["uk", "us"]), default="uk",
              show_default=True)
@click.option("--reform", required=True,
              help='PolicyEngine reform JSON, e.g. \'{"gov.hmrc.income_tax.rates.uk[0].rate":0.21}\' '
                   "(same shape as `pe-macro population-impact`).")
# Not a click.Choice: a Choice rejects `svar`, `frbus`, `hank` and `define`
# with a bare usage error listing the valid names, which is exactly the
# unhelpful message core.score_reform was given explicit refusals to replace.
# Accepting any string and letting core refuse means a user who reasonably
# tries --model svar is told there is no reform bridge by design, and which
# tool to use instead. Unknown names still hit core's enum error.
@click.option("--model", required=True,
              help="Scoring model: og (OG-UK steady state; slow), obr (OBR "
                   "emulator via the microsim static-costing bridge), "
                   "microsim (PolicyEngine population costing, no macro "
                   "feedback), or og+microsim (dynamic scoring: OG earnings "
                   "overlay on the microsim; UK, local-only).")
@click.option("--year", default=2026, show_default=True, help="Reform start year.")
@click.option("--max-iter", default=250, show_default=True,
              help="og only: solver iteration cap per steady-state solve.")
@click.option("--years", default=5, show_default=True,
              help="obr only: costing window length in years.")
@click.option("--dataset", default=None,
              help="obr/microsim only: microdata dataset name override.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def score(country, reform, model, year, max_iter, years, dataset, as_json):
    """Score a PolicyEngine reform with a scoring model of the suite.

    One reform vocabulary: the same {parameter_path: value} dict as
    `pe-macro population-impact`. Every result carries a common `score`
    block for cross-model comparison (`pe-macro compare`). For raw OBR
    variable shocks in model units, use `pe-macro obr-shock`.
    """
    try:
        res = core.score_reform(
            country=country, reform=_json_opt(reform, "reform"), model=model,
            start_year=year, max_iter=max_iter, years=years, dataset=dataset,
        )
    except (NotImplementedError, ValueError, ImportError, RuntimeError) as e:
        raise click.ClickException(str(e)) from e
    if as_json:
        _emit_json(res)
        return
    if model == "og":
        _echo_og_impact(res)
    else:
        _echo_score_block(res["score"])


def _echo_score_block(score: dict) -> None:
    """Render one common ScoreResult block as a table."""
    click.echo(f"{score['model']} ({score['model_class']}, "
               f"{score['country'].upper()}, {score['horizon']})")
    click.echo(f"Reform: {score['reform']}\n")
    rows = []
    for name, q in score["quantities"].items():
        rows.append({
            "quantity": name,
            "delta_bn": q.get("delta_bn"),
            "delta_pct": q.get("delta_pct"),
            "units": q["units"],
        })
    click.echo(_table(rows, ["quantity", "delta_bn", "delta_pct", "units"]))
    for label, items in (("Assumptions", score.get("assumptions") or []),
                         ("Caveats", score.get("caveats") or [])):
        if items:
            click.echo(f"\n{label}:")
            for it in items:
                click.echo(f"  - {it}")


@main.command("report")
@click.argument("input_file", type=click.Path(exists=True, dir_okay=False),
                required=False)
@click.option("--format", "output_format",
              type=click.Choice(["json", "markdown"]), default="markdown",
              show_default=True)
def report(input_file, output_format):
    """Render a stored ScoreResult (or a response containing ``score``).

    Reads INPUT_FILE, or standard input when omitted.
    """
    source = open(input_file, encoding="utf-8") if input_file else sys.stdin
    try:
        payload = json.load(source)
    except (OSError, json.JSONDecodeError) as e:
        raise click.ClickException(f"could not read result JSON: {e}") from e
    finally:
        if input_file:
            source.close()
    try:
        if output_format == "json":
            _emit_json(reporting.build_report(payload))
        else:
            click.echo(reporting.render_markdown(payload), nl=False)
    except ValueError as e:
        raise click.ClickException(f"invalid ScoreResult: {e}") from e


@main.command()
@click.option("--country", type=click.Choice(["uk", "us"]), default="uk",
              show_default=True)
@click.option("--reform", required=True,
              help='PolicyEngine reform JSON (same shape as `pe-macro score`).')
@click.option("--models", default="microsim,obr", show_default=True,
              help="Comma-separated scoring models (og, obr, microsim).")
@click.option("--year", default=2026, show_default=True, help="Reform start year.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON list of ScoreResults.")
def compare(country, reform, models, year, as_json):
    """Run one reform through supported adapters, with comparability warnings.

    Runs `score` once per model and renders one table from the common
    ScoreResult blocks (PolicyEngine/macro#10)."""
    reform_dict = _json_opt(reform, "reform")
    scores = []
    for model in [m.strip() for m in models.split(",") if m.strip()]:
        try:
            res = core.score_reform(
                country=country, reform=reform_dict, model=model,
                start_year=year,
            )
        except (NotImplementedError, ValueError, ImportError, RuntimeError) as e:
            raise click.ClickException(f"{model}: {e}") from e
        scores.append(res["score"])
    if as_json:
        _emit_json(scores)
        return
    click.echo(f"Reform: {reform_dict}  ({country.upper()}, from {year})\n")
    rows = []
    for s in scores:
        for name, q in s["quantities"].items():
            rows.append({
                "model": s["model"],
                "class": s["model_class"],
                "horizon": s["horizon"],
                "quantity": name,
                "delta_bn": q.get("delta_bn"),
                "delta_pct": q.get("delta_pct"),
                "units": q["units"],
                "time_basis": q["time_basis"],
                "comparability": q["comparability"],
            })
    click.echo(_table(rows, ["model", "class", "horizon", "quantity",
                             "delta_bn", "delta_pct", "units", "time_basis",
                             "comparability"]))
    click.echo("\nThese results use different horizons and mechanisms. "
               "Treat related-not-like-for-like rows as complementary: they "
               "must not be added, averaged, or ranked.")


@main.command("obr-shock")
@click.option("--var", required=True, help="Policy variable to shock (see `pe-macro variables`).")
@click.option("--shock", required=True, type=float,
              help="Shock size; units depend on the variable (£m/quarter for CGG, decimal for TCPRO).")
@click.option("--periods", default=12, show_default=True, help="Quarters the shock is applied.")
@click.option("--name", default=None, help="Label for the reform.")
@click.option("--investment-closure/--no-investment-closure", default=None,
              help="Cost-of-capital investment channel; omit for the safe "
                   "per-variable default (on for TCPRO, off otherwise).")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def obr_shock(var, shock, periods, name, investment_closure, as_json):
    """Shock one OBR variable directly, in model units (escape hatch)."""
    res = core.obr_shock(
        var=var, shock=shock, periods=periods, name=name,
        investment_closure=investment_closure,
    )
    if as_json:
        _emit_json(res)
        return
    click.echo(f"Reform: {res['name']}  (var={res['var']}, shock={res['shock']:+g}, "
               f"periods={res['periods']}, investment_closure={res['investment_closure']})")
    click.echo(_table(res["results"],
                      ["period", "delta_gdp_bn", "pct_gdp", "delta_cons_m", "delta_if_m"]))
    click.echo(f"\nCumulative GDP effect over shocked periods: "
               f"£{res['cumulative_delta_gdp_bn_over_shock_periods']}bn")
    click.echo(f"Peak GDP effect: {res['peak_pct_gdp']}% in {res['peak_period']}")


@main.command()
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def variables(as_json):
    """List commonly shocked OBR policy variables."""
    res = core.obr_list_variables()
    if as_json:
        _emit_json(res)
        return
    click.echo(_table(res, ["var", "description", "units", "investment_closure"]))


@main.command("frbus-shock")
@click.option("--var", required=True,
              help="Lever to shock (see `pe-macro frbus-variables`).")
@click.option("--shock", required=True, type=float,
              help="Shock size; UNITS DIFFER PER LEVER (pp for rffintay_aerr, "
                   "decimal rate for trp_aerr, log points of quarterly growth "
                   "for egfe_aerr/ecnia_aerr).")
@click.option("--start", default=core.FRBUS_DEFAULT_START, show_default=True,
              help="First shocked quarter, e.g. 2026Q1.")
@click.option("--periods", default=1, show_default=True,
              help="Quarters the shock is held (1 = single-quarter impulse).")
@click.option("--horizon", default=core.FRBUS_DEFAULT_HORIZON, show_default=True,
              help="Quarters simulated and reported.")
@click.option("--policy-rule", default="inertial_taylor", show_default=True,
              type=click.Choice(sorted(core.FRBUS_POLICY_RULES)),
              help="Monetary policy reaction; changes the answer materially.")
@click.option("--variable", "variables", multiple=True,
              help="Extra model variable to report (repeatable).")
@click.option("--name", default=None, help="Label for the experiment.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def frbus_shock(var, shock, start, periods, horizon, policy_rule, variables,
                name, as_json):
    """Shock one FRB/US variable, in model units (US escape hatch)."""
    res = core.frbus_shock(
        var=var, shock=shock, start=start, periods=periods, horizon=horizon,
        policy_rule=policy_rule, variables=list(variables) or None, name=name,
    )
    if as_json:
        _emit_json(res)
        return
    click.echo(f"Experiment: {res['name']}  (var={res['var']}, "
               f"shock={res['shock']:+g}, start={res['start']}, "
               f"periods={res['periods']}, rule={res['policy_rule']})")
    click.echo(f"Units: {res['units']}")
    columns = ["period"] + [k for k in res["results"][0] if k != "period"]
    click.echo(_table(res["results"], columns))
    click.echo("\nPeak absolute deviations:")
    for v, peak in res["peaks"].items():
        click.echo(f"  {v:10s} {peak['value']:+.4f} in {peak['period']}"
                   f"   ({res['series_meaning'][v]})")
    if res.get("warning"):
        click.echo(f"\nWARNING: {res['warning']}")


def _echo_incidence(res: dict, shock_line: str) -> None:
    ea = res["economic_assumptions"]
    click.echo(shock_line)
    click.echo(f"Earnings factor: {ea['earnings_factor']}   "
               f"Labour-supply factor: {ea['labour_supply_factor']}   "
               f"applied: {res['application']['applied']}\n")
    micro = res["microsim"]
    click.echo(micro["headline"])
    sym = "£" if micro["country"] == "uk" else "$"
    click.echo(f"Budget change: {sym}{micro['budgetary_impact_bn']}bn/year "
               f"({micro['budgetary_impact_basis']})")
    click.echo(f"Household net income change: "
               f"{sym}{micro['household_net_income_change_bn']}bn/year")
    click.echo(f"Winners: {micro['winners']:,}   Losers: {micro['losers']:,}\n")
    click.echo(_table(micro["decile_impacts"],
                      ["decile", "avg_income_change", "relative_change_pct",
                       "count_better_off", "count_worse_off"]))
    for label, items in (("Assumptions", res.get("assumptions") or []),
                         ("Caveats", res.get("caveats") or [])):
        click.echo(f"\n{label}:")
        for it in items:
            click.echo(f"  - {it}")


@main.command("frbus-shock-incidence")
@click.option("--var", required=True,
              help="FRB/US lever to shock (see `pe-macro frbus-variables`).")
@click.option("--shock", required=True, type=float,
              help="Shock size in the lever's model units (units differ "
                   "per lever; see `pe-macro frbus-variables`).")
@click.option("--year", default=2027, show_default=True,
              help="Calendar year whose four quarters are averaged into the "
                   "earnings overlay and scored by the microsim.")
@click.option("--start", default=core.FRBUS_DEFAULT_START, show_default=True,
              help="First shocked quarter, e.g. 2026Q1.")
@click.option("--periods", default=1, show_default=True,
              help="Quarters the shock is held.")
@click.option("--horizon", default=core.FRBUS_DEFAULT_HORIZON,
              show_default=True,
              help="Quarters simulated (must cover the incidence year).")
@click.option("--policy-rule", default="inertial_taylor", show_default=True,
              type=click.Choice(sorted(core.FRBUS_POLICY_RULES)),
              help="Monetary policy reaction.")
@click.option("--income-concept", default="wage_bill", show_default=True,
              type=click.Choice(["wage", "wage_bill"]),
              help="wage = compensation per hour only; wage_bill also "
                   "carries the hours change (applied uniformly).")
@click.option("--dataset", default=None,
              help="Microdata dataset name override.")
@click.option("--frbus-payload", "frbus_payload_path", default=None,
              help="Path to a pre-computed `pe-macro frbus-shock --json` "
                   "result (run with pl/lhp/leh in --variables), so the "
                   "FRB/US solve and the microsim run in separate processes "
                   "on memory-constrained machines.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def frbus_shock_incidence(var, shock, year, start, periods, horizon,
                          policy_rule, income_concept, dataset,
                          frbus_payload_path, as_json):
    """Who bears a FRB/US shock, at household resolution (experimental).

    Runs frbus_shock, turns the year's mean pl/lhp deviations into a
    pre-tax earnings factor, scales the US microsim's employment-income
    inputs (no reform on either side), and reports the automatic-stabilizer
    budget change plus decile impacts. NOT reform scoring.
    """
    frbus_payload = _load_payload_file(frbus_payload_path, "--frbus-payload")
    try:
        res = core.frbus_shock_incidence(
            var=var, shock=shock, year=year, start=start, periods=periods,
            horizon=horizon, policy_rule=policy_rule,
            income_concept=income_concept, dataset=dataset,
            frbus_payload=frbus_payload,
        )
    except (ValueError, ImportError, RuntimeError) as e:
        raise click.ClickException(str(e)) from e
    if as_json:
        _emit_json(res)
        return
    _echo_incidence(res, (
        f"FRB/US shock incidence: {res['frbus']['name']} "
        f"(year {res['year']}, income_concept={res['income_concept']})"
    ))


@main.command("hank-shock-incidence")
@click.option("--kind", required=True,
              type=click.Choice([k["kind"] for k in core.HANK_SHOCK_KINDS]),
              help="Shock kind (see `pe-macro hank-summary` for units).")
@click.option("--size", required=True, type=float,
              help="Impact size in model units (units differ per kind).")
@click.option("--year", default=2026, show_default=True,
              help="Calendar year whose four quarters are averaged into the "
                   "earnings overlay and scored by the microsim.")
@click.option("--persistence", default=0.9, show_default=True, type=float,
              help="AR(1) decay of the shock path (in [0, 1)).")
@click.option("--horizon", default=core.HANK_DEFAULT_HORIZON,
              show_default=True,
              help="Quarters simulated (must cover the incidence year).")
@click.option("--variant", default="two_asset", show_default=True,
              type=click.Choice(list(core.HANK_VARIANTS)))
@click.option("--income-concept", default="wage_bill", show_default=True,
              type=click.Choice(["wage", "wage_bill"]),
              help="wage = real wage w only; wage_bill also carries the "
                   "labor N change (applied uniformly).")
@click.option("--start-year", "start_year", default=2026, show_default=True,
              help="Calendar year the shock's quarter 0 maps to.")
@click.option("--dataset", default=None,
              help="Microdata dataset name override.")
@click.option("--hank-payload", "hank_payload_path", default=None,
              help="Path to a pre-computed `pe-macro hank-shock --json` "
                   "result, so the HANK solve and the microsim run in "
                   "separate processes on memory-constrained machines.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def hank_shock_incidence(kind, size, year, persistence, horizon, variant,
                         income_concept, start_year, dataset,
                         hank_payload_path, as_json):
    """Who bears a US HANK shock, at household resolution (experimental).

    Runs hank_shock (which surfaces the pre-tax real wage w and labor N
    IRFs), turns the year's mean deviations into a pre-tax earnings factor,
    scales the US microsim's employment-income inputs (no reform), and
    reports the automatic-stabilizer budget change plus decile impacts.
    Stylized calibrated model; NOT a forecaster, NOT reform scoring.
    """
    hank_payload = _load_payload_file(hank_payload_path, "--hank-payload")
    try:
        res = core.hank_shock_incidence(
            kind=kind, size=size, year=year, persistence=persistence,
            horizon=horizon, variant=variant, income_concept=income_concept,
            start_year=start_year, dataset=dataset,
            hank_payload=hank_payload,
        )
    except (ValueError, ImportError, RuntimeError) as e:
        raise click.ClickException(str(e)) from e
    if as_json:
        _emit_json(res)
        return
    _echo_incidence(res, (
        f"US HANK shock incidence: {res['hank']['name']} "
        f"(year {res['year']}, income_concept={res['income_concept']})"
    ))


@main.command("svar-inflation-incidence")
@click.option("--year", default=2027, show_default=True,
              help="Forecast year whose CPI gap drives the following "
                   "April's uprating.")
@click.option("--horizons", default=12, show_default=True,
              help="SVAR forecast horizon in quarters (must cover the year).")
@click.option("--draws", default=2000, show_default=True,
              help="SVAR posterior draws (first call takes minutes).")
@click.option("--reference", default="obr", show_default=True,
              type=click.Choice(["obr", "target"]),
              help="CPI reference path: the March 2026 EFO (obr) or a flat "
                   "2.0% (target).")
@click.option("--dataset", default=None,
              help="Microdata dataset name override.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def svar_inflation_incidence(year, horizons, draws, reference, dataset,
                             as_json):
    """Cost and incidence of the SVAR-vs-reference CPI gap via uprating.

    Compares the UK SVAR's median CPI path for the year against the
    reference, scales a short curated list of statutorily CPI-uprated
    benefit parameters by the gap from the following 6 April, and scores
    that real reform with the UK population microsim. Excludes the state
    pension triple lock and frozen tax thresholds (stated in caveats).
    """
    try:
        res = core.svar_inflation_incidence(
            year=year, horizons=horizons, draws=draws, reference=reference,
            dataset=dataset,
        )
    except (ValueError, ImportError, RuntimeError) as e:
        raise click.ClickException(str(e)) from e
    if as_json:
        _emit_json(res)
        return
    click.echo(f"UK SVAR inflation-uprating incidence, {res['year']} gap -> "
               f"April {res['uprating_year']} uprating")
    click.echo(f"SVAR CPI {res['svar_cpi_yoy_pct']}% vs reference "
               f"{res['reference_cpi_yoy_pct']}% ({res['reference_description']}): "
               f"gap {res['cpi_gap_pp']:+}pp\n")
    click.echo(_table(res["parameters"],
                      ["path", "description", "unit", "baseline_value",
                       "counterfactual_value"]))
    micro = res["microsim"]
    click.echo(f"\n{res['headline']}")
    click.echo(f"Winners: {micro['winners']:,}   Losers: {micro['losers']:,}\n")
    click.echo(_table(micro["decile_impacts"],
                      ["decile", "avg_income_change", "relative_change_pct",
                       "count_better_off", "count_worse_off"]))
    for label, items in (("Assumptions", res["assumptions"]),
                         ("Caveats", res["caveats"])):
        click.echo(f"\n{label}:")
        for it in items:
            click.echo(f"  - {it}")


@main.command("frbus-variables")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def frbus_variables(as_json):
    """List the shockable FRB/US levers and their units."""
    res = core.frbus_list_variables()
    if as_json:
        _emit_json(res)
        return
    click.echo(_table(res, ["var", "description", "units", "typical_shock",
                            "requires_policy_rule"]))


@main.command("frbus-summary")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def frbus_summary(as_json):
    """FRB/US model metadata and validation provenance (instant)."""
    res = core.frbus_summary()
    if as_json:
        _emit_json(res)
        return
    click.echo(f"{res['model']} — {res['implementation']}")
    click.echo(f"  {res['equations']} endogenous equations, "
               f"{res['data_vintage']}, {res['expectations']}")
    click.echo(f"  source: {res.get('source', res.get('source_error'))}\n")
    click.echo("Policy rules:")
    for rule in res["policy_rules"]:
        click.echo(f"  {rule['rule']:18s} {rule['description']}")
    val = res["validation"]
    click.echo("\nValidation:")
    click.echo(f"  tracking invariant: {val['tracking_invariant']['value']:.1e} "
               f"(gate {val['tracking_invariant']['gate']:.0e})")
    click.echo(f"  vs pyfrbus 1.0.0:   {val['vs_vendor_pyfrbus']['value']:.1e} "
               f"(gate {val['vs_vendor_pyfrbus']['gate']:.0e})")
    mon = val["monetary_tightening_properties"]
    click.echo(f"  {mon['shock']}: xgdp trough {mon['xgdp_trough_pct']}%, "
               f"lur peak {mon['lur_peak_pp']}pp")
    click.echo(f"\nReform bridge: {res['reform_bridge']}")


@main.command("hank-shock")
@click.option("--kind", required=True,
              type=click.Choice([k["kind"] for k in core.HANK_SHOCK_KINDS]),
              help="Shock kind (see `pe-macro hank-summary` for units).")
@click.option("--size", required=True, type=float,
              help="Impact size in model units; UNITS DIFFER PER KIND "
                   "(quarterly-rate level for monetary, level of G for "
                   "fiscal_spending, level of Z for productivity).")
@click.option("--persistence", default=0.9, show_default=True, type=float,
              help="AR(1) decay of the shock path (in [0, 1)).")
@click.option("--horizon", default=core.HANK_DEFAULT_HORIZON, show_default=True,
              help="Quarters reported.")
@click.option("--variant", default="two_asset", show_default=True,
              type=click.Choice(list(core.HANK_VARIANTS)),
              help="two_asset (the paper model) or one_asset (fast, no "
                   "capital, monetary/productivity only).")
@click.option("--distribution", "include_distribution", is_flag=True,
              help="two_asset only: add MPC-by-quartile, hand-to-mouth share "
                   "and the first-order impact consumption response by "
                   "wealth quartile.")
@click.option("--name", default=None, help="Label for the experiment.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def hank_shock(kind, size, persistence, horizon, variant,
               include_distribution, name, as_json):
    """Run a stylized US HANK shock (validated replication; not a forecaster)."""
    try:
        res = core.hank_shock(
            kind=kind, size=size, persistence=persistence, horizon=horizon,
            variant=variant, include_distribution=include_distribution,
            name=name,
        )
    except (ValueError, ImportError) as e:
        raise click.ClickException(str(e)) from e
    if as_json:
        _emit_json(res)
        return
    click.echo(f"Experiment: {res['name']}  (kind={res['kind']}, "
               f"size={res['size']:+g}, persistence={res['persistence']:g}, "
               f"variant={res['variant']})")
    click.echo(f"Units: {res['units']}")
    click.echo(f"Framing: {res['framing']}")
    columns = list(res["results"][0])
    click.echo(_table(res["results"], columns))
    click.echo("\nPeak absolute deviations:")
    for v, peak in res["peaks"].items():
        click.echo(f"  {v:3s} {peak['value']:+.4f} at quarter {peak['quarter']}"
                   f"   ({res['series_meaning'][v]})")
    if res.get("distributional"):
        d = res["distributional"]
        click.echo("\nDistributional (first-order approximation):")
        click.echo(f"  aggregate quarterly MPC (liquid): "
                   f"{d['aggregate_quarterly_mpc_liquid']}")
        click.echo(f"  MPC by liquid-wealth quartile:    "
                   f"{d['mpc_by_liquid_quartile']}")
        click.echo(f"  hand-to-mouth share:              "
                   f"{d['hand_to_mouth_share']}")
        click.echo(f"  impact dC (%) by wealth quartile: "
                   f"{d['impact_consumption_response_pct_by_wealth_quartile']}")
    if res.get("warning"):
        click.echo(f"\nWARNING: {res['warning']}")


@main.command("hank-summary")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def hank_summary(as_json):
    """US HANK model metadata, shock catalogue and scope limits (instant)."""
    res = core.hank_summary()
    if as_json:
        _emit_json(res)
        return
    click.echo(f"{res['model']} — {res['implementation']}")
    click.echo(f"  upstream: {res['upstream']}")
    click.echo(f"  {res['framing']}\n")
    click.echo("Variants:")
    for variant, desc in res["variants"].items():
        click.echo(f"  {variant:10s} {desc}")
    click.echo("\nShock kinds:")
    for k in res["shock_kinds"]:
        click.echo(f"  {k['kind']:16s} ({'/'.join(k['variants'])}) {k['units']}")
    click.echo(f"\n{res['no_tax_or_transfer_instrument']}")
    click.echo(f"\nValidation: {res['validation']['suite']} "
               f"({res['validation']['note']})")
    click.echo(f"\nReform bridge: {res['reform_bridge']}")


@main.command()
@click.option("--horizons", default=12, show_default=True, help="Forecast horizon in quarters.")
@click.option("--draws", default=2000, show_default=True,
              help="Posterior draws (more = slower, smoother; 2000 takes "
                   "~2 min and ~3500 reaches importance-weight ESS >= 100).")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def forecast(horizons, draws, as_json):
    """UK SVAR forecast: YoY GDP growth and CPI inflation with bands."""
    res = core.svar_forecast(horizons=horizons, draws=draws)
    if as_json:
        _emit_json(res)
        return
    click.echo(f"UK SVAR forecast from {res['forecast_origin']} "
               f"({res['draws']} draws, {res['accepted_draws']} accepted, ESS {res['ess']})")
    for msg in res.get("warnings", []):
        click.secho(f"WARNING: {msg}", fg="yellow", err=True)
    for key, label in [("gdp_growth_yoy", "YoY GDP growth (%)"),
                       ("cpi_inflation_yoy", "YoY CPI inflation (%)")]:
        click.echo(f"\n{label}")
        click.echo(_table(res[key], ["quarter", "median", "lo68", "hi68", "lo90", "hi90"]))


@main.command()
@click.option("--draws", default=2000, show_default=True, help="Posterior draws.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def shocks(draws, as_json):
    """P(sign) of the identified structural shocks in the latest quarter."""
    res = core.svar_latest_shocks(draws=draws)
    if as_json:
        _emit_json(res)
        return
    click.echo(f"Structural shocks in {res['quarter']} "
               f"({res['draws']} draws, {res['accepted_draws']} accepted, ESS {res['ess']})")
    for msg in res.get("warnings", []):
        click.secho(f"WARNING: {msg}", fg="yellow", err=True)
    click.echo()
    click.echo(_table(res["shocks"], ["shock", "p_positive", "p_negative"]))
    click.echo()
    for s in res["shocks"]:
        click.echo(f"- {s['reading']}")


@main.command()
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def summary(as_json):
    """Headline SVAR results parsed from the repo's committed summaries (instant)."""
    res = core.svar_summary()
    if set(res) == {"error"}:
        raise click.ClickException(res["error"])
    rep = res.get("replication", {})
    fr_check = res.get("forecast_revision", {})
    if "error" in rep and "error" in fr_check:
        raise click.ClickException(
            "no parseable SVAR results — replication: "
            f"{rep['error']}; forecast revision: {fr_check['error']}"
        )
    if as_json:
        _emit_json(res)
        return
    click.echo("Replication (results/summary.md)")
    if "error" in rep:
        click.echo(f"  error: {rep['error']}", err=True)
    for ln in rep.get("metadata", []):
        click.echo(f"  {ln}")
    fevd = rep.get("fevd_1yr_headline", [])
    if fevd:
        click.echo("\nFEVD at 1-year horizon (sum of medians, renormalised)")
        click.echo(_table(fevd, list(fevd[0].keys())))
    groups = rep.get("fevd_1yr_group_shares", [])
    if groups:
        click.echo("\nGroup share formed per draw — the paper-comparable "
                   "statistic is the mean column")
        click.echo(_table(groups, list(groups[0].keys())))
    if rep.get("fevd_note"):
        click.echo(f"\n  {rep['fevd_note']}")
    fr = res.get("forecast_revision", {})
    click.echo("\nForecast-revision exercise (results/forecast_summary.md)")
    if "error" in fr:
        click.echo(f"  error: {fr['error']}", err=True)
    for ln in fr.get("metadata", []):
        click.echo(f"  {ln}")
    signs = fr.get("latest_shock_signs", [])
    if signs:
        click.echo("\nP(sign) of the identified shocks at T")
        click.echo(_table(signs, list(signs[0].keys())))
    for ln in fr.get("composite_irf", []):
        click.echo(f"  {ln}")


def _json_opt(value, name):
    if value is None:
        return None
    try:
        return json.loads(value)
    except json.JSONDecodeError as e:
        raise click.BadParameter(f"--{name} must be valid JSON: {e}") from e


_PEOPLE_HELP = 'JSON list of person dicts, e.g. \'[{"age":35,"employment_income":50000}]\'.'
_REFORM_HELP = (
    'JSON reform dict, e.g. \'{"gov.hmrc.income_tax.rates.uk[0].rate":0.25}\'. '
    'A value may also be a single-effective-date dict, e.g. '
    '\'{"gov.hmrc.income_tax.rates.uk[0].rate":{"2026-01-01":0.25}}\'; '
    'date ranges ("2026-01-01.2029-12-31") are not supported.'
)


def _pe_common_options(fn):
    for opt in reversed([
        click.option("--country", type=click.Choice(["uk", "us"]), required=True),
        click.option("--people", required=True, help=_PEOPLE_HELP),
        click.option("--year", default=2026, show_default=True),
        click.option("--benunit", default=None, help="UK only: JSON benunit dict."),
        click.option("--tax-unit", "tax_unit", default=None,
                     help='US only: JSON tax unit dict, e.g. \'{"filing_status":"SINGLE"}\'.'),
        click.option("--household", default=None,
                     help='JSON household dict, e.g. \'{"state_code_str":"CA"}\' (US).'),
        click.option("--json", "as_json", is_flag=True, help="Emit JSON."),
    ]):
        fn = opt(fn)
    return fn


def _echo_summary(label: str, summary: dict, sym: str) -> None:
    click.echo(label)
    for k, v in summary.items():
        if isinstance(v, list):
            v = ", ".join(f"{sym}{x:,.0f}" for x in v)
        elif isinstance(v, (int, float)):
            v = f"{sym}{v:,.0f}"
        click.echo(f"  {k:32} {v}")


@main.command()
@_pe_common_options
@click.option("--reform", default=None, help=_REFORM_HELP)
def household(country, people, year, benunit, tax_unit, household, as_json, reform):
    """Calculate taxes and benefits for a household (PolicyEngine)."""
    res = core.pe_household(
        country=country,
        people=_json_opt(people, "people"),
        year=year,
        reform=_json_opt(reform, "reform"),
        benunit=_json_opt(benunit, "benunit"),
        tax_unit=_json_opt(tax_unit, "tax-unit"),
        household=_json_opt(household, "household"),
    )
    if as_json:
        _emit_json(res)
        return
    sym = "£" if res["country"] == "uk" else "$"
    click.echo(f"PolicyEngine {res['country'].upper()} household, {res['year']}"
               + (f" (reform: {res['reform']})" if res["reform"] else ""))
    _echo_summary("Summary:", res["summary"], sym)


@main.command("household-impact")
@_pe_common_options
@click.option("--reform", required=True, help=_REFORM_HELP)
def household_impact(country, people, year, benunit, tax_unit, household, as_json, reform):
    """Baseline-vs-reform impact of a reform on one household (PolicyEngine)."""
    res = core.pe_household_impact(
        country=country,
        people=_json_opt(people, "people"),
        reform=_json_opt(reform, "reform"),
        year=year,
        benunit=_json_opt(benunit, "benunit"),
        tax_unit=_json_opt(tax_unit, "tax-unit"),
        household=_json_opt(household, "household"),
    )
    if as_json:
        _emit_json(res)
        return
    sym = "£" if res["country"] == "uk" else "$"
    click.echo(f"PolicyEngine {res['country'].upper()} reform impact, {res['year']}")
    click.echo(f"Reform: {res['reform']}\n")
    _echo_summary("Baseline:", res["baseline"], sym)
    _echo_summary("\nWith reform:", res["with_reform"], sym)
    _echo_summary("\nChange:", {k: v for k, v in res["change"].items() if v is not None}, sym)


@main.command("population-impact")
@click.option("--country", type=click.Choice(["uk", "us"]), default="uk",
              show_default=True)
@click.option("--reform", required=True, help=_REFORM_HELP)
@click.option("--year", default=2026, show_default=True)
@click.option("--dataset", default=None,
              help="Dataset name (default: enhanced_frs_2023_24 for UK).")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def population_impact(country, reform, year, dataset, as_json):
    """Population-level reform score (PolicyEngine microsimulation).

    Budgetary impact in £bn/$bn per year plus decile impacts. First UK run
    downloads private microdata (set HUGGING_FACE_TOKEN); afterwards a score
    takes tens of seconds.
    """
    res = core.pe_population_impact(
        country=country, reform=_json_opt(reform, "reform"),
        year=year, dataset=dataset,
    )
    if as_json:
        _emit_json(res)
        return
    sym = "£" if res["country"] == "uk" else "$"
    click.echo(f"PolicyEngine {res['country'].upper()} population impact, "
               f"{res['year']} ({res['dataset']}, "
               f"{res['n_households']:,} households)")
    click.echo(f"Reform: {res['reform']}\n")
    click.echo(res["headline"])
    click.echo(f"Budgetary impact: {sym}{res['budgetary_impact_bn']}bn/year "
               f"({res['budgetary_impact_basis']})")
    click.echo(f"Household net income change: "
               f"{sym}{res['household_net_income_change_bn']}bn/year")
    click.echo(f"Winners: {res['winners']:,}   Losers: {res['losers']:,}\n")
    click.echo(_table(res["decile_impacts"],
                      ["decile", "avg_income_change", "relative_change_pct",
                       "count_better_off", "count_worse_off"]))


@main.command()
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def parameters(as_json):
    """List curated PolicyEngine reform parameters (verified paths)."""
    res = core.pe_list_common_parameters()
    if as_json:
        _emit_json(res)
        return
    cols = ["country", "path", "description", "unit"]
    for extra in ("baseline_value", "live"):
        if any(extra in r for r in res):
            cols.append(extra)
    click.echo(_table(res, cols))
    dead = [r for r in res if r.get("live") is False]
    if dead:
        click.echo(
            f"\nWARNING: {len(dead)} parameter(s) failed live resolution "
            "(static catalogue shown for them):", err=True,
        )
        for r in dead:
            click.echo(f"  {r['path']}: {r.get('live_error')}", err=True)


def _echo_og_impact(res: dict) -> None:
    click.echo("OG-UK steady-state reform score")
    click.echo(f"Reform: {res['reform']} (from {res['start_year']})")
    click.echo(f"Assumptions: {res['assumptions']}\n")
    imp = res["impact"]
    rows = []
    for k in ("gdp", "consumption", "investment", "government",
              "tax_revenue", "debt"):
        rows.append({
            "aggregate": k,
            "level (£bn)": imp["levels_bn"][k],
            "change (£bn)": imp["changes_bn"][f"{k}_change"],
            "change (%)": imp["changes_pct"][f"{k}_pct"],
        })
    click.echo(_table(rows, ["aggregate", "level (£bn)", "change (£bn)", "change (%)"]))
    ir = imp["interest_rate"]
    click.echo(f"\nInterest rate: {ir['baseline']} -> {ir['reform']}")


@main.command("og-score")
@click.option("--reform", required=True, help=_REFORM_HELP)
@click.option("--year", default=2026, show_default=True, help="Reform start year.")
@click.option("--max-iter", default=250, show_default=True,
              help="Max solver iterations for each steady-state solve.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def og_score(reform, year, max_iter, as_json):
    """Score a reform with the OG-UK model (alias for `score --model og`; slow: ~10 min)."""
    try:
        res = core.og_score_reform(
            reform=_json_opt(reform, "reform"), start_year=year,
            max_iter=max_iter,
        )
    except ValueError as e:
        raise click.ClickException(str(e)) from e
    if as_json:
        _emit_json(res)
        return
    _echo_og_impact(res)


@main.command("dynamic-score")
@click.option("--reform", required=True, help=_REFORM_HELP)
@click.option("--og-payload", "og_payload_path", default=None,
              type=click.Path(exists=True, dir_okay=False),
              help="Path to a pe-macro og-score --json output produced in a "
                   "separate OG environment (two-env pipeline; see "
                   "PSLmodels/OG-UK#68).")
@click.option("--start-year", "start_year", default=2026, show_default=True,
              help="Reform start year (OG solve and microsim year).")
@click.option("--max-iter", default=250, show_default=True,
              help="Max solver iterations for each OG steady-state solve.")
@click.option("--dataset", default=None,
              help="Microdata dataset name override.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def dynamic_score(reform, og_payload_path, start_year, max_iter, dataset,
                  as_json):
    """Dynamic population score: OG-UK macro overlay on the microsim.

    Alias for `score --model og+microsim` (UK only). With --og-payload it
    consumes a pre-computed OG solve (two-environment pipeline); without,
    it solves OG in-process, which requires an oguk-compatible env.
    """
    og_payload = None
    if og_payload_path is not None:
        try:
            with open(og_payload_path) as f:
                og_payload = json.load(f)
        except (OSError, json.JSONDecodeError) as e:
            raise click.ClickException(
                f"--og-payload {og_payload_path}: not readable JSON ({e}); "
                "pass the unmodified output of `pe-macro og-score --json`"
            ) from e
        if not isinstance(og_payload, dict):
            raise click.ClickException(
                "--og-payload must contain a JSON object (the og-score "
                f"result), got {type(og_payload).__name__}"
            )
    try:
        res = core.dynamic_population_reform_impact(
            country="uk", reform=_json_opt(reform, "reform"),
            year=start_year, max_iter=max_iter, dataset=dataset,
            og_payload=og_payload,
        )
    except (ValueError, ImportError, RuntimeError) as e:
        raise click.ClickException(str(e)) from e
    if as_json:
        _emit_json(res)
        return
    ea = res["economic_assumptions"]
    click.echo("Dynamic score: OG-UK overlay + PolicyEngine microsim")
    click.echo(f"Earnings factor: {ea['earnings_factor']}   "
               f"Labour-supply factor: {ea['labour_supply_factor']}   "
               f"r: {ea['interest_rate_baseline']} -> "
               f"{ea['interest_rate_reform']}\n")
    _echo_score_block(res["score"])
    micro = res["microsim"]
    click.echo(f"\n{micro['headline']}")
    click.echo(_table(micro["decile_impacts"],
                      ["decile", "avg_income_change", "relative_change_pct",
                       "count_better_off", "count_worse_off"]))


@main.command("og-baseline")
@click.option("--year", default=2026, show_default=True, help="Start year.")
@click.option("--max-iter", default=250, show_default=True,
              help="Max solver iterations.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def og_baseline(year, max_iter, as_json):
    """Baseline OG-UK steady state (slow: ~5 min; model units)."""
    res = core.og_baseline(start_year=year, max_iter=max_iter)
    if as_json:
        _emit_json(res)
        return
    click.echo(f"OG-UK baseline steady state, start year {res['start_year']}")
    click.echo(f"Assumptions: {res['assumptions']}\n")
    for k, v in res["steady_state_model_units"].items():
        click.echo(f"  {k:12} {v}")


@main.command("define-scenarios")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def define_scenarios(as_json):
    """List DEFINE-UK climate-policy scenarios (local-only, experimental)."""
    res = core.define_list_scenarios()
    if as_json:
        _emit_json(res)
        return
    if not res.get("available"):
        click.echo(res["how_to_run"])
        return
    click.echo(_table(res["scenarios"], ["name", "block", "description"]))


@main.command("define-scenario")
@click.argument("name")
@click.option("--horizon-years", default=15, show_default=True,
              help="Annual delta horizon.")
@click.option("--incidence", is_flag=True,
              help="Include the incidence variables (YD_HH, P, WR) so the "
                   "output is a valid --define-payload for define-incidence.")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def define_scenario(name, horizon_years, incidence, as_json):
    """DEFINE-UK scenario deltas vs baseline (local-only, experimental).

    Returns annualised DELTAS only — baseline levels are deliberately not
    reported; read the caveats in the output.
    """
    res = core.define_scenario(
        name, horizon_years=horizon_years, incidence=incidence,
    )
    if as_json:
        _emit_json(res)
        return
    if not res.get("available"):
        click.echo(res["how_to_run"])
        if "error" in res:
            click.echo(f"\n{res['error']}")
        return
    click.echo(f"DEFINE-UK scenario deltas vs baseline: {res['scenario']} "
               f"({res['block']})")
    click.echo(res["description"])
    for c in res["caveats"]:
        click.echo(f"  ! {c}")
    years = res["years"]
    for var, paths in res["variables"].items():
        click.echo(f"\n{var} (delta level / delta %):")
        for y, lv, pc in zip(years, paths["delta_level"], paths["delta_pct"]):
            click.echo(f"  {y}  {lv:+.3f}  {pc:+.3f}%")


@main.command("define-incidence")
@click.argument("scenario")
@click.option("--year", default=2030, show_default=True,
              help="Calendar year of the delta path to apply.")
@click.option("--income-concept", default="disposable_income",
              show_default=True,
              type=click.Choice(["disposable_income", "real_wage"]),
              help="Which DEFINE series drives the earnings factor.")
@click.option("--dataset", default=None,
              help="Microdata dataset name override.")
@click.option("--define-payload", "payload_path", default=None,
              type=click.Path(exists=True, dir_okay=False),
              help="Path to a `pe-macro define-scenario <name> --incidence "
                   "--json` output produced where the adapter is installed "
                   "(two-env pipeline: only numbers travel).")
@click.option("--json", "as_json", is_flag=True, help="Emit JSON.")
def define_incidence(scenario, year, income_concept, dataset, payload_path,
                     as_json):
    """Household incidence of a DEFINE-UK scenario (experimental).

    Who bears the climate-policy scenario, household by household: the
    scenario's income deltas scale the microsim's employment-income inputs
    and the change flows through the automatic stabilizers. NOT a reform
    score.
    """
    payload = None
    if payload_path:
        with open(payload_path) as fh:
            try:
                payload = json.load(fh)
            except json.JSONDecodeError as e:
                raise click.ClickException(
                    f"--define-payload is not valid JSON ({e}); pass the "
                    "unmodified output of `pe-macro define-scenario "
                    f"{scenario} --incidence --json`"
                ) from e
    try:
        res = core.define_scenario_incidence(
            scenario, year=year, income_concept=income_concept,
            dataset=dataset, define_payload=payload,
        )
    except ValueError as e:
        raise click.ClickException(str(e)) from e
    if as_json:
        _emit_json(res)
        return
    if not res.get("available", True):
        click.echo(res["how_to_run"])
        return
    click.echo(f"DEFINE-UK scenario incidence: {scenario} ({year}, "
               f"{income_concept})")
    for c in res["caveats"]:
        click.echo(f"  ! {c}")
    app = res["application"]
    click.echo(f"earnings factor: {app['earnings_factor']:.5f} "
               f"(applied: {app['applied']})")
    micro = res["microsim"]
    for k in ("budgetary_impact_bn", "winners_share", "losers_share"):
        if k in micro:
            click.echo(f"  {k:22} {micro[k]}")


if __name__ == "__main__":
    main()