"""Reference implementation of five daily-OHLC realized volatility estimators.

Conventions (identical to FlashAlpha's /v1/volatility/realized documentation):
  window of n bars, each bar i has O_i, H_i, L_i, C_i, and C_{i-1} is the close before bar i
  o_i = ln(O_i / C_{i-1})   overnight (close-to-open) log return
  c_i = ln(C_i / O_i)       open-to-close log return
  r_i = ln(C_i / C_{i-1}) = o_i + c_i   close-to-close log return
  u_i = ln(H_i / O_i), d_i = ln(L_i / O_i)

  close_to_close   sample variance of r (mean removed, denominator n-1)
  parkinson        sum(ln(H/L)^2) / (4 n ln 2)
  garman_klass     mean(0.5 ln(H/L)^2 - (2 ln 2 - 1) c^2)
  rogers_satchell  mean(u (u - c) + d (d - c))
  yang_zhang       s_o^2 + k s_c^2 + (1 - k) RS, s_o^2 and s_c^2 mean removed with denominator n-1,
                   k = 0.34 / (1.34 + (n + 1) / (n - 1))

Every function returns a DAILY variance. annualize() turns it into percent volatility: 100 * sqrt(252 * var).
A window of n bars needs n + 1 closes (the close before the first bar), for every estimator that uses C_{i-1}.
"""
import math
import numpy as np

LN2 = math.log(2.0)
GK_C = 2.0 * LN2 - 1.0
NAMES = ("close_to_close", "parkinson", "garman_klass", "rogers_satchell", "yang_zhang")
LABELS = {"close_to_close": "Close-to-close", "parkinson": "Parkinson", "garman_klass": "Garman-Klass",
          "rogers_satchell": "Rogers-Satchell", "yang_zhang": "Yang-Zhang"}


def yz_k(n):
    return 0.34 / (1.34 + (n + 1.0) / (n - 1.0))


def components(o, h, l, c, prev_c):
    """Per-bar log terms. Inputs are equal-length arrays; prev_c[i] is the close before bar i."""
    o, h, l, c, prev_c = (np.log(np.asarray(x, dtype=float)) for x in (o, h, l, c, prev_c))
    return components_from_logs(o, h, l, c, prev_c)


def components_from_logs(o, h, l, c, prev_c):
    """Same as components() but the inputs are already natural-log prices."""
    ov = o - prev_c
    oc = c - o
    u = h - o
    d = l - o
    hl = h - l
    return {"o": ov, "c": oc, "r": ov + oc, "u": u, "d": d, "hl2": hl * hl,
            "gk": 0.5 * hl * hl - GK_C * oc * oc, "rs": u * (u - oc) + d * (d - oc)}


def window_variances(comp):
    """Daily variance of each estimator over ONE window (all bars in comp)."""
    n = len(comp["o"])
    if n < 2:
        raise ValueError("need at least 2 bars")
    k = yz_k(n)
    s_o = np.var(comp["o"], ddof=1)
    s_c = np.var(comp["c"], ddof=1)
    rs = float(np.mean(comp["rs"]))
    return {"close_to_close": float(np.var(comp["r"], ddof=1)),
            "parkinson": float(np.sum(comp["hl2"]) / (4.0 * n * LN2)),
            "garman_klass": float(np.mean(comp["gk"])),
            "rogers_satchell": rs,
            "yang_zhang": float(s_o + k * s_c + (1.0 - k) * rs)}


def rolling_variances(comp, n):
    """Rolling daily variance for every estimator; element t uses bars t-n+1..t (NaN before that)."""
    import pandas as pd
    df = pd.DataFrame({key: comp[key] for key in ("o", "c", "r", "hl2", "gk", "rs")})
    roll = df.rolling(n, min_periods=n)
    k = yz_k(n)
    s_o, s_c, m = roll["o"].var(ddof=1), roll["c"].var(ddof=1), roll.mean()
    return pd.DataFrame({"close_to_close": roll["r"].var(ddof=1),
                         "parkinson": m["hl2"] / (4.0 * LN2),
                         "garman_klass": m["gk"],
                         "rogers_satchell": m["rs"],
                         "yang_zhang": s_o + k * s_c + (1.0 - k) * m["rs"]})


def annualize(var):
    return 100.0 * np.sqrt(252.0 * np.asarray(var, dtype=float))


def yz_minus_cc_parts(comp):
    """YZ - CC split into its two exact parts over one window (same bars, same denominators):
    (1 - k)(RS - s_c^2)  and  -2 Cov(o, c)."""
    n = len(comp["o"])
    k = yz_k(n)
    s_c = np.var(comp["c"], ddof=1)
    cov = np.cov(comp["o"], comp["c"], ddof=1)[0, 1]
    return {"range_part": float((1 - k) * (np.mean(comp["rs"]) - s_c)), "cov_part": float(-2 * cov)}