"""
signals/technical.py — Technical signal engine.
RSI, MACD, EMA200, Bollinger Bands, VWAP (session-aware), ATR percentile filter.
Correlation guard for portfolio-level risk.
Returns TechnicalScore per contracts.py.
"""
from __future__ import annotations

import datetime
import logging
import zoneinfo
from typing import Literal

import numpy as np
import pandas as pd

import config
from contracts import TechnicalScore, PriceTarget

logger = logging.getLogger("trading_system.technical")

# ── Indicator Functions ──────────────────────────────────────────────────────

def compute_rsi(series: pd.Series, period: int = 14) -> pd.Series:
    """Relative Strength Index."""
    delta = series.diff()
    gain = delta.where(delta > 0, 0.0)
    loss = (-delta).where(delta < 0, 0.0)
    avg_gain = gain.ewm(alpha=1 / period, min_periods=period).mean()
    avg_loss = loss.ewm(alpha=1 / period, min_periods=period).mean()
    rs = avg_gain / avg_loss.replace(0, np.nan)
    rsi = 100 - (100 / (1 + rs))
    return rsi


def compute_macd(
    series: pd.Series,
    fast: int = 12,
    slow: int = 26,
    signal: int = 9,
) -> tuple[pd.Series, pd.Series, pd.Series]:
    """MACD line, signal line, histogram."""
    ema_fast = series.ewm(span=fast, min_periods=fast).mean()
    ema_slow = series.ewm(span=slow, min_periods=slow).mean()
    macd_line = ema_fast - ema_slow
    signal_line = macd_line.ewm(span=signal, min_periods=signal).mean()
    histogram = macd_line - signal_line
    return macd_line, signal_line, histogram


def compute_bollinger_bands(
    series: pd.Series,
    period: int = 20,
    std_dev: float = 2.0,
) -> tuple[pd.Series, pd.Series, pd.Series]:
    """Upper, middle, lower Bollinger Bands."""
    middle = series.rolling(window=period).mean()
    std = series.rolling(window=period).std()
    upper = middle + std_dev * std
    lower = middle - std_dev * std
    return upper, middle, lower


def compute_atr(df: pd.DataFrame, period: int = 14) -> pd.Series:
    """Average True Range on OHLC data."""
    high = df["high"]
    low = df["low"]
    close = df["close"].shift(1)
    tr1 = high - low
    tr2 = (high - close).abs()
    tr3 = (low - close).abs()
    tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
    atr = tr.ewm(alpha=1 / period, min_periods=period).mean()
    return atr


def compute_ema(series: pd.Series, period: int) -> pd.Series:
    """Exponential Moving Average."""
    return series.ewm(span=period, min_periods=period).mean()


def compute_vwap(df: pd.DataFrame) -> pd.Series:
    """Session-aware VWAP. Resets at 09:30 ET every trading day.

    Partitions by trading session date in US/Eastern timezone.
    VWAP = cumsum(typical_price * volume) / cumsum(volume)
    """
    if df.empty:
        return pd.Series(dtype=float)

    work = df.copy()

    # Get timestamps — handle both index-based and column-based
    if isinstance(work.index, pd.DatetimeIndex):
        idx = work.index
    elif "timestamp" in work.columns:
        idx = pd.DatetimeIndex(work["timestamp"])
    else:
        return pd.Series(np.nan, index=work.index)

    if idx.tz is None:
        idx = idx.tz_localize("UTC")
    eastern = idx.tz_convert(zoneinfo.ZoneInfo("US/Eastern"))
    
    mask = (
        ((eastern.hour > 9) | ((eastern.hour == 9) & (eastern.minute >= 30))) &
        ((eastern.hour < 16) | ((eastern.hour == 16) & (eastern.minute == 0)))
    )
    
    work = work[mask].copy()
    if work.empty:
        return pd.Series(np.nan, index=df.index)
        
    eastern = eastern[mask]
    work["session_date"] = eastern.date

    typical_price = (work["high"] + work["low"] + work["close"]) / 3
    work["tp_vol"] = typical_price * work["volume"]

    vwap_values = []
    for _, group in work.groupby("session_date"):
        cum_tp_vol = group["tp_vol"].cumsum()
        cum_vol = group["volume"].cumsum().replace(0, np.nan)
        session_vwap = cum_tp_vol / cum_vol
        vwap_values.append(session_vwap)

    if vwap_values:
        result = pd.concat(vwap_values).reindex(df.index)
    else:
        result = pd.Series(np.nan, index=df.index)
    return result


