"""
Replicate Neely, Rapach, Tu and Zhou (2014) from their own data and results files.

    python replicate_neely.py "path/to/Neely2014"

The folder must contain Returns_econ_tech_data.xls and Returns_econ_tech_results.xls
from the authors' replication package. The script

  1. rebuilds the 14 technical indicators from the S&P 500 index and volume in the
     data file, with the rules in their Generate_variables.m, and compares every
     indicator-month with the results file;
  2. re-estimates the bivariate in-sample regressions of Table 2 (next-month log
     equity premium, 1951:01-2011:12, White t-statistics) and compares the slope,
     t-statistic and R-squared with the results file.
"""

import sys
from pathlib import Path

import numpy as np
import pandas as pd

from update_technical import PRICE_RULES as NAMES_PRICE, VOLUME_RULES as NAMES_VOLUME, price_rules, volume_rules


def main(folder: Path) -> None:
    data = pd.read_excel(folder / "Returns_econ_tech_data.xls", sheet_name="Monthly", header=None, skiprows=1)
    data = data[[0, 1, 17]].set_axis(["yyyymm", "P", "vol"], axis=1).dropna(subset=["yyyymm"])
    data["yyyymm"] = data["yyyymm"].astype(int)
    data = data.set_index("yyyymm")
    p = data.loc[192701:201112, "P"].astype(float)
    v = data.loc[195001:201112, "vol"].astype(float)

    mine = pd.concat([price_rules(p), volume_rules(p.loc[195001:], v)], axis=1)

    res = pd.read_excel(folder / "Returns_econ_tech_results.xls", sheet_name="Technical indicators", header=None)
    orig = res.iloc[1:1021, :15].copy()
    orig.columns = ["yyyymm"] + NAMES_PRICE + NAMES_VOLUME
    orig = orig.dropna(subset=["yyyymm"]).astype(float).set_index(orig["yyyymm"].dropna().astype(int))
    orig = orig.drop(columns="yyyymm")
    m = mine.reindex(orig.index)
    has = orig.notna()
    same = (orig == m)
    extra = (m.notna() & orig.isna()).values.sum()
    print("1. Indicators")
    print(f"   indicator-months in the results file: {int(has.values.sum()):,}; identical: {same.values[has.values].mean():.2%}")
    print(f"   ({extra} extra months here are the 9-month rules' first three months, which the original code leaves blank)")
    for c in orig.columns:
        print(f"   {c:9s} from {orig[c].first_valid_index()}: {same.loc[has[c], c].mean():.2%} identical")

    # Table 2: y(t+1) on indicator(t), 1950:12-2011:12, White (HC0) standard errors
    eq = pd.read_excel(folder / "Returns_econ_tech_results.xls", sheet_name="Equity premium", header=None)
    y = 100 * pd.Series(eq.iloc[288:1021, 1].astype(float).values, index=mine.loc[195012:201112].index)
    est = pd.read_excel(folder / "Returns_econ_tech_results.xls", sheet_name="In-sample estimates", header=None)
    paper = est.iloc[33:47, [0, 1, 2, 4]].set_axis(["name", "slope", "t", "r2"], axis=1).set_index("name")
    print("\n2. Table 2, bivariate in-sample regressions, 1951:01-2011:12")
    print(f"   {'':9s} {'slope':>15s} {'t-stat':>15s} {'R2 (%)':>15s}")
    worst = 0.0
    for c in NAMES_PRICE + NAMES_VOLUME:
        x = mine.loc[195012:201112, c].values[:-1]
        yy = y.values[1:]
        X = np.column_stack([np.ones_like(x), x])
        b = np.linalg.lstsq(X, yy, rcond=None)[0]
        e = yy - X @ b
        XtX = np.linalg.inv(X.T @ X)
        V = XtX @ (X.T * e ** 2) @ X @ XtX
        t = b[1] / np.sqrt(V[1, 1])
        r2 = 1 - e @ e / ((yy - yy.mean()) @ (yy - yy.mean()))
        q = paper.loc[c]
        worst = max(worst, abs(b[1] - q.slope), abs(t - q.t), abs(100 * (r2 - q.r2)))
        print(f"   {c:9s} {b[1]:7.3f} ({q.slope:6.3f}) {t:7.2f} ({q.t:6.2f}) {100 * r2:7.2f} ({100 * q.r2:6.2f})")
    print(f"   largest absolute difference from the paper's results file: {worst:.1e}")


if __name__ == "__main__":
    main(Path(sys.argv[1]))