"""
Technical indicator calculations for stock OHLCV DataFrames.
All functions operate in-place and return the enriched DataFrame.
"""

import logging
import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


def add_moving_averages(df: pd.DataFrame) -> pd.DataFrame:
    """Add MA5/10/20/60/120/240 columns based on closing price."""
    for period in (5, 10, 20, 60, 120, 240):
        df[f"ma{period}"] = df["close"].rolling(window=period, min_periods=period).mean()
    return df


def add_ma_cross_signals(df: pd.DataFrame) -> pd.DataFrame:
    """
    Add golden/death cross boolean features for MA5xMA20 and MA20xMA60.
    Golden cross = short MA crosses above long MA (1 on crossover day, else 0).
    Death cross = short MA crosses below long MA.
    """
    for short, long in ((5, 20), (20, 60)):
        ma_s = df.get(f"ma{short}")
        ma_l = df.get(f"ma{long}")
        if ma_s is None or ma_l is None:
            df[f"golden_cross_{short}_{long}"] = np.nan
            df[f"death_cross_{short}_{long}"] = np.nan
            continue
        above = (ma_s > ma_l).fillna(False).astype(bool)
        prev_above = above.shift(1).fillna(False).astype(bool)
        df[f"golden_cross_{short}_{long}"] = (above & ~prev_above).astype(float)
        df[f"death_cross_{short}_{long}"] = (~above & prev_above).astype(float)
    return df


def add_ma_bull_alignment(df: pd.DataFrame) -> pd.DataFrame:
    mas = [df.get(f"ma{p}") for p in (5, 10, 20, 60)]
    if any(m is None for m in mas):
        df["ma_bull_alignment"] = 0.0
        df["ma_bear_alignment"] = 0.0
        df["ma_alignment_days"] = 0.0
        return df
    ma5, ma10, ma20, ma60 = mas
    bull = ((ma5 > ma10) & (ma10 > ma20) & (ma20 > ma60)).astype(float)
    bear = ((ma5 < ma10) & (ma10 < ma20) & (ma20 < ma60)).astype(float)
    df["ma_bull_alignment"] = bull
    df["ma_bear_alignment"] = bear
    state = bull - bear  # 1, -1, or 0
    group = state.ne(state.shift()).cumsum()
    days = state.groupby(group).cumcount().add(1).astype(float)
    df["ma_alignment_days"] = days * state
    return df


def add_bias_rates(df: pd.DataFrame) -> pd.DataFrame:
    close = df.get("close", pd.Series(dtype=float))
    for period in (20, 60):
        ma = df.get(f"ma{period}")
        if ma is not None:
            df[f"bias_{period}"] = ((close - ma) / ma.replace(0, np.nan)).fillna(0.0)
    return df


