from typing import Tuple
import numpy as np
import pandas as pd
import math
from pathlib import Path
import warnings

# Force all numpy/pandas runtime warnings to raise an exception instead
warnings.simplefilter("error", RuntimeWarning)


def compute_rsi(series, period=14):
    delta = series.diff()
    gain = delta.clip(lower=0).rolling(period).mean()
    loss = (-delta.clip(upper=0)).rolling(period).mean()
    rs = gain / (loss + 1e-9)
    return 100 - (100 / (1 + rs))


def compute_atr(high, low, close, period=14):
    high_low = high - low
    high_close = (high - close.shift(1)).abs()
    low_close = (low - close.shift(1)).abs()

    true_range = pd.concat([high_low, high_close, low_close], axis=1).max(axis=1)

    atr = true_range.ewm(alpha=1 / period, adjust=False).mean()

    return atr


def compute_aroon(high, low, period=25):
    """
    Returns
    -------
    aroon_up : pd.Series
    aroon_down : pd.Series
    """

    days_since_high = high.rolling(period).apply(
        lambda x: period - 1 - np.argmax(x),
        raw=True,
    )

    days_since_low = low.rolling(period).apply(
        lambda x: period - 1 - np.argmin(x),
        raw=True,
    )

    aroon_up = ((period - days_since_high) / period) * 100

    aroon_down = ((period - days_since_low) / period) * 100

    return aroon_up, aroon_down


def compute_parkinson_volatility(high, low, window=20):
    """
    Parkinson volatility estimator.

    Parameters
    ----------
    high : pd.Series
    low : pd.Series
    window : int

    Returns
    -------
    pd.Series
        Rolling Parkinson volatility.
    """
    # Replaces == 0 with a tolerance check
    new_low = np.where(np.isclose(low, 0.0, atol=1e-9), 1e-9, low)
    log_hl_sq = np.log(high / new_low).pow(2)
    return np.sqrt(log_hl_sq.rolling(window).sum() / (4 * window * np.log(2)))


def compute_adx(high, low, close, period=14) -> Tuple[pd.Series, pd.Series, pd.Series]:
    """
    Compute ADX, DI+ and DI- using Wilder's smoothing.

    Parameters
    ----------
    high : pd.Series
    low : pd.Series
    close : pd.Series
    period : int, default=14

    Returns
    -------
    adx : pd.Series
    di_plus : pd.Series
    di_minus : pd.Series
    """

    # ----- Directional Movement -----
    up_move = high.diff()
    down_move = -low.diff()

    plus_dm = pd.Series(
        np.where(
            (up_move > down_move) & (up_move > 0),
            up_move,
            0.0,
        ),
        index=high.index,
    )

    minus_dm = pd.Series(
        np.where(
            (down_move > up_move) & (down_move > 0),
            down_move,
            0.0,
        ),
        index=high.index,
    )

    # ----- True Range -----
    tr = compute_atr(high, low, close, period)

    atr = tr.ewm(alpha=1 / period, adjust=False).mean()

    # ----- Wilder smoothing -----
    plus_dm_smoothed = plus_dm.ewm(
        alpha=1 / period,
        adjust=False,
    ).mean()

    minus_dm_smoothed = minus_dm.ewm(
        alpha=1 / period,
        adjust=False,
    ).mean()

    # ----- Directional Indicators -----
    di_plus = (plus_dm_smoothed / (atr + 1e-9)) * 100
    di_minus = (minus_dm_smoothed / (atr + 1e-9)) * 100

    # ----- Directional Index -----
    dx = ((di_plus - di_minus).abs() / (di_plus + di_minus + 1e-9)) * 100

    # ----- Average Directional Index -----
    adx = dx.ewm(
        alpha=1 / period,
        adjust=False,
    ).mean()

    return adx, di_plus, di_minus


def mean_absolute_deviation(x):
    return np.mean(np.abs(x - np.mean(x)))


