"""
rural_road_cba.py
-----------------
Worked example by Ahmed Beddi: economic cost-benefit analysis of moving an unpaved rural
road network from deferred maintenance to a routine + periodic maintenance regime.

Synthetic, illustrative network. Method in the spirit of the World Bank Roads Economic
Decision model (RED): consumer-surplus benefits to normal traffic, rule-of-half for
generated traffic, roughness-driven vehicle operating costs (VOC) and speeds, and agency
costs in economic prices. Parameters marked "illustrative" in parameters.csv are not
taken from a published source. This is not an appraisal of any real road.

Run:  python rural_road_cba.py        (prints results, writes parameters.csv,
                                       cashflows_base.csv, monte_carlo_summary.csv)

The interactive demo on the website (CbaModel.tsx) implements exactly these formulas
with the same defaults.
"""
from __future__ import annotations

import math
from dataclasses import dataclass, replace, asdict

import numpy as np
import pandas as pd


@dataclass(frozen=True)
class Params:
    # Network and traffic
    length_km: float = 400.0          # network length (km)
    aadt0: float = 120.0              # average daily traffic in year 1 (vehicles/day)
    growth: float = 0.04              # normal traffic growth per year
    elasticity: float = 1.0           # generated traffic: elasticity to transport cost
    # Road-user costs (economic prices, US$)
    voc_a0: float = 0.064             # VOC intercept (US$/veh-km)
    voc_beta: float = 0.013           # VOC slope (US$/veh-km per IRI m/km)
    speed_a: float = 75.0             # speed intercept (km/h)
    speed_b: float = 2.5              # speed loss per IRI point (km/h)
    speed_min: float = 20.0           # speed floor (km/h)
    occupancy: float = 3.0            # persons per vehicle
    vot: float = 1.0                  # value of time (US$/person-hour)
    # Roughness (IRI, m/km)
    iri0: float = 9.0                 # condition in year 1 (both cases)
    iri_good: float = 6.0             # condition after regravelling / reconstruction
    rise_with: float = 0.4            # IRI rise per year under routine maintenance
    rise_without: float = 1.2         # IRI rise per year under deferred maintenance
    iri_trigger: float = 16.0         # without case: reconstruction when IRI reaches this
    # Agency costs (financial prices, US$/km)
    routine: float = 1500.0           # with case: routine maintenance per km per year
    regravel: float = 18000.0         # with case: periodic regravelling per km
    regravel_cycle: int = 5           # years between regravelling (first in year 1)
    spot: float = 500.0               # without case: spot repairs per km per year
    recon: float = 30000.0            # without case: reconstruction per km
    recon_life: int = 10              # years of service assumed for residual value
    admin_share: float = 0.08         # programme administration, share of with-case works
    admin_fixed: float = 200000.0     # programme administration, fixed per year (network)
    cf: float = 0.85                  # financial-to-economic conversion factor (works)
    overrun: float = 0.10             # cost overrun applied to with-case costs
    # Appraisal
    rate: float = 0.06                # economic discount rate
    horizon: int = 20                 # years (year 1 = works year)


def iri_paths(p: Params):
    """Roughness by year for the with and without cases, plus works years."""
    T = p.horizon
    with_iri, without_iri = np.zeros(T), np.zeros(T)
    regravel_years = [t for t in range(1, T + 1) if (t - 1) % p.regravel_cycle == 0]
    recon_years = []
    for t in range(1, T + 1):
        # with case: year 1 is the catch-up works year, still at the initial condition
        if t == 1:
            with_iri[0] = p.iri0
        else:
            last_rg = max(y for y in regravel_years if y < t)
            with_iri[t - 1] = p.iri_good + p.rise_with * (t - last_rg - 1)
        # without case: deterioration; reconstruction in the year the trigger is reached,
        # condition restored from the following year
        if t == 1:
            w = p.iri0
        elif (t - 1) in recon_years:
            w = p.iri_good
        else:
            w = without_iri[t - 2] + p.rise_without
        without_iri[t - 1] = w
        if w >= p.iri_trigger - 1e-9:
            recon_years.append(t)
    return with_iri, without_iri, regravel_years, recon_years


def unit_cost(p: Params, iri):
    """Generalised cost per vehicle-km (US$): VOC + travel time."""
    iri = np.asarray(iri, dtype=float)
    voc = p.voc_a0 + p.voc_beta * iri
    speed = np.maximum(p.speed_min, p.speed_a - p.speed_b * iri)
    time = p.occupancy * p.vot / speed
    return voc, time


