"""
Fast polars-native drop-in replacements for vnpy's slow rolling operators.

This module provides native polars implementations that override vnpy's
rolling_map-based operators through the EXPRESSION_FUNCTIONS extension point.
Optimized for 10M+ row datasets on 2-core NAS infrastructure.

Registration: Call register_fast_ops() once before batch evaluation.
Idempotent and thread-safe.
"""

import sys
from typing import cast

import polars as pl
import numpy as np

# Inject vnpy source path (follow sanguo_factor/universe.py pattern)
_VNPY_SRC = "/Users/chufeng/.openclaw/sanguo_projects/sanguo_vnpy_v2/vnpy_v4.4.0"
if _VNPY_SRC not in sys.path:
    sys.path.insert(0, _VNPY_SRC)

from vnpy.alpha.dataset.utility import DataProxy, EXPRESSION_FUNCTIONS


def fast_ts_rank(feature: DataProxy, window: int) -> DataProxy:
    """Percentile rank of current value within window [0,1] (polars native).

    Replicates: scipy.stats.percentileofscore(s, s[-1]) / 100
    Strategy: Count values <= current, divide by window size.
    Handles NaN by treating NaN > everything (consistent with scipy).
    Returns None for partial windows (like original).
    """
    current = pl.col("data")

    # Count values <= current for each shift k in 0..window-1
    # Fill null with False to treat null comparisons as not <=
    le_count = pl.sum_horizontal([
        pl.when(current.shift(k).over("vt_symbol").le(current).fill_null(False))
        .then(pl.lit(1))
        .otherwise(pl.lit(0))
        for k in range(window)
    ])

    # Rank = count / window
    result = pl.when(pl.int_range(pl.len()).over("vt_symbol") >= window - 1) \
        .then(le_count / pl.lit(window)) \
        .otherwise(None)

    df = feature.df.select(
        pl.col("datetime"),
        pl.col("vt_symbol"),
        result.alias("data")
    )

    return DataProxy(df)


def fast_ts_corr(feature1: DataProxy, feature2: DataProxy, window: int) -> DataProxy:
    """Correlation between two features over rolling window (optimized native).

    Uses correlation identity: corr(x,y) = cov(x,y) / (std_x * std_y)
    """
    df_merged = feature1.df.join(feature2.df, on=["datetime", "vt_symbol"])

    # Use native polars rolling correlation for exact match
    corr_result = pl.rolling_corr(
        pl.col("data"),
        pl.col("data_right"),
        window_size=window,
        min_samples=1
    ).over("vt_symbol")

    df = df_merged.select(
        pl.col("datetime"),
        pl.col("vt_symbol"),
        pl.when(corr_result.is_infinite() | corr_result.is_nan())
        .then(None)
        .otherwise(corr_result)
        .alias("data")
    )

    return DataProxy(df)


def fast_ts_cov(feature1: DataProxy, feature2: DataProxy, window: int) -> DataProxy:
    """Covariance between two features over rolling window (native).

    Uses identity: cov(x,y) = corr(x,y) * std_x * std_y
    """
    # Get correlation and standard deviations, then compute covariance
    corr_result = fast_ts_corr(feature1, feature2, window)

    # Extract std deviations from the merged dataframe
    df_merged = feature1.df.join(feature2.df, on=["datetime", "vt_symbol"])

    std_x = pl.col("data").rolling_std(window, min_samples=1, ddof=0).over("vt_symbol")
    std_y = pl.col("data_right").rolling_std(window, min_samples=1, ddof=0).over("vt_symbol")

    # Covariance = corr * std_x * std_y
    df = df_merged.with_columns([
        std_x.alias("std_x"),
        std_y.alias("std_y"),
        corr_result.df["data"].alias("corr")
    ])

    df = df.select(
        pl.col("datetime"),
        pl.col("vt_symbol"),
        (pl.col("corr") * pl.col("std_x") * pl.col("std_y")).alias("data")
    )

    # Handle infinite/NaN values
    df = df.select(
        pl.col("datetime"),
        pl.col("vt_symbol"),
        pl.when(pl.col("data").is_infinite() | pl.col("data").is_nan())
        .then(None)
        .otherwise(pl.col("data"))
        .alias("data")
    )

    return DataProxy(df)


def fast_ts_decay_linear(feature: DataProxy, window: int) -> DataProxy:
    """Linear decay weighted average: weights (w, w-1, ..., 1) / sum (polars native).

    Optimized by expanding the weighted sum as:
    Σ_k (k+1) * shift(k) for k in 0..window-1
    where shift(0) gets weight 1 (last element), shift(window-1) gets weight w (first element)
    Returns None for partial windows (like original).
    """
    current = pl.col("data")

    # Calculate weighted sum: Σ_k (k+1) * shift(k) for k in 0..window-1
    # shift(k) goes backward in time, so higher k = earlier element = higher weight
    weighted_sum = pl.sum_horizontal([
        pl.lit(k + 1) * current.shift(k).over("vt_symbol")
        for k in range(window)
    ])

    denominator = window * (window + 1) // 2

    result = pl.when(pl.int_range(pl.len()).over("vt_symbol") >= window - 1) \
        .then(weighted_sum / pl.lit(denominator)) \
        .otherwise(None)

    df = feature.df.select(
        pl.col("datetime"),
        pl.col("vt_symbol"),
        result.alias("data")
    )

    return DataProxy(df)