# ── ATR Percentile Filter ───────────────────────────────────────────────────

def atr_percentile_filter(
    df_5m: pd.DataFrame,
    window_days: int | None = None,
) -> tuple[bool, float, float]:
    """Check if current ATR is above the 20th percentile of the rolling window.

    Returns: (passes_filter, current_atr, threshold_20pct)
    """
    window_days = window_days or config.ATR_PERCENTILE_WINDOW_DAYS
    atr_series = compute_atr(df_5m, period=14)
    if atr_series.dropna().empty:
        return True, 0.0, 0.0

    bars_per_day = 78  # ~6.5 hours of 5m bars
    window_bars = window_days * bars_per_day
    window = atr_series.dropna().tail(window_bars)

    if len(window) < bars_per_day:
        return True, float(atr_series.iloc[-1]), 0.0

    current_atr = float(atr_series.iloc[-1])
    threshold = float(np.percentile(window.values, 20))
    passes = current_atr >= threshold
    return passes, current_atr, threshold


# ── EMA200 Trend Filter ─────────────────────────────────────────────────────

def ema200_trend_bias(daily_df: pd.DataFrame) -> int:
    """Compute EMA200 trend bias on daily bars.

    Returns: +1 (bullish), -1 (bearish), 0 (disabled/neutral)
    """
    if daily_df.empty or len(daily_df) < 200:
        if not daily_df.empty:
            logger.warning(
                "Daily bars < 200 (%d), EMA200 filter disabled.", len(daily_df)
            )
        return 0

    ema200 = compute_ema(daily_df["close"], 200)
    last_close = float(daily_df["close"].iloc[-1])
    last_ema = float(ema200.iloc[-1])

    if np.isnan(last_ema):
        return 0
    if last_close > last_ema:
        return 1
    elif last_close < last_ema:
        return -1
    return 0


# ── Correlation Guard ────────────────────────────────────────────────────────

def build_correlation_matrix(
    symbols: list[str],
    get_daily_bars_fn,
    window: int = 60,
) -> pd.DataFrame:
    """Build correlation matrix on 60-day rolling returns."""
    returns = {}
    for sym in symbols:
        try:
            df = get_daily_bars_fn(sym)
            if df is not None and len(df) >= window:
                returns[sym] = df["close"].pct_change().tail(window)
        except Exception:
            continue
    if not returns:
        return pd.DataFrame()
    return pd.DataFrame(returns).corr()


def correlation_guard(
    symbol: str,
    direction: Literal["buy", "sell"],
    open_positions: dict[str, str],
    corr_matrix: pd.DataFrame,
    threshold: float = 0.85,
) -> tuple[bool, str | None]:
    """Check if a proposed trade is blocked by correlation with open positions.

    Args:
        symbol: Proposed trade symbol
        direction: "buy" or "sell"
        open_positions: {symbol: "buy"|"sell"} for currently open positions
        corr_matrix: Correlation matrix from build_correlation_matrix
        threshold: Correlation threshold

    Returns: (allowed, skip_reason)
    """
    if corr_matrix.empty or symbol not in corr_matrix.columns:
        return True, None

    for pos_sym, pos_dir in open_positions.items():
        if pos_sym not in corr_matrix.columns:
            continue
        corr = corr_matrix.loc[symbol, pos_sym]
        if abs(corr) > threshold and direction == pos_dir:
            reason = f"correlated_position_same_direction: corr({symbol},{pos_sym})={corr:.3f}"
            return False, reason

    return True, None