def add_rsi(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
    """
    Add RSI (Relative Strength Index) column.
    Uses Wilder's smoothing (EWM with alpha = 1/period).
    """
    delta = df["close"].diff()
    gain = delta.clip(lower=0)
    loss = -delta.clip(upper=0)

    # Wilder's smoothing: adjust=False, com=period-1  →  alpha = 1/(period)
    avg_gain = gain.ewm(com=period - 1, adjust=False, min_periods=period).mean()
    avg_loss = loss.ewm(com=period - 1, adjust=False, min_periods=period).mean()

    rs = avg_gain / avg_loss.replace(0, np.nan)
    df["rsi"] = 100 - (100 / (1 + rs))
    return df


def add_macd(df: pd.DataFrame) -> pd.DataFrame:
    """
    Add MACD line, signal line (9-EMA of MACD), and histogram.
    Standard parameters: fast=12, slow=26, signal=9.
    """
    ema12 = df["close"].ewm(span=12, adjust=False).mean()
    ema26 = df["close"].ewm(span=26, adjust=False).mean()
    df["macd"] = ema12 - ema26
    df["macd_signal"] = df["macd"].ewm(span=9, adjust=False).mean()
    df["macd_hist"] = df["macd"] - df["macd_signal"]
    return df


def add_bollinger_bands(
    df: pd.DataFrame, period: int = 20, std: float = 2.0
) -> pd.DataFrame:
    """
    Add Bollinger Bands: middle (SMA), upper (SMA + 2σ), lower (SMA - 2σ).
    Also computes %B = (close - lower) / (upper - lower).
    """
    sma = df["close"].rolling(window=period, min_periods=period).mean()
    rolling_std = df["close"].rolling(window=period, min_periods=period).std()
    df["bb_middle"] = sma
    df["bb_upper"] = sma + std * rolling_std
    df["bb_lower"] = sma - std * rolling_std
    band_width = df["bb_upper"] - df["bb_lower"]
    df["bb_pct_b"] = (df["close"] - df["bb_lower"]) / band_width.replace(0, np.nan)
    return df


def add_kd(df: pd.DataFrame, period: int = 9) -> pd.DataFrame:
    """
    Add Stochastic Oscillator K and D lines.
    K = 3-period SMA of raw %K (popular Taiwan variant uses 3,3 smoothing).
    D = 3-period SMA of K.
    Raw %K = (close - lowest_low(period)) / (highest_high(period) - lowest_low(period)) * 100
    """
    low_min = df["low"].rolling(window=period, min_periods=period).min()
    high_max = df["high"].rolling(window=period, min_periods=period).max()
    range_ = (high_max - low_min).replace(0, np.nan)
    raw_k = (df["close"] - low_min) / range_ * 100

    # Smooth K and D with 3-period SMA (Taiwan convention)
    df["k"] = raw_k.rolling(window=3, min_periods=1).mean()
    df["d"] = df["k"].rolling(window=3, min_periods=1).mean()
    return df


def add_volume_indicators(df: pd.DataFrame) -> pd.DataFrame:
    """
    Add VolumeMA20 and volume_ratio (volume / VolumeMA20).
    """
    df["volume_ma20"] = df["volume"].rolling(window=20, min_periods=20).mean()
    df["volume_ratio"] = df["volume"] / df["volume_ma20"].replace(0, np.nan)
    return df


def add_atr(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
    """
    Add ATR (Average True Range) and atr_ratio (ATR / close).
    True Range = max(H-L, |H-Prev_Close|, |L-Prev_Close|) — accounts for gaps.
    atr_ratio normalises ATR to price level.
    """
    prev_close = df["close"].shift(1)
    tr1 = df["high"] - df["low"]
    tr2 = (df["high"] - prev_close).abs()
    tr3 = (df["low"] - prev_close).abs()
    tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
    df["atr"] = tr.rolling(window=period, min_periods=period).mean()
    df["atr_ratio"] = df["atr"] / df["close"].replace(0, np.nan)
    return df


def add_obv(df: pd.DataFrame) -> pd.DataFrame:
    """
    Add OBV (On-Balance Volume) and obv_trend (5-day normalised slope).
    OBV = cumulative sum of volume * sign(close change).
    obv_trend = 5-day simple diff of OBV, divided by mean OBV to normalise.
    """
    close_change = df["close"].diff()
    direction = np.sign(close_change).fillna(0)
    df["obv"] = (df["volume"] * direction).cumsum()

    # 5-day slope: simple difference normalised by rolling mean absolute OBV
    obv_diff = df["obv"].diff(5)
    obv_scale = df["obv"].abs().rolling(window=20, min_periods=5).mean().replace(0, np.nan)
    df["obv_trend"] = obv_diff / obv_scale
    return df


def add_volatility_regime(df: pd.DataFrame, period: int = 20) -> pd.DataFrame:
    """Add rolling return volatility as a market regime indicator."""
    returns = df["close"].pct_change()
    df["volatility_20d"] = returns.rolling(window=period, min_periods=period).std()
    return df


def add_cross_asset_tw(
    df: pd.DataFrame,
    taiex_close: "pd.Series | None",
    usdtwd_close: "pd.Series | None",
    sox_close: "pd.Series | None" = None,
    tnx_close: "pd.Series | None" = None,
) -> pd.DataFrame:
    """
    Add Taiwan cross-asset features to df:
      - taiex_return_5d  : TAIEX 5-day return
      - taiex_ma20_ratio : TAIEX close / 20-day MA
      - usdtwd_return_5d : USD/TWD 5-day change
      - sox_ret_1d       : SOX prior-day return (non-leaking; US closes before TW opens)
      - sox_ret_5d       : SOX 5-day return
      - sox_ma20_ratio   : SOX close / SOX 20-day MA
      - tnx_level        : US 10Y yield level (%)
      - tnx_change_5d    : 5-day change in yield

    All Series are indexed by 'YYYY-MM-DD' strings from fetch_cross_asset_tw().
    Pass None for any series to get NaN columns (neutral-imputed later).
    """
    date_col = df["date"].astype(str) if "date" in df.columns else None
    nan_col = pd.Series(np.nan, index=df.index)

    if taiex_close is not None and date_col is not None:
        taiex_dict = taiex_close.to_dict()
        aligned = date_col.map(taiex_dict).astype(float)
        df["taiex_return_5d"]  = aligned.pct_change(5)
        df["taiex_return_20d"] = aligned.pct_change(20)
        taiex_ma20  = aligned.rolling(20,  min_periods=20).mean()
        taiex_ma200 = aligned.rolling(200, min_periods=100).mean()
        df["taiex_ma20_ratio"]  = aligned / taiex_ma20.replace(0, np.nan)
        df["taiex_ma200_ratio"] = aligned / taiex_ma200.replace(0, np.nan)
    else:
        df["taiex_return_5d"]   = nan_col
        df["taiex_return_20d"]  = nan_col
        df["taiex_ma20_ratio"]  = nan_col
        df["taiex_ma200_ratio"] = nan_col

    if usdtwd_close is not None and date_col is not None:
        fx_dict = usdtwd_close.to_dict()
        aligned_fx = date_col.map(fx_dict).astype(float)
        df["usdtwd_return_5d"] = aligned_fx.pct_change(5)
    else:
        df["usdtwd_return_5d"] = nan_col

    if sox_close is not None and date_col is not None:
        sox_dict = sox_close.to_dict()
        aligned_sox = date_col.map(sox_dict).astype(float)
        df["sox_ret_1d"] = aligned_sox.pct_change(1).shift(1)  # prior-day return, non-leaking
        df["sox_ret_5d"] = aligned_sox.pct_change(5)
        sox_ma20 = aligned_sox.rolling(20, min_periods=20).mean()
        df["sox_ma20_ratio"] = aligned_sox / sox_ma20.replace(0, np.nan)
    else:
        df["sox_ret_1d"] = nan_col
        df["sox_ret_5d"] = nan_col
        df["sox_ma20_ratio"] = nan_col

    if tnx_close is not None and date_col is not None:
        tnx_dict = tnx_close.to_dict()
        aligned_tnx = date_col.map(tnx_dict).astype(float)
        df["tnx_level"] = aligned_tnx
        df["tnx_change_5d"] = aligned_tnx.diff(5)
    else:
        df["tnx_level"] = nan_col
        df["tnx_change_5d"] = nan_col

    return df


def add_amihud_illiquidity(df: pd.DataFrame) -> pd.DataFrame:
    close = df.get("close", pd.Series(dtype=float))
    volume = df.get("volume", pd.Series(dtype=float))
    ret = close.pct_change().abs()
    dollar_vol = (close * volume).replace(0, np.nan)
    # Raw Amihud: |return| / dollar_volume (in billions to normalize scale)
    amihud_raw = (ret / (dollar_vol / 1e9)).fillna(0.0)
    # Rolling 20-day average — stable estimate
    df["amihud_illiquidity"] = amihud_raw.rolling(20, min_periods=5).mean().fillna(0.0)
    # Z-score vs 60-day rolling mean/std — relative illiquidity
    roll_mean = df["amihud_illiquidity"].rolling(60, min_periods=20).mean()
    roll_std  = df["amihud_illiquidity"].rolling(60, min_periods=20).std().replace(0, np.nan)
    df["amihud_zscore"] = ((df["amihud_illiquidity"] - roll_mean) / roll_std).fillna(0.0).clip(-3, 3)
    return df


def add_fracdiff_features(df: pd.DataFrame) -> pd.DataFrame:
    """
    Add fractionally differenced log(close) features at d=0.35 and d=0.40.
    Standard AFML sign convention: w_0=1, w_k = -w_{k-1}*(d-k+1)/k
    (d=1 recovers first difference). Applied to log(close) for scale-invariance.
    Output is z-scored with a 60-day rolling window to remove price-level trend
    and make the feature stationary and cross-stock comparable.
    window=30 lags; first window-1 rows padded with 0.0.
    Pure numpy — no external fracdiff package required.
    """
    try:
        log_close = np.log(df["close"].values.astype(float))
        n = len(log_close)
        window = 30
        zscore_window = 60
        for d, col in [(0.35, "fracdiff_close_35"), (0.40, "fracdiff_close_40")]:
            weights = np.empty(window)
            weights[0] = 1.0
            for k in range(1, window):
                weights[k] = -weights[k - 1] * (d - k + 1) / k
            full = np.convolve(log_close, weights)
            result = full[0:n].copy()
            result[:window - 1] = 0.0
            # Rolling z-score to remove price-level trend
            series = pd.Series(result)
            roll_mean = series.rolling(zscore_window, min_periods=zscore_window).mean()
            roll_std = series.rolling(zscore_window, min_periods=zscore_window).std().replace(0, np.nan)
            zscored = ((series - roll_mean) / roll_std).fillna(0.0).clip(-3, 3).values
            df[col] = zscored
    except Exception:
        df["fracdiff_close_35"] = 0.0
        df["fracdiff_close_40"] = 0.0
    return df


def add_calendar_features(df: pd.DataFrame) -> pd.DataFrame:
    dates = pd.to_datetime(df["date"])
    dow   = dates.dt.dayofweek   # 0=Mon … 4=Fri
    month = dates.dt.month
    dom   = dates.dt.day

    df["dow_sin"]                = np.sin(2 * np.pi * dow / 5)
    df["dow_cos"]                = np.cos(2 * np.pi * dow / 5)
    df["month_sin"]              = np.sin(2 * np.pi * month / 12)
    df["month_cos"]              = np.cos(2 * np.pi * month / 12)
    df["is_options_expiry_week"] = ((dom >= 15) & (dom <= 21)).astype(int)
    df["is_earnings_season"]     = month.isin([3, 4, 8, 9]).astype(int)
    return df


def add_all_indicators(df: pd.DataFrame) -> pd.DataFrame:
    """
    Apply all technical indicators to the DataFrame.
    Operates on a copy to avoid mutating the caller's data.
    Returns the enriched DataFrame (NaN rows are NOT dropped here —
    let the predictor handle that).
    """
    df = df.copy()
    try:
        df = add_moving_averages(df)
    except Exception as exc:
        logger.warning("add_moving_averages failed: %s", exc)
    try:
        df = add_ma_cross_signals(df)
    except Exception as exc:
        logger.warning("add_ma_cross_signals failed: %s", exc)
    try:
        df = add_ma_bull_alignment(df)
    except Exception as exc:
        logger.warning("add_ma_bull_alignment failed: %s", exc)
    try:
        df = add_bias_rates(df)
    except Exception as exc:
        logger.warning("add_bias_rates failed: %s", exc)
    try:
        df = add_rsi(df)
    except Exception as exc:
        logger.warning("add_rsi failed: %s", exc)
    try:
        df = add_macd(df)
    except Exception as exc:
        logger.warning("add_macd failed: %s", exc)
    try:
        df = add_bollinger_bands(df)
    except Exception as exc:
        logger.warning("add_bollinger_bands failed: %s", exc)
    try:
        df = add_kd(df)
    except Exception as exc:
        logger.warning("add_kd failed: %s", exc)
    try:
        df = add_volume_indicators(df)
    except Exception as exc:
        logger.warning("add_volume_indicators failed: %s", exc)
    try:
        df = add_atr(df)
    except Exception as exc:
        logger.warning("add_atr failed: %s", exc)
    try:
        df = add_obv(df)
    except Exception as exc:
        logger.warning("add_obv failed: %s", exc)
    try:
        df = add_volatility_regime(df)
    except Exception as exc:
        logger.warning("add_volatility_regime failed: %s", exc)
    try:
        df = add_amihud_illiquidity(df)
    except Exception as exc:
        logger.warning("add_amihud_illiquidity failed: %s", exc)
    df = add_fracdiff_features(df)
    try:
        df = add_calendar_features(df)
    except Exception as exc:
        logger.warning("add_calendar_features failed: %s", exc)
    return df