"""
Oil supply, demand and risk shocks of Ready (2018), monthly, 1986 to date.

Ready, R. C. (2018). Oil prices and the stock market. Review of Finance 22(1), 155-176,
splits each month's change in the oil price into three orthogonal parts using asset prices:

  risk shock    v   the part explained by the unexpected change in the VIX
  demand shock  d   the part explained by the return on oil producers' shares, once the
                    VIX innovation is taken out (producers gain when demand rises but are
                    roughly hedged against supply shocks)
  supply shock  s   the rest

With X = (oil price change, producer return, VIX innovation)' and Z = (s, d, v)', the
paper sets X = A Z with A = [1 1 1; 0 a22 a23; 0 0 a33] and chooses A so that the three
shocks are uncorrelated (Section 3.1, equations 13 and 14). The shocks add up to the
oil price change exactly: s + d + v = oil price change in every month.

Data (all free, downloaded when the script runs)
  Oil price          WTI spot, Cushing, month-end, dollars a barrel (FRED DCOILWTICO; the
                     paper used the second-nearest NYMEX contract, which the EIA stopped
                     publishing in April 2024; over 1986-2024 the two give supply shocks
                     with a correlation of 0.97 and identical demand and risk shocks)
  Producer return    Fama-French 49-industry "Oil" portfolio, value-weighted, from Kenneth
                     French's data library (the paper used the Datastream World Integrated
                     Oil & Gas Producers index)
  VIX                CBOE VIX, month-end close (FRED VIXCLS), with the older VXO (FRED
                     VXOCLS) for 1986-1989, the years before the VIX itself starts
  Recessions         NBER recession months (FRED USREC), for the chart only
Returns and price changes are logs; the VIX innovation is the residual of an ARMA(1,1)
fitted to the log of the VIX over the whole sample, as in the paper.

Output (next to this script): oil_shocks_monthly.csv, one row per month
  month, oil_price, dp (log change in the oil price), xi_vix (VIX innovation),
  s, d, v (the shocks; s + d + v = dp), dp_cum, s_cum, d_cum, v_cum (cumulated from
  January 1986), s_realtime, d_realtime, v_realtime (the same shocks when the ARMA and
  the covariance matrix are estimated only on data up to each month, from a first window
  of 120 months), recession.
The full-sample shocks are re-estimated from scratch every month, so earlier values move
a little as the sample grows; the realtime columns never change once published.

Cite: Ready, R. C. (2018). Oil prices and the stock market. Review of Finance 22(1),
155-176. https://doi.org/10.1093/rof/rfw071

Requires: Python 3.9+, numpy, pandas, scipy.   Run: python update_oilshocks.py
"""

import datetime as dt
import io
import re
import zipfile
from pathlib import Path
from urllib.request import Request, urlopen

import numpy as np
import pandas as pd
from scipy.optimize import minimize

HERE = Path(__file__).resolve().parent
BROWSER = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                         "(KHTML, like Gecko) Chrome/126 Safari/537.36"}
START = pd.Period("1986-01", freq="M")   # first month with a VIX-type index
VIX_FROM = pd.Period("1990-01", freq="M")  # VIX from here, VXO before
FIRST_WINDOW = 120                        # months in the first realtime estimate
FRENCH = "https://mba.tuck.dartmouth.edu/pages/faculty/ken.french/ftp/"


# ---------------------------------------------------------------- data

def fetch(url: str) -> bytes:
    with urlopen(Request(url, headers=BROWSER), timeout=120) as resp:
        return resp.read()


def fred_month_end(series: str) -> pd.Series:
    """Last available daily value in each month of a FRED series."""
    raw = pd.read_csv(io.BytesIO(fetch(f"https://fred.stlouisfed.org/graph/fredgraph.csv?id={series}")),
                      na_values=["."])
    raw.columns = ["date", "value"]
    raw["date"] = pd.to_datetime(raw["date"])
    raw = raw.dropna()
    s = raw.groupby(raw["date"].dt.to_period("M"))["value"].last()
    return s.rename(series)


def fred_monthly(series: str) -> pd.Series:
    """A FRED series that is already monthly (dated the first of the month)."""
    raw = pd.read_csv(io.BytesIO(fetch(f"https://fred.stlouisfed.org/graph/fredgraph.csv?id={series}")),
                      na_values=["."])
    raw.columns = ["date", "value"]
    raw.index = pd.to_datetime(raw["date"]).dt.to_period("M")
    return raw["value"].rename(series)


def french_monthly(zip_name: str, column: str) -> pd.Series:
    """One column of the first monthly block (value-weighted returns, in percent) of a
    Kenneth French CSV file, as a decimal return."""
    with zipfile.ZipFile(io.BytesIO(fetch(FRENCH + zip_name))) as z:
        text = z.read(z.namelist()[0]).decode("latin-1")
    lines = text.splitlines()
    # The first block starts at the first header row (",Agric,Food ,..." or ",Mkt-RF,...").
    start = next(i for i, ln in enumerate(lines) if ln.startswith(",") and column in [c.strip() for c in ln.split(",")])
    header = [c.strip() for c in lines[start].split(",")]
    rows = {}
    for ln in lines[start + 1:]:
        if not re.match(r"^\s*\d{6},", ln):
            break
        cells = [c.strip() for c in ln.split(",")]
        rows[pd.Period(f"{cells[0][:4]}-{cells[0][4:]}", freq="M")] = float(cells[header.index(column)])
    s = pd.Series(rows, name=column)
    s[s <= -99] = np.nan  # -99.99 marks missing
    return s / 100