def compute_cci(high, low, close, period=20):
    """
    Commodity Channel Index (CCI)
    """

    typical_price = (high + low + close) / 3

    sma = typical_price.rolling(period).mean()

    mean_deviation = typical_price.rolling(period).apply(
        mean_absolute_deviation,
        raw=True,
    )

    cci = (typical_price - sma) / (0.015 * (mean_deviation + 1e-9))

    return cci


def compute_stochastic_k(high, low, close, period=14):
    """
    Stochastic Oscillator %K
    """

    highest_high = high.rolling(period).max()
    lowest_low = low.rolling(period).min()

    stochastic_k = ((close - lowest_low) / (highest_high - lowest_low + 1e-9)) * 100

    return stochastic_k


def compute_macd(close, fast_period=12, slow_period=26, signal_period=9):
    """
    Compute MACD, Signal Line and Histogram.

    Parameters
    ----------
    close : pd.Series
        Closing prices.
    fast_period : int, default=12
    slow_period : int, default=26
    signal_period : int, default=9

    Returns
    -------
    macd : pd.Series
    signal : pd.Series
    histogram : pd.Series
    """

    ema_fast = close.ewm(
        span=fast_period,
        adjust=False,
    ).mean()

    ema_slow = close.ewm(
        span=slow_period,
        adjust=False,
    ).mean()

    macd = ema_fast - ema_slow

    signal = macd.ewm(
        span=signal_period,
        adjust=False,
    ).mean()

    histogram = macd - signal

    return macd, signal, histogram


def compute_bollinger(close, period=20, num_std=2):
    """
    Compute Bollinger Band features.

    Parameters
    ----------
    close : pd.Series
        Closing prices.
    period : int, default=20
        Rolling window for SMA and standard deviation.
    num_std : float, default=2
        Number of standard deviations for the bands.

    Returns
    -------
    bb_width : pd.Series
        Normalized Bollinger Band width.

    bb_position : pd.Series
        Position of the close within the bands.
        0   -> Lower Band
        0.5 -> Middle Band
        1   -> Upper Band

    bb_squeeze : pd.Series
        Width normalized by its rolling mean.
        <1 : Bands tighter than usual.
        >1 : Bands wider than usual.
    """

    middle = close.rolling(period).mean()
    std = close.rolling(period).std()

    upper = middle + num_std * std
    lower = middle - num_std * std

    # Normalized width
    bb_width = (upper - lower) / (middle + 1e-9)

    # Position inside the bands
    bb_position = (close - lower) / (upper - lower + 1e-9)

    # Relative squeeze
    bb_squeeze = bb_width / (bb_width.rolling(period).mean() + 1e-9)

    return bb_width, bb_position, bb_squeeze


import numpy as np
import pandas as pd


def compute_volume_features(
    high,
    low,
    close,
    volume,
    volume_ma_period=20,
    mfi_period=14,
):
    """
    Compute volume-based features.

    Returns
    -------
    volume_ma20
    volume_ratio
    obv
    vwap
    mfi
    """

    # -------------------------------------------------
    # Volume Moving Average
    # -------------------------------------------------

    volume_ma = volume.rolling(volume_ma_period).mean()

    volume_ratio = volume / (volume_ma + 1e-9)

    # -------------------------------------------------
    # OBV
    # -------------------------------------------------

    price_change = close.diff()

    obv = np.sign(price_change).fillna(0).mul(volume).cumsum()

    # -------------------------------------------------
    # VWAP (Cumulative)
    # -------------------------------------------------

    typical_price = (high + low + close) / 3

    vwap = (typical_price * volume).cumsum() / (volume.cumsum() + 1e-9)

    # -------------------------------------------------
    # Money Flow Index (MFI)
    # -------------------------------------------------

    raw_money_flow = typical_price * volume

    positive_flow = raw_money_flow.where(
        typical_price > typical_price.shift(1),
        0.0,
    )

    negative_flow = raw_money_flow.where(
        typical_price < typical_price.shift(1),
        0.0,
    )

    positive_sum = positive_flow.rolling(mfi_period).sum()

    negative_sum = negative_flow.rolling(mfi_period).sum()

    money_ratio = positive_sum / (negative_sum + 1e-9)

    mfi = 100 - (100 / (1 + money_ratio))

    return (
        volume_ma,
        volume_ratio,
        obv,
        vwap,
        mfi,
    )