def residual_fraction(start_year: int, life: int, horizon: int) -> float:
    """Straight-line share of an asset's life remaining after the horizon."""
    remaining = start_year + life - 1 - horizon
    return max(0.0, min(life, remaining)) / life


def cashflows(p: Params) -> pd.DataFrame:
    T = p.horizon
    t = np.arange(1, T + 1)
    iri_w, iri_wo, rg_years, rc_years = iri_paths(p)
    voc_w, time_w = unit_cost(p, iri_w)
    voc_wo, time_wo = unit_cost(p, iri_wo)
    c_w, c_wo = voc_w + time_w, voc_wo + time_wo

    aadt = p.aadt0 * (1 + p.growth) ** (t - 1)
    vkm = 365.0 * aadt * p.length_km                       # normal veh-km per year
    gen_vkm = p.elasticity * vkm * (c_wo - c_w) / c_wo     # generated veh-km

    b_voc = vkm * (voc_wo - voc_w)
    b_time = vkm * (time_wo - time_w)
    b_gen = 0.5 * gen_vkm * (c_wo - c_w)

    # Agency costs, financial prices
    works_w = np.full(T, p.routine * p.length_km)
    for y in rg_years:
        works_w[y - 1] += p.regravel * p.length_km
    admin = p.admin_share * works_w + p.admin_fixed
    k = p.cf * (1 + p.overrun)
    c_prog = k * works_w                                    # economic
    c_admin = k * admin

    works_wo = np.full(T, p.spot * p.length_km)
    for y in rc_years:
        works_wo[y - 1] += p.recon * p.length_km
    b_avoided = p.cf * works_wo

    # Residual value at the horizon (straight line), with minus without
    rv_w = sum(k * p.regravel * p.length_km * residual_fraction(y, p.regravel_cycle, T) for y in rg_years)
    rv_wo = sum(p.cf * p.recon * p.length_km * residual_fraction(y, p.recon_life, T) for y in rc_years)
    b_resid = np.zeros(T)
    b_resid[-1] = rv_w - rv_wo

    df = pd.DataFrame({
        "year": t, "iri_with": iri_w, "iri_without": iri_wo, "aadt": aadt,
        "voc_saving_per_vkm": voc_wo - voc_w, "time_saving_per_vkm": time_wo - time_w,
        "b_voc": b_voc, "b_time": b_time, "b_generated": b_gen, "b_avoided": b_avoided,
        "b_residual": b_resid, "c_programme": c_prog, "c_admin": c_admin,
    })
    df["benefits"] = df[["b_voc", "b_time", "b_generated", "b_avoided", "b_residual"]].sum(axis=1)
    df["costs"] = df["c_programme"] + df["c_admin"]
    df["net"] = df["benefits"] - df["costs"]
    df["discount_factor"] = (1 + p.rate) ** (-t)
    df["pv_net"] = df["net"] * df["discount_factor"]
    return df


def npv_at(df: pd.DataFrame, rate: float) -> float:
    t = df["year"].to_numpy()
    return float(np.sum(df["net"].to_numpy() / (1 + rate) ** t))


def eirr(df: pd.DataFrame) -> float:
    """Lowest discount rate at which NPV changes sign: scan -20%..100% in 1-point steps,
    then bisect inside the first bracket. NaN if NPV does not change sign in that range."""
    grid = [-0.20 + 0.01 * i for i in range(121)]
    for lo, hi in zip(grid[:-1], grid[1:]):
        f_lo, f_hi = npv_at(df, lo), npv_at(df, hi)
        if f_lo == 0:
            return lo
        if f_lo * f_hi < 0:
            for _ in range(100):
                mid = (lo + hi) / 2
                f_mid = npv_at(df, mid)
                if f_lo * f_mid <= 0:
                    hi = mid
                else:
                    lo, f_lo = mid, f_mid
            return (lo + hi) / 2
    return float("nan")


def summary(p: Params) -> dict:
    df = cashflows(p)
    d = df["discount_factor"]
    pv_b = float((df["benefits"] * d).sum())
    pv_c = float((df["costs"] * d).sum())
    year1_cost = float(df.loc[0, "costs"])
    return {
        "npv": pv_b - pv_c,
        "eirr": eirr(df),
        "bcr": pv_b / pv_c,
        "pv_benefits": pv_b,
        "pv_costs": pv_c,
        "fyrr": float(df.loc[1, "net"]) / year1_cost,   # first full year of service / works-year cost
        "pv_voc": float((df["b_voc"] * d).sum()),
        "pv_time": float((df["b_time"] * d).sum()),
        "pv_generated": float((df["b_generated"] * d).sum()),
        "pv_avoided": float((df["b_avoided"] * d).sum()),
        "pv_residual": float((df["b_residual"] * d).sum()),
        "pv_programme": float((df["c_programme"] * d).sum()),
        "pv_admin": float((df["c_admin"] * d).sum()),
    }