# ---------------------------------------------------------------- method

def arma11_residuals(y: np.ndarray) -> np.ndarray:
    """Residuals of y_t = c + phi y_{t-1} + theta e_{t-1} + e_t, estimated by conditional
    least squares (e_1 = 0). Same algorithm as the MATLAB script."""
    def resid(p):
        c, phi, theta = p
        e = np.zeros(len(y))
        for t in range(1, len(y)):
            e[t] = y[t] - c - phi * y[t - 1] - theta * e[t - 1]
        return e

    def ssr(p):
        if abs(p[2]) >= 1:
            return 1e12
        return float(np.sum(resid(p)[1:] ** 2))

    x0 = np.array([y.mean() * 0.1, 0.9, 0.0])
    best = minimize(ssr, x0, method="Nelder-Mead", options={"xatol": 1e-8, "fatol": 1e-10, "maxiter": 20000})
    best = minimize(ssr, best.x, method="Nelder-Mead", options={"xatol": 1e-9, "fatol": 1e-12, "maxiter": 20000})
    return resid(best.x)


def decompose(X: np.ndarray) -> np.ndarray:
    """Columns of X: oil price change, producer return, VIX innovation. Returns the shocks
    Z = inv(A) X with A = [1 1 1; 0 a22 a23; 0 0 a33] chosen so that inv(A) cov(X) inv(A)'
    is diagonal (Ready 2018, equations 13-14). The producer return is demeaned first, so
    the demand and risk shocks average zero and the supply shock, the remainder, carries
    the average drift in the oil price; the covariances, and so A, are unaffected."""
    S = np.cov(X.T)
    a33 = S[2, 2] / S[0, 2]
    a23 = S[1, 2] / S[0, 2]
    c = S[0, 1] - S[1, 2] * S[0, 2] / S[2, 2]
    e = S[1, 1] - S[1, 2] ** 2 / S[2, 2]
    a22 = e / c
    A = np.array([[1.0, 1.0, 1.0], [0.0, a22, a23], [0.0, 0.0, a33]])
    Xc = X.copy()
    Xc[:, 1] = Xc[:, 1] - Xc[:, 1].mean()
    return np.linalg.solve(A, Xc.T).T


# ---------------------------------------------------------------- build

def build() -> pd.DataFrame:
    vix = fred_month_end("VIXCLS")
    vxo = fred_month_end("VXOCLS")
    wti = fred_month_end("DCOILWTICO")
    rec = fred_monthly("USREC")
    prod = french_monthly("49_Industry_Portfolios_CSV.zip", "Oil")

    vol = pd.concat([vxo[vxo.index < VIX_FROM], vix[vix.index >= VIX_FROM]])
    d = pd.DataFrame({
        "oil_price": wti,
        "dp": np.log(wti).diff(),
        "r_prod": np.log1p(prod),
        "lvix": np.log(vol),
    }).loc[START:]
    d = d.loc[:d[["dp", "r_prod", "lvix"]].dropna().index[-1]]  # last month with every input
    gaps = d.index[d[["dp", "r_prod", "lvix"]].isna().any(axis=1)][1:]  # the first month has no price change
    if len(gaps):
        raise RuntimeError(f"Missing inputs in {list(map(str, gaps))}")

    d["xi_vix"] = arma11_residuals(d["lvix"].values)
    use = d.index[1:]  # the first month has no price change and no innovation
    X = d.loc[use, ["dp", "r_prod", "xi_vix"]].values
    d.loc[use, ["s", "d", "v"]] = decompose(X)
    for c in ["dp", "s", "d", "v"]:
        d[c + "_cum"] = d[c].fillna(0).cumsum()

    # Realtime: everything re-estimated on data up to each month.
    for c in ["s", "d", "v"]:
        d[c + "_realtime"] = np.nan
    for k in range(FIRST_WINDOW, len(d) + 1):
        sub = d.iloc[:k]
        xi = arma11_residuals(sub["lvix"].values)
        Xk = np.column_stack([sub["dp"].values[1:], sub["r_prod"].values[1:], xi[1:]])
        d.loc[sub.index[-1], ["s_realtime", "d_realtime", "v_realtime"]] = decompose(Xk)[-1]

    d["recession"] = rec.reindex(d.index).astype("Int64")
    d.loc[d.index[0], "xi_vix"] = np.nan
    out = d[["oil_price", "dp", "xi_vix", "s", "d", "v", "dp_cum", "s_cum", "d_cum", "v_cum",
             "s_realtime", "d_realtime", "v_realtime", "recession"]].copy()
    out.index = out.index.astype(str)
    out.index.name = "month"
    return out


def main() -> None:
    out = build()
    out.to_csv(HERE / "oil_shocks_monthly.csv", float_format="%.6g")
    last = out.dropna(subset=["s"]).iloc[-1]
    share = out[["s", "d", "v"]].var() / out["dp"].var()
    print(f"Wrote oil_shocks_monthly.csv ({len(out)} months, {out.index[0]} to {out.index[-1]})")
    print(f"Share of the variance of oil price changes: supply {share['s']:.0%}, demand {share['d']:.0%}, risk {share['v']:.0%}")
    print(f"{out.index[-1]}: oil price change {last['dp']:+.1%} = supply {last['s']:+.1%} + demand {last['d']:+.1%} + risk {last['v']:+.1%}")


if __name__ == "__main__":
    main()