def compute_candlestick_features(
    open_,
    high,
    low,
    close,
    doji_threshold=0.1,
):
    """
    Compute candlestick-based features.

    Parameters
    ----------
    open_ : pd.Series
    high : pd.Series
    low : pd.Series
    close : pd.Series
    doji_threshold : float, default=0.1
        Maximum body percentage to classify as a Doji.

    Returns
    -------
    body_percent
    upper_shadow_percent
    lower_shadow_percent
    gap_up
    gap_down
    inside_day
    outside_day
    doji
    """

    candle_range = (high - low).replace(0, np.nan)

    # ---------------------------------------------------------
    # Body
    # ---------------------------------------------------------

    body = (close - open_).abs()

    body_percent = body / candle_range

    # ---------------------------------------------------------
    # Upper Shadow
    # ---------------------------------------------------------

    upper_shadow = high - np.maximum(open_, close)

    upper_shadow_percent = upper_shadow / candle_range

    # ---------------------------------------------------------
    # Lower Shadow
    # ---------------------------------------------------------

    lower_shadow = np.minimum(open_, close) - low

    lower_shadow_percent = lower_shadow / candle_range

    # ---------------------------------------------------------
    # Gap Up / Gap Down
    # ---------------------------------------------------------

    previous_high = high.shift(1)
    previous_low = low.shift(1)

    gap_up = (low > previous_high).astype(int)

    gap_down = (high < previous_low).astype(int)

    # ---------------------------------------------------------
    # Inside / Outside Day
    # ---------------------------------------------------------

    inside_day = ((high < previous_high) & (low > previous_low)).astype(int)

    outside_day = ((high > previous_high) & (low < previous_low)).astype(int)

    # ---------------------------------------------------------
    # Doji
    # ---------------------------------------------------------

    doji = (body_percent <= doji_threshold).astype(int)

    return (
        body_percent,
        upper_shadow_percent,
        lower_shadow_percent,
        gap_up,
        gap_down,
        inside_day,
        outside_day,
        doji,
    )


def compute_relative_position(high, low, close, period=252):
    """
    Compute relative position features.

    Parameters
    ----------
    high : pd.Series
    low : pd.Series
    close : pd.Series
    period : int, default=252
        Number of trading days representing one year.

    Returns
    -------
    distance_from_52w_high : pd.Series
    distance_from_52w_low : pd.Series
    rolling_drawdown : pd.Series
    """

    # -------------------------------------------------
    # 52-week High / Low
    # -------------------------------------------------

    rolling_high = high.rolling(period).max()
    rolling_low = low.rolling(period).min()

    distance_from_52w_high = (close - rolling_high) / (rolling_high + 1e-9)

    distance_from_52w_low = (close - rolling_low) / (rolling_low + 1e-9)

    # -------------------------------------------------
    # Rolling Drawdown
    # -------------------------------------------------

    rolling_drawdown = (close - rolling_high) / (rolling_high + 1e-9)

    return (
        distance_from_52w_high,
        distance_from_52w_low,
        rolling_drawdown,
    )