def npv(p: Params) -> float:
    return summary(p)["npv"]


# ---------------------------------------------------------------- sensitivity
# (field, label, low, high) - ranges are those in parameters.csv
SENSITIVITY = [
    ("aadt0", "Traffic, year 1 (veh/day)", 70.0, 200.0),
    ("growth", "Traffic growth (per year)", 0.01, 0.06),
    ("voc_beta", "VOC slope (US$/veh-km per IRI)", 0.009, 0.020),
    ("vot", "Value of time (US$/person-hour)", 0.5, 1.5),
    ("rise_without", "Deterioration without maintenance (IRI/yr)", 0.8, 1.8),
    ("recon", "Reconstruction cost (US$/km)", 20000.0, 40000.0),
    ("regravel", "Regravelling cost (US$/km)", 14000.0, 26000.0),
    ("overrun", "Cost overrun", 0.0, 0.50),
    ("rate", "Discount rate", 0.03, 0.12),
]


def tornado(p: Params) -> pd.DataFrame:
    base = npv(p)
    rows = []
    for field, label, lo, hi in SENSITIVITY:
        n_lo = npv(replace(p, **{field: lo}))
        n_hi = npv(replace(p, **{field: hi}))
        rows.append({"parameter": field, "label": label, "low": lo, "high": hi,
                     "npv_low": n_lo, "npv_high": n_hi, "swing": abs(n_hi - n_lo), "base_npv": base})
    return pd.DataFrame(rows).sort_values("swing", ascending=False).reset_index(drop=True)


def _bisect(f, lo, hi, iters=200):
    f_lo, f_hi = f(lo), f(hi)
    if f_lo * f_hi > 0:
        return float("nan")
    for _ in range(iters):
        mid = (lo + hi) / 2
        f_mid = f(mid)
        if f_lo * f_mid <= 0:
            hi = mid
        else:
            lo, f_lo = mid, f_mid
    return (lo + hi) / 2


# (field, search lower bound, search upper bound)
SWITCHING = [
    ("aadt0", 1.0, 1000.0),
    ("growth", -0.10, 0.20),
    ("voc_beta", 0.0, 0.10),
    ("vot", 0.0, 20.0),
    ("rise_without", 0.0, 1.2),
    ("recon", 0.0, 30000.0),
    ("regravel", 18000.0, 200000.0),
    ("overrun", 0.0, 5.0),
]


def switching_values(p: Params) -> pd.DataFrame:
    rows = []
    for field, lo, hi in SWITCHING:
        sv = _bisect(lambda x: npv(replace(p, **{field: x})), lo, hi)
        base = getattr(p, field)
        rows.append({"parameter": field, "base": base, "switching_value": sv,
                     "change_pct": (sv / base - 1) * 100 if base else float("nan")})
    s = summary(p)
    # Uniform shocks to all benefits / all costs
    rows.append({"parameter": "all_benefits", "base": 1.0, "switching_value": s["pv_costs"] / s["pv_benefits"],
                 "change_pct": (s["pv_costs"] / s["pv_benefits"] - 1) * 100})
    rows.append({"parameter": "all_costs", "base": 1.0, "switching_value": s["pv_benefits"] / s["pv_costs"],
                 "change_pct": (s["pv_benefits"] / s["pv_costs"] - 1) * 100})
    rows.append({"parameter": "rate", "base": p.rate, "switching_value": s["eirr"],
                 "change_pct": (s["eirr"] / p.rate - 1) * 100})
    return pd.DataFrame(rows)


# ---------------------------------------------------------------- Monte Carlo
# Triangular (low, mode, high); discount rate held at its base value.
MC_DISTS = {
    "aadt0": (70.0, 120.0, 200.0),
    "growth": (0.01, 0.04, 0.06),
    "voc_beta": (0.009, 0.013, 0.020),
    "vot": (0.5, 1.0, 1.5),
    "rise_without": (0.8, 1.2, 1.8),
    "recon": (20000.0, 30000.0, 40000.0),
    "regravel": (14000.0, 18000.0, 26000.0),
    "overrun": (0.0, 0.10, 0.50),
}