def fast_ts_slope(feature: DataProxy, window: int) -> DataProxy:
    """OLS slope over rolling window: cov(x,t) / var(t) where t = 0..w-1 (native).

    Replicates the optimized original formula but with cleaner native polars syntax.
    x = time index (0, 1, 2, ..., w-1)
    slope = cov(y, x) / var(x)
    """
    n = window
    mean_x = (n - 1) / 2.0
    var_x = (n**2 - 1) / 12.0  # Variance of [0, 1, ..., n-1]

    # E[y*x] using weighted sum
    mean_yx_expr = pl.sum_horizontal([
        i * pl.col("data").shift(window - 1 - i).over("vt_symbol")
        for i in range(n)
    ]) / n

    mean_y = pl.col("data").rolling_mean(window, min_samples=window).over("vt_symbol")

    # cov(y, x) = E[yx] - E[y]E[x]
    cov_yx = mean_yx_expr - mean_y * mean_x

    slope = cov_yx / var_x

    df = feature.df.select(
        pl.col("datetime"),
        pl.col("vt_symbol"),
        slope.alias("data")
    )

    return DataProxy(df)


def fast_ts_rsquare(feature: DataProxy, window: int) -> DataProxy:
    """R-squared of linear regression over rolling window (native).

    r² = slope² * var(x) / var(y)
    where var(x) is constant for window size.
    """
    n = window
    mean_x = (n - 1) / 2.0
    var_x = (n**2 - 1) / 12.0  # Variance of [0, 1, ..., n-1]

    # E[y*x]
    mean_yx_expr = pl.sum_horizontal([
        i * pl.col("data").shift(window - 1 - i).over("vt_symbol")
        for i in range(n)
    ]) / n

    mean_y = pl.col("data").rolling_mean(window, min_samples=window).over("vt_symbol")

    cov_yx = mean_yx_expr - mean_y * mean_x
    var_y = pl.col("data").rolling_var(window, min_samples=window, ddof=0).over("vt_symbol")

    # r² = cov²(x,y) / (var(x) * var(y))
    rsquare = (cov_yx.pow(2)) / (var_x * var_y)

    df = feature.df.select(
        pl.col("datetime"),
        pl.col("vt_symbol"),
        pl.when(rsquare.is_infinite() | rsquare.is_nan())
        .then(None)
        .otherwise(rsquare)
        .alias("data")
    )

    return DataProxy(df)


def fast_ts_resi(feature: DataProxy, window: int) -> DataProxy:
    """Residuals from linear regression over rolling window (native).

    residual = y - (intercept + slope * x_last)
    where x_last = window - 1 (time index of last point)
    """
    n = window
    mean_x = (n - 1) / 2.0
    var_x = (n**2 - 1) / 12.0
    x_last = n - 1

    # E[y*x]
    mean_yx_expr = pl.sum_horizontal([
        i * pl.col("data").shift(window - 1 - i).over("vt_symbol")
        for i in range(n)
    ]) / n

    mean_y = pl.col("data").rolling_mean(window, min_samples=window).over("vt_symbol")

    cov_yx = mean_yx_expr - mean_y * mean_x
    slope = cov_yx / var_x
    intercept = mean_y - slope * mean_x

    # residual = y - (intercept + slope * x_last)
    residual = pl.col("data") - (intercept + slope * x_last)

    df = feature.df.select(
        pl.col("datetime"),
        pl.col("vt_symbol"),
        residual.alias("data")
    )

    return DataProxy(df)


def fast_ts_quantile(feature: DataProxy, window: int, quantile: float) -> DataProxy:
    """Quantile value over rolling window (native polars).

    Uses rolling_quantile with linear interpolation to match original behavior.
    Returns None for partial windows (like original).
    """
    current = pl.col("data")

    # Use native rolling_quantile with linear interpolation
    quantile_result = current.rolling_quantile(
        quantile=quantile,
        interpolation="linear",
        window_size=window
    ).over("vt_symbol")

    # Apply window gating - null for partial windows
    result = pl.when(pl.int_range(pl.len()).over("vt_symbol") >= window - 1) \
        .then(quantile_result) \
        .otherwise(None)

    df = feature.df.select(
        pl.col("datetime"),
        pl.col("vt_symbol"),
        result.alias("data")
    )

    return DataProxy(df)


def register_fast_ops() -> list[str]:
    """Register fast polars operators into vnpy's EXPRESSION_FUNCTIONS.

    Returns list of overridden operator names for verification.
    Idempotent: safe to call multiple times.

    Usage:
        overrides = register_fast_ops()
        print(f"Registered {len(overrides)} fast operators")
    """
    operators = {
        "ts_rank": fast_ts_rank,
        "ts_corr": fast_ts_corr,
        "ts_cov": fast_ts_cov,
        "ts_decay_linear": fast_ts_decay_linear,
        "ts_slope": fast_ts_slope,
        "ts_rsquare": fast_ts_rsquare,
        "ts_resi": fast_ts_resi,
        "ts_quantile": fast_ts_quantile,
    }

    # Register all operators
    for name, func in operators.items():
        EXPRESSION_FUNCTIONS[name] = func

    return list(operators.keys())


# Auto-registration on import is intentional for batch_eval usage
# But also expose explicit registration for testing
__all__ = [
    "register_fast_ops",
    "fast_ts_rank",
    "fast_ts_corr",
    "fast_ts_cov",
    "fast_ts_decay_linear",
    "fast_ts_slope",
    "fast_ts_rsquare",
    "fast_ts_resi",
    "fast_ts_quantile",
]