def build_features_ohlcv(ohlcv_df) -> pd.Series:
    """
    this will build all the features we can build from ohlcv.
    Input:  raw OHLCV per stock + market data
    Output: feature matrix, one row per (symbol, date)
    """
    features = []
    symbols = ohlcv_df["symbol"].unique() if "symbol" in ohlcv_df.columns else None

    if symbols is None or len(symbols) == 0:
        print("No tickers were found, exiting")
        return []

    for symbol in symbols:
        df = (
            ohlcv_df[ohlcv_df["symbol"] == symbol].copy() if symbol else ohlcv_df.copy()
        )
        df = df.sort_index()
        # price based features ----------------------------------------------------------
        # log returns
        df["log_ret_1d"] = np.log(df["close"] / df["close"].shift(1))
        df["log_ret_3d"] = np.log(df["close"] / df["close"].shift(3))
        df["log_ret_5d"] = np.log(df["close"] / df["close"].shift(5))
        df["log_ret_10d"] = np.log(df["close"] / df["close"].shift(10))
        df["log_ret_20d"] = np.log(df["close"] / df["close"].shift(20))
        df["log_ret_60d"] = np.log(df["close"] / df["close"].shift(60))

        # simple returns -----------------------------------------------------------------
        for period in [1, 3, 5, 10, 20, 60, 120]:
            df[f"ret_{period}d"] = df["close"].pct_change(period)

        # simple moving average
        df["sma_20d"] = df["close"].rolling(20).mean()
        df["sma_50d"] = df["close"].rolling(50).mean()
        df["sma_200d"] = df["close"].rolling(200).mean()

        # Exponential Moving Average (EMA) 20,50
        df["ema_20d"] = df["close"].ewm(span=20, adjust=False).mean()
        df["ema_50d"] = df["close"].ewm(span=50, adjust=False).mean()

        # price position
        # close_sma20_ratio close_sma50_ratio close_sma200_ratio close_ema20_ratio close_ema50_ratio high_20_position low_20_position
        # distance_from_52w_high distance_from_52w_low
        df["close_sma20_ratio"] = df["close"] / df["sma_20d"]
        df["close_sma50_ratio"] = df["close"] / df["sma_50d"]
        df["close_sma200_ratio"] = df["close"] / df["sma_200d"]
        df["close_ema20_ratio"] = df["close"] / df["ema_20d"]
        df["high_20"] = df["high"].rolling(20).max()
        df["high_20_position"] = df["close"] / (df["high_20"] + 1e-9)
        df["low_20"] = df["low"].rolling(20).min()
        df["low_20_position"] = (df["close"] - df["low_20"]) / (df["low_20"] + 1e-9)
        df["high_52w"] = df["high"].rolling(252).max()
        df["close_to_52w_high"] = (df["close"] - df["high_52w"]) / (
            df["high_52w"] + 1e-9
        )
        df["low_52w"] = df["low"].rolling(252).min()
        df["close_to_52w_low"] = (df["close"] - df["low_52w"]) / (df["low_52w"] + 1e-9)
        df["position_in_20d_range"] = (df["close"] - df["low_20"]) / (
            df["high_20"] - df["low_20"] + 1e-9
        )
        # Momentum — ROC: yes for 10, 20, 60. Momentum 10, 20 yes. PPO yes, APO no (PPO is just normalized APO, keep one)
        # ROC
        for period in [10, 20, 60]:
            df[f"roc_{period}"] = df["close"].pct_change(period) * 100

        # Momentum
        for period in [10, 20]:
            df[f"momentum_{period}"] = df["close"] - df["close"].shift(period)
        # PPO
        df["ppo"] = ((df["ema_20d"] - df["ema_50d"]) / (df["ema_50d"] + 1e-9)) * 100

        # Distance from moving averages (normalized)
        df["dist_sma20"] = (df["close"] - df["sma_20d"]) / df["sma_20d"]
        df["dist_sma50"] = (df["close"] - df["sma_50d"]) / df["sma_50d"]
        df["dist_sma200"] = (df["close"] - df["sma_200d"]) / df["sma_200d"]

        # Volatility
        # Volatility — rolling_std: yes for 10, 20, 60. Skip 5 (noise).
        # ATR(Average True Range) 14 yes, ATR21 no (redundant). ATR_percent yes.
        # Parkinson yes (uses high/low, genuinely different from close-to-close std).
        # True Range no (ATR already captures it)
        for period in [10, 20, 60]:
            df[f"vol_{period}"] = df["ret_1d"].rolling(period).std()

        df["vol_ratio"] = df["vol_10"] / df["vol_20"]  # vol regime

        df["atr_14"] = compute_atr(df["high"], df["low"], df["close"])

        df["atr_percent"] = (df["atr_14"] / (df["close"] + 1e-9)) * 100

        df["parkinson_volatility"] = compute_parkinson_volatility(df["high"], df["low"])

        #  Trend Strength — ADX14 yes, ADX20 no (redundant). DI+ and DI- both yes. Aroon Up and Down both yes.
        df["adx_14"], df["di_plus"], df["di_minus"] = compute_adx(
            df["high"],
            df["low"],
            df["close"],
        )

        df["aroon_up"], df["aroon_down"] = compute_aroon(
            df["high"],
            df["low"],
        )
        # Oscillators — RSI14 yes, RSI7
        # CCI20, Stochastic K
        # RSI
        df["rsi_14"] = compute_rsi(df["close"], 14)
        df["rsi_7"] = compute_rsi(df["close"], 7)
        df["cci_20"] = compute_cci(
            df["high"],
            df["low"],
            df["close"],
        )
        df["stochastic_k"] = compute_stochastic_k(
            df["high"],
            df["low"],
            df["close"],
        )

        #  MACD — MACD yes, Signal yes, Histogram yes (all three — histogram is the most predictive of the three)
        df["macd"], df["macd_signal"], df["macd_histogram"] = compute_macd(df["close"])
        # Bollinger — BB Width yes, BB Position yes, BB Squeeze yes.
        # Skip Upper and Lower raw values (BB Position already captures where price sits, raw levels aren't meaningful cross-sectionally)
        df["bb_width"], df["bb_position"], df["bb_squeeze"] = compute_bollinger(
            df["close"]
        )
        # Range features
        df["intraday_range"] = (df["high"] - df["low"]) / df["close"]
        df["gap"] = (df["open"] - df["close"].shift(1)) / df["close"].shift(1)
        df["close_position"] = (df["close"] - df["low"]) / (
            df["high"] - df["low"] + 1e-9
        )  # 0=low, 1=high

        # Volume — Volume MA20 yes, MA50 no (redundant).
        # Volume Ratio yes. OBV yes.
        # VWAP yes.
        # MFI yes (combines price + volume, genuinely different).
        # Chaikin Money Flow no (redundant with MFI)
        # ── Volume features ───────────────────────────────────
        (
            df["volume_ma20"],
            df["volume_ratio"],
            df["obv"],
            df["vwap"],
            df["mfi"],
        ) = compute_volume_features(
            df["high"],
            df["low"],
            df["close"],
            df["volume"],
        )

        # Candlestick — Body % yes, Upper Shadow % yes, Lower Shadow % yes, Gap Up yes, Gap Down yes,
        # Inside Day yes, Outside Day yes. Doji yes.
        # Skip Hammer, Shooting Star, Bullish/Bearish Engulfing — these are rule-based patterns XGBoost will reconstruct itself from body/shadow/gap features anyway.
        # Adding them explicitly is redundant.
        (
            df["body_percent"],
            df["upper_shadow_percent"],
            df["lower_shadow_percent"],
            df["gap_up"],
            df["gap_down"],
            df["inside_day"],
            df["outside_day"],
            df["doji"],
        ) = compute_candlestick_features(
            df["open"],
            df["high"],
            df["low"],
            df["close"],
        )
        # Relative Position — 52w high yes, 52w low yes, Rolling Max no (52w high covers it), Rolling Min no (same). Rolling Drawdown yes.
        (
            df["distance_from_52w_high"],
            df["distance_from_52w_low"],
            df["rolling_drawdown"],
        ) = compute_relative_position(
            df["high"],
            df["low"],
            df["close"],
        )

        # ── Calendar features ─────────────────────────────────
        # FIX: Direct column-level datetime extraction
        df["day_of_week"] = df["timestamp"].dt.dayofweek
        df["month"] = df["timestamp"].dt.month
        df["is_month_end"] = df["timestamp"].dt.is_month_end.astype(int)

        df["symbol"] = symbol
        features.append(df)

    features_df = pd.concat(features)
    return features_df