def monte_carlo(p: Params, n: int = 10000, seed: int = 20261103) -> pd.DataFrame:
    rng = np.random.default_rng(seed)
    draws = {k: rng.triangular(lo, mode, hi, size=n) for k, (lo, mode, hi) in MC_DISTS.items()}
    out = np.empty(n)
    for i in range(n):
        out[i] = npv(replace(p, **{k: float(v[i]) for k, v in draws.items()}))
    df = pd.DataFrame(draws)
    df["npv"] = out
    return df


# ---------------------------------------------------------------- parameter table
PARAMETER_TABLE = [
    # field, label, unit, low, high, status, source, url
    ("length_km", "Network length", "km", None, None, "illustrative", "Synthetic network; results scale linearly with length", ""),
    ("aadt0", "Average daily traffic, year 1", "veh/day", 70, 200, "cited range", "RED designed for low-volume roads of about 50-200 veh/day (SSATP Technical Note 18)", "https://www.ssatp.org/sites/ssatp/files/publication/ATTN18.pdf"),
    ("growth", "Normal traffic growth", "per year", 0.01, 0.06, "illustrative", "", ""),
    ("elasticity", "Generated-traffic elasticity to transport cost", "-", None, None, "cited", "RED user guide case studies use a price elasticity of 1.0 (SSATP Working Paper 78)", "https://documents1.worldbank.org/curated/en/764611468202186160/pdf/349290ENGLISH0ssatpwp78.pdf"),
    ("voc_a0", "VOC intercept", "US$/veh-km", None, None, "derived", "Linear chord (IRI 6-16) of the example car VOC-roughness cubic in the RED user guide", "https://documents1.worldbank.org/curated/en/764611468202186160/pdf/349290ENGLISH0ssatpwp78.pdf"),
    ("voc_beta", "VOC slope per IRI point", "US$/veh-km per IRI", 0.009, 0.020, "derived; range illustrative", "Same chord; applied to an average vehicle; a real appraisal uses HDM-4 VOC by vehicle class at national prices", "https://documents1.worldbank.org/curated/en/764611468202186160/pdf/349290ENGLISH0ssatpwp78.pdf"),
    ("speed_a", "Speed at zero roughness", "km/h", None, None, "illustrative", "Linear speed-roughness approximation", ""),
    ("speed_b", "Speed loss per IRI point", "km/h per IRI", None, None, "illustrative", "", ""),
    ("speed_min", "Speed floor", "km/h", None, None, "illustrative", "", ""),
    ("occupancy", "Occupants per vehicle", "persons", None, None, "cited", "RED user guide case studies: car, 3 occupants", "https://documents1.worldbank.org/curated/en/764611468202186160/pdf/349290ENGLISH0ssatpwp78.pdf"),
    ("vot", "Value of time", "US$/person-hour", 0.5, 1.5, "cited base; range illustrative", "RED case studies: car passengers US$1.0/h, bus passengers US$0.5/h. Non-work time default 0.3 x household income per head (World Bank transport toolkit)", "https://its.leeds.ac.uk/projects/WBToolkit/Note10.htm"),
    ("iri0", "Roughness, year 1", "IRI m/km", None, None, "cited range", "RED case study gravel road at 7-10 IRI; HDM-III gravel roughness 3.5-20 IRI (Transport Note RT-1)", "https://ssatp.org/sites/default/files/publication/InfrastructureNote_rt1.pdf"),
    ("iri_good", "Roughness after works", "IRI m/km", None, None, "illustrative", "Within the RT-1 range", "https://ssatp.org/sites/default/files/publication/InfrastructureNote_rt1.pdf"),
    ("rise_with", "IRI rise per year, maintained", "IRI/yr", None, None, "illustrative", "", ""),
    ("rise_without", "IRI rise per year, deferred", "IRI/yr", 0.8, 1.8, "illustrative", "", ""),
    ("iri_trigger", "Reconstruction trigger", "IRI m/km", None, None, "illustrative", "Within the RT-1 range", "https://ssatp.org/sites/default/files/publication/InfrastructureNote_rt1.pdf"),
    ("routine", "Routine maintenance", "US$/km/yr", None, None, "illustrative", "Order of magnitude: TI-UP enquiry (up to about US$1,000/km/yr, 1998-2004 prices); RED case study all-weather earth road US$2,900/km/yr incl. grading and regravelling", "https://assets.publishing.service.gov.uk/media/57a08bd7e5274a27b2000dc3/TI_UP_HD_Feb2007_Cost_of_Roads_in_Africa.pdf"),
    ("regravel", "Periodic regravelling", "US$/km", 14000, 26000, "illustrative", "Order of magnitude: TI-UP enquiry regravelling US$9,000-13,000/km (2000 prices)", "https://assets.publishing.service.gov.uk/media/57a08bd7e5274a27b2000dc3/TI_UP_HD_Feb2007_Cost_of_Roads_in_Africa.pdf"),
    ("regravel_cycle", "Regravelling cycle", "years", None, None, "illustrative", "RED case study earth road option: regravelling every six years", "https://documents1.worldbank.org/curated/en/764611468202186160/pdf/349290ENGLISH0ssatpwp78.pdf"),
    ("spot", "Spot repairs, deferred regime", "US$/km/yr", None, None, "illustrative", "", ""),
    ("recon", "Reconstruction (rehabilitation) of gravel road", "US$/km", 20000, 40000, "illustrative", "Order of magnitude: TI-UP rough guide US$10,000-15,000/km for gravel roads (older prices)", "https://assets.publishing.service.gov.uk/media/57a08bd7e5274a27b2000dc3/TI_UP_HD_Feb2007_Cost_of_Roads_in_Africa.pdf"),
    ("recon_life", "Service life of reconstruction (residual value)", "years", None, None, "illustrative", "", ""),
    ("admin_share", "Administration, variable", "share of works", None, None, "illustrative", "", ""),
    ("admin_fixed", "Administration, fixed", "US$/yr", None, None, "illustrative", "", ""),
    ("cf", "Conversion factor, financial to economic", "-", None, None, "cited", "RED user guide case studies use an economic/financial cost factor of 0.85", "https://documents1.worldbank.org/curated/en/764611468202186160/pdf/349290ENGLISH0ssatpwp78.pdf"),
    ("overrun", "Cost overrun on programme costs", "share", 0.0, 0.5, "illustrative", "Optimism-bias adjustment in the sense of HM Treasury Green Book para. 2.25", "https://assets.publishing.service.gov.uk/media/6645c709bd01f5ed32793cbc/Green_Book_2022__updated_links_.pdf"),
    ("rate", "Economic discount rate", "per year", 0.03, 0.12, "cited", "World Bank (2016) discounting note: 6% central, 2-10% plausible range; RED case studies use 12%", "https://documents1.worldbank.org/curated/en/099610503022315638/pdf/IDU02e978a2a023f704c140bb1b082595e67e593.pdf"),
    ("horizon", "Evaluation period", "years", None, None, "cited", "RED evaluation period of 2-20 years", "https://documents1.worldbank.org/curated/en/764611468202186160/pdf/349290ENGLISH0ssatpwp78.pdf"),
]


def parameter_table(p: Params = Params()) -> pd.DataFrame:
    vals = asdict(p)
    rows = []
    for field, label, unit, lo, hi, status, source, url in PARAMETER_TABLE:
        rows.append({"parameter": field, "label": label, "unit": unit, "base": vals[field],
                     "low": lo, "high": hi, "status": status, "source": source, "url": url})
    return pd.DataFrame(rows)


def fiscal(p: Params) -> dict:
    """Recurrent budget implications, financial prices (no conversion factor)."""
    ov = 1 + p.overrun
    routine = p.routine * p.length_km * ov
    periodic_avg = p.regravel * p.length_km / p.regravel_cycle * ov
    admin_avg = (p.admin_share * (p.routine + p.regravel / p.regravel_cycle) * p.length_km) * ov + p.admin_fixed * ov
    total = routine + periodic_avg + admin_avg
    return {"routine": routine, "periodic_avg": periodic_avg, "admin_avg": admin_avg,
            "total_per_year": total, "per_km": total / p.length_km}


def main():
    p = Params()
    s = summary(p)
    print("Base case")
    for k, v in s.items():
        print(f"  {k:14s} {v:,.4f}")
    print("\nSwitching values")
    print(switching_values(p).to_string(index=False))
    print("\nTornado")
    print(tornado(p)[["label", "npv_low", "npv_high", "swing"]].to_string(index=False))
    mc = monte_carlo(p)
    print(f"\nMonte Carlo: mean {mc.npv.mean():,.0f}, P5 {mc.npv.quantile(.05):,.0f}, "
          f"P50 {mc.npv.median():,.0f}, P95 {mc.npv.quantile(.95):,.0f}, P(NPV<0) {(mc.npv < 0).mean():.3f}")
    print("\nFiscal", fiscal(p))
    parameter_table(p).to_csv("parameters.csv", index=False)
    cashflows(p).to_csv("cashflows_base.csv", index=False)


if __name__ == "__main__":
    main()
