"""STA06 — Kalman Local Level+Slope Trend
Hypothesis: Run a causal Kalman filter on log price with local level + slope states.
The slope state gives a smooth, causal estimate of local trend direction.
Long when filtered slope > 0, flat otherwise (long-only, crypto-style).
Vol-targeted position like TP01.

Grid: 2 observation-noise / process-noise ratio settings × 2 TFs = 4 total cells.
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/alt")
import altlib as al
import numpy as np


def kalman_slope(log_price: np.ndarray, q_level: float = 1e-4, q_slope: float = 1e-6,
                 r_obs: float = 1e-2) -> np.ndarray:
    """
    Causal Kalman local-level + slope filter on log_price.

    State: x = [level, slope]
    Transition: level_{t+1} = level_t + slope_t
                slope_{t+1} = slope_t
    Observation: y_t = level_t + noise

    Parameters:
        q_level: process noise variance for the level
        q_slope: process noise variance for the slope
        r_obs:   observation noise variance

    Returns slope array (same length as log_price), causal at each i.
    """
    n = len(log_price)
    slope_out = np.zeros(n)

    # State transition matrix F
    F = np.array([[1.0, 1.0],
                  [0.0, 1.0]])

    # Process noise covariance Q
    Q = np.array([[q_level, 0.0],
                  [0.0, q_slope]])

    # Observation matrix H (we observe only the level)
    H = np.array([[1.0, 0.0]])

    # Observation noise variance R
    R = np.array([[r_obs]])

    # Initialize state and covariance
    x = np.array([[log_price[0]], [0.0]])   # [level, slope]
    P = np.eye(2) * 1.0

    for i in range(n):
        # --- Predict ---
        x_pred = F @ x
        P_pred = F @ P @ F.T + Q

        # --- Update with observation y[i] ---
        y = np.array([[log_price[i]]])
        S = H @ P_pred @ H.T + R
        K = P_pred @ H.T @ np.linalg.inv(S)
        x = x_pred + K @ (y - H @ x_pred)
        P = (np.eye(2) - K @ H) @ P_pred

        # Record slope (state[1]) at this bar — causal (uses data up to i)
        slope_out[i] = x[1, 0]

    return slope_out


def make_target(q_slope: float):
    """Factory: return a target_fn for a given Kalman noise configuration."""
    def target_fn(df):
        c = df["close"].values.astype(float)
        lp = np.log(c)

        # Kalman filter slope — fully causal recursive
        # q_level scales with q_slope for coherence
        q_level = q_slope * 100.0   # level noise 100x slope noise
        r_obs = 1e-2                # observation noise fixed

        slope = kalman_slope(lp, q_level=q_level, q_slope=q_slope, r_obs=r_obs)

        # Direction: long when slope > 0, flat otherwise
        direction = np.where(slope > 0, 1.0, 0.0)

        # Vol-target the position (TP01 style)
        pos = al.vol_target(direction, df, target_vol=0.20, vol_win_days=30, leverage_cap=2.0)
        return pos

    return target_fn


if __name__ == "__main__":
    # Small grid: 2 q_slope values (controls filter responsiveness)
    # Low q_slope = smoother/slower filter; high q_slope = more responsive
    configs = [
        ("q_slope=1e-6", 1e-6),   # slow, smooth
        ("q_slope=1e-5", 1e-5),   # medium
    ]

    results = []
    for label, q_slope in configs:
        print(f"\n--- Running STA06 config: {label} ---")
        rep = al.study_weights(
            f"STA06-Kalman-{label}",
            make_target(q_slope),
            tfs=("1d", "12h"),
        )
        print(al.fmt(rep))
        print("JSON:", al.as_json(rep))
        results.append((label, q_slope, rep))

    # Pick best config by min_asset_holdout_sharpe across all cells
    best_label, best_q, best_rep = max(
        results,
        key=lambda x: x[2]["verdict"].get("best_holdout_sharpe", -99)
    )
    print(f"\n=== BEST CONFIG: {best_label} ===")
    print(al.fmt(best_rep))
    print("JSON:", al.as_json(best_rep))