def clean_ohlcv(df):
    """
    Apply before any feature engineering.
    """
    original_len = len(df)

    # ── Layer 1: Drop zero/negative prices ───────────────────
    # Any OHLCV value of 0 is invalid
    price_cols = ["open", "high", "low", "close"]
    zero_mask = (df[price_cols] <= 0).any(axis=1)
    df = df[~zero_mask]
    print(f"Dropped {zero_mask.sum()} rows with zero/negative prices")

    # ── Layer 2: Drop impossible OHLC relationships ───────────
    invalid_ohlc = (
        (df["high"] < df["low"])  # high below low
        | (df["high"] < df["close"])  # high below close
        | (df["high"] < df["open"])  # high below open
        | (df["low"] > df["close"])  # low above close
        | (df["low"] > df["open"])  # low above open
    )
    df = df[~invalid_ohlc]
    print(f"Dropped {invalid_ohlc.sum()} rows with invalid OHLC relationships")

    # ── Layer 3: Drop symbols with insufficient history ───────
    # A symbol needs at least 250 rows (≈1 year) for 200d SMA to warm up
    symbol_counts = df.groupby("symbol")["close"].count()
    valid_symbols = symbol_counts[symbol_counts >= 250].index
    dropped_symbols = symbol_counts[symbol_counts < 250].index.tolist()
    df = df[df["symbol"].isin(valid_symbols)]
    print(f"Dropped {len(dropped_symbols)} symbols with < 250 trading days")
    print(
        f"Dropped symbols: {dropped_symbols[:10]}{'...' if len(dropped_symbols) > 10 else ''}"
    )

    print(
        f"\nTotal rows: {original_len:,} → {len(df):,} "
        f"({original_len - len(df):,} removed)"
    )

    return df