# ── Signal Score Composition ─────────────────────────────────────────────────
#
# Weights (must sum to 1.0):
#   RSI(14) divergence from 50:    0.20
#   MACD histogram sign + slope:   0.20
#   BB position (% from mean):     0.15
#   VWAP deviation:                0.20
#   EMA200 trend bias:             0.15
#   ATR percentile (pass/fail):    0.10 (0 if below threshold)

SIGNAL_WEIGHTS = {
    "rsi": 0.20,
    "macd": 0.20,
    "bb": 0.15,
    "vwap": 0.20,
    "ema200": 0.15,
    "atr": 0.10,
}
assert abs(sum(SIGNAL_WEIGHTS.values()) - 1.0) < 0.001, "Signal weights must sum to 1.0"


def compute_technical(
    symbol: str,
    df_5m: pd.DataFrame,
    daily_df: pd.DataFrame | None = None,
    open_positions: dict[str, str] | None = None,
    corr_matrix: pd.DataFrame | None = None,
) -> tuple[TechnicalScore, PriceTarget]:
    """Compute full technical score and price target for a symbol.

    Args:
        symbol: Stock symbol
        df_5m: 5-minute OHLCV DataFrame (needs ~100+ bars)
        daily_df: Daily OHLCV DataFrame (needs 200+ bars for EMA200)
        open_positions: Currently open positions {symbol: direction}
        corr_matrix: Correlation matrix for guard

    Returns: (TechnicalScore, PriceTarget)
    """
    now = datetime.datetime.now(datetime.timezone.utc)
    sub_signals: dict[str, float] = {}
    n_computed = 0
    skip_reason = None

    # Minimum data check
    if len(df_5m) < 30:
        return (
            TechnicalScore(
                symbol=symbol, timestamp=now, score=0.0,
                sub_signals={}, skip_reason="insufficient_data",
                n_signals_computed=0,
            ),
            PriceTarget(symbol=symbol, current_price=0.0, predicted_direction="NEUTRAL"),
        )

    close = df_5m["close"]
    current_price = float(close.iloc[-1])

    # ── 1. RSI (weight 0.20) ──
    try:
        rsi = compute_rsi(close, 14)
        last_rsi = float(rsi.iloc[-1])
        if not np.isnan(last_rsi):
            # Normalize: RSI < 30 → +1 (oversold = bullish), RSI > 70 → -1 (overbought = bearish)
            rsi_score = (50 - last_rsi) / 50  # range roughly -1 to +1
            rsi_score = max(-1.0, min(1.0, rsi_score))
            sub_signals["rsi"] = rsi_score
            n_computed += 1
    except Exception as e:
        logger.debug("RSI computation failed for %s: %s", symbol, e)

    # ── 2. MACD (weight 0.20) ──
    try:
        _, _, histogram = compute_macd(close)
        hist_vals = histogram.dropna()
        if len(hist_vals) >= 2:
            last_hist = float(hist_vals.iloc[-1])
            prev_hist = float(hist_vals.iloc[-2])
            # Sign gives direction, slope gives momentum
            sign = 1.0 if last_hist > 0 else -1.0
            slope = 1.0 if last_hist > prev_hist else -1.0
            macd_score = 0.5 * sign + 0.5 * slope
            macd_score = max(-1.0, min(1.0, macd_score))
            sub_signals["macd"] = macd_score
            n_computed += 1
    except Exception as e:
        logger.debug("MACD computation failed for %s: %s", symbol, e)

    # ── 3. Bollinger Bands (weight 0.15) ──
    try:
        upper, middle, lower = compute_bollinger_bands(close)
        last_close = current_price
        last_upper = float(upper.iloc[-1])
        last_lower = float(lower.iloc[-1])
        last_middle = float(middle.iloc[-1])
        if not any(np.isnan([last_upper, last_lower, last_middle])):
            band_width = last_upper - last_lower
            if band_width > 0:
                # Position within bands: -1 at upper (sell), +1 at lower (buy)
                bb_pos = (last_middle - last_close) / (band_width / 2)
                bb_score = max(-1.0, min(1.0, bb_pos))
                sub_signals["bb"] = bb_score
                n_computed += 1
    except Exception as e:
        logger.debug("BB computation failed for %s: %s", symbol, e)

    # ── 4. VWAP (weight 0.20) ──
    try:
        vwap = compute_vwap(df_5m)
        last_vwap = float(vwap.iloc[-1])
        if not np.isnan(last_vwap) and last_vwap > 0:
            # Price below VWAP = bullish (buy opportunity), above = bearish
            deviation = (last_vwap - current_price) / last_vwap
            vwap_score = max(-1.0, min(1.0, deviation * 20))  # Scale
            sub_signals["vwap"] = vwap_score
            n_computed += 1
    except Exception as e:
        logger.debug("VWAP computation failed for %s: %s", symbol, e)

    # ── 5. EMA200 Trend Bias (weight 0.15) ──
    try:
        if daily_df is not None and not daily_df.empty:
            trend = ema200_trend_bias(daily_df)
            sub_signals["ema200"] = float(trend)
            n_computed += 1
        else:
            sub_signals["ema200"] = 0.0
    except Exception as e:
        logger.debug("EMA200 computation failed for %s: %s", symbol, e)

    # ── 6. ATR Percentile Filter (weight 0.10) ──
    try:
        atr_pass, current_atr, atr_threshold = atr_percentile_filter(df_5m)
        if atr_pass:
            sub_signals["atr"] = 1.0
        else:
            sub_signals["atr"] = 0.0
            skip_reason = "atr_below_20th_percentile"
        n_computed += 1
    except Exception as e:
        logger.debug("ATR computation failed for %s: %s", symbol, e)
        current_atr = 0.0

    # ── Correlation Guard ──
    if skip_reason is None and open_positions and corr_matrix is not None:
        weighted_score = sum(
            sub_signals.get(k, 0.0) * SIGNAL_WEIGHTS[k] for k in SIGNAL_WEIGHTS
        )
        proposed_dir = "buy" if weighted_score > 0 else "sell"
        allowed, corr_reason = correlation_guard(
            symbol, proposed_dir, open_positions, corr_matrix
        )
        if not allowed:
            skip_reason = corr_reason

    # ── Compute weighted score ──
    if n_computed < 3:
        score = 0.0
        if skip_reason is None:
            skip_reason = "insufficient_data"
    else:
        score = sum(
            sub_signals.get(k, 0.0) * SIGNAL_WEIGHTS[k] for k in SIGNAL_WEIGHTS
        )
        score = max(-1.0, min(1.0, score))

    # ── Build Price Target ──
    atr_val = current_atr if current_atr > 0 else current_price * 0.005
    if score > 0:
        predicted_dir = "UP"
        entry = current_price
        stop = current_price - 1.0 * atr_val
        tp = current_price + 2.0 * atr_val
    elif score < 0:
        predicted_dir = "DOWN"
        entry = current_price
        stop = current_price + 1.0 * atr_val
        tp = current_price - 2.0 * atr_val
    else:
        predicted_dir = "NEUTRAL"
        entry = None
        stop = None
        tp = None

    expected_profit = 0.0
    if entry and tp:
        qty_estimate = (config.ALLOCATED_CAPITAL * config.RISK_PER_TRADE_PCT / 100) / atr_val if atr_val > 0 else 0
        expected_profit = abs(tp - entry) * qty_estimate

    price_target = PriceTarget(
        symbol=symbol,
        current_price=current_price,
        predicted_direction=predicted_dir,
        entry_price=entry,
        stop_loss_price=stop,
        take_profit_price=tp,
        expected_profit_usd=round(expected_profit, 2),
        confidence=abs(score),
    )

    tech_score = TechnicalScore(
        symbol=symbol,
        timestamp=now,
        score=score,
        sub_signals=sub_signals,
        skip_reason=skip_reason,
        n_signals_computed=n_computed,
    )

    return tech_score, price_target