"""
Checks the shocks against Ready (2018) over the paper's sample, January 1986 to December 2011.

Rebuilds the decomposition on that sample alone, with the public data used by
update_oilshocks.py, and prints the figures that correspond to the paper's Table III
(summary statistics of the data and the shocks) and Table IV, Panel A (the US market
return regressed on the shocks, White standard errors). The paper's values are printed
beside them. Needs update_oilshocks.py in the same folder.

Run: python replicate_ready.py
"""

import numpy as np
import pandas as pd

import update_oilshocks as U

PAPER = {  # Ready (2018), Tables III and IV, Panel A
    "sd": {"dp": 0.34, "s": 0.30, "d": 0.15, "v": 0.04},
    "share": {"s": 0.78, "d": 0.21, "v": 0.01},
    "corr_dp": {"s": 0.88, "d": 0.45, "v": 0.11},
    "beta": {"d": (0.370, 0.046), "s": (-0.102, 0.021), "xi_vix": (-0.184, 0.012)},
    "r2": 0.604, "r2_uni": {"d": 0.124, "s": 0.036, "xi_vix": 0.444},
}


def ols(y, X):
    X = np.column_stack([np.ones(len(y)), X])
    b = np.linalg.lstsq(X, y, rcond=None)[0]
    e = y - X @ b
    XtXi = np.linalg.inv(X.T @ X)
    V = XtXi @ (X.T * e ** 2) @ X @ XtXi  # White (1980)
    r2 = 1 - e.var() / y.var()
    return b[1:], np.sqrt(np.diag(V))[1:], r2


def main() -> None:
    vix = U.fred_month_end("VIXCLS")
    vxo = U.fred_month_end("VXOCLS")
    wti = U.fred_month_end("DCOILWTICO")
    prod = U.french_monthly("49_Industry_Portfolios_CSV.zip", "Oil")
    mkt = U.french_monthly("F-F_Research_Data_Factors_CSV.zip", "Mkt-RF") + U.french_monthly("F-F_Research_Data_Factors_CSV.zip", "RF")
    vol = pd.concat([vxo[vxo.index < U.VIX_FROM], vix[vix.index >= U.VIX_FROM]])
    d = pd.DataFrame({"dp": np.log(wti).diff(), "r_prod": np.log1p(prod), "r_mkt": np.log1p(mkt),
                      "lvix": np.log(vol)}).loc["1986-01":"2011-12"]
    d["xi_vix"] = U.arma11_residuals(d["lvix"].values)
    d = d.iloc[1:]
    d[["s", "d", "v"]] = U.decompose(d[["dp", "r_prod", "xi_vix"]].values)
    ann = 12 ** 0.5

    print(f"Sample {d.index[0]} to {d.index[-1]}, {len(d)} months (paper: 1986-2011)\n")
    print("Table III, Panel B                       here    paper")
    for c in ["dp", "s", "d", "v"]:
        print(f"  annualised sd of {c:<20} {d[c].std() * ann:6.2f}   {PAPER['sd'][c]:6.2f}")
    for c in ["s", "d", "v"]:
        print(f"  share of variance, {c:<18} {d[c].var() / d['dp'].var():6.2f}   {PAPER['share'][c]:6.2f}")
    for c in ["s", "d", "v"]:
        print(f"  correlation with dp, {c:<16} {d[c].corr(d['dp']):6.2f}   {PAPER['corr_dp'][c]:6.2f}")
    print("\nTable IV, Panel A: US market return on d, s and the VIX innovation")
    b, se, r2 = ols(d["r_mkt"].values, d[["d", "s", "xi_vix"]].values)
    for k, name in enumerate(["d", "s", "xi_vix"]):
        pb, pse = PAPER["beta"][name]
        print(f"  {name:<8} {b[k]:7.3f} ({se[k]:.3f})   paper {pb:7.3f} ({pse:.3f})")
    print(f"  R-squared {r2:.3f}   paper {PAPER['r2']:.3f}")
    for name in ["d", "s", "xi_vix"]:
        _, _, r2u = ols(d["r_mkt"].values, d[[name]].values)
        print(f"  univariate R-squared, {name:<8} {r2u:.3f}   paper {PAPER['r2_uni'][name]:.3f}")
    b0, se0, r20 = ols(d["r_mkt"].values, d[["dp"]].values)
    print(f"  on the raw oil price change: {b0[0]:.3f} ({se0[0]:.3f}), R-squared {r20:.3f}   paper 0.031 (0.027), 0.004")


if __name__ == "__main__":
    main()