def sanity_check(df):
    print(f"Rows: {len(df):,}")
    print(f"Symbols: {df['symbol'].nunique()}")
    print(f"Date range: {df.index.min()} → {df.index.max()}")
    print(f"Null counts:\n{df[['open','high','low','close','volume']].isnull().sum()}")
    print(f"Min close: {df['close'].min()}")
    print(f"Any zero close: {(df['close'] <= 0).any()}")


def identify_bad_df(df):
    bad = (df["close"] <= 0) | (df["close"].shift(1) <= 0)

    print(df.loc[bad, ["symbol", "open", "high", "low", "close"]])


# input path for OHLCV
INPUT_PATH = Path(
    "H:/Developer/stock_model/Dataset/Processed_dataset/feature_stores/us_stock_store.parquet"
)
# output path
OUTPUT_PATH = Path(
    "H:/Developer/stock_model/Dataset/Processed_dataset/feature_stores/us_stock_feature_store.parquet"
)


def start_save_feature_store():
    ohlcv_df = pd.read_parquet(INPUT_PATH)
    clean_df = clean_ohlcv(ohlcv_df)
    ohlcv_feature_store = build_features_ohlcv(clean_df)
    # ohlcv_feature_store = pd.concat(ohlcv_feature_store, ignore_index=True)
    # save the new store
    # Ensure destination tracking directory paths exist natively
    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)

    # Export out to structural Parquet architecture
    ohlcv_feature_store.to_parquet(OUTPUT_PATH, index=False)

    print(f"\n✅ Success! OHLCV Feature Store created at: {OUTPUT_PATH}")
    print(f"Total rows recorded: {len(ohlcv_feature_store)}")
    print(
        f"Timeline Range: {ohlcv_feature_store['timestamp'].min().strftime('%Y-%m-%d')} to {ohlcv_feature_store['timestamp'].max().strftime('%Y-%m-%d')}"
    )


if __name__ == "__main__":
    start_save_feature_store()