import pandas as pd
import numpy as np
from pathlib import Path


def compute_zscore(series: pd.Series, period=252):
    mean = series.rolling(period).mean()
    std = series.rolling(period).std()

    return (series - mean) / (std + 1e-9)


def compute_percentile(series: pd.Series, period=252):
    return series.rolling(period).apply(
        lambda x: (x <= x.iloc[-1]).mean(),
        raw=False,
    )


def build_vix_store_features(vix_df: pd.Series):
    df = vix_df.sort_index().copy()

    df["vix_return_1d"] = df["close"].pct_change(1)
    df["vix_return_5d"] = df["close"].pct_change(5)
    df["vix_return_20d"] = df["close"].pct_change(20)

    df["vix_ma20"] = df["close"].rolling(20).mean()

    df["vix_volatility"] = df["vix_return_1d"].rolling(20).std()

    df["vix_zscore"] = compute_zscore(df["close"])
    df["vix_percentile"] = compute_percentile(df["close"])

    df["high_vol_regime"] = (df["vix_zscore"] > 1).astype("int8")
    df["low_vol_regime"] = (df["vix_zscore"] < -1).astype("int8")

    df["extreme_fear"] = (df["vix_zscore"] > 2).astype("int8")
    df["extreme_calm"] = (df["vix_zscore"] < -2).astype("int8")

    return df


def clean_nifty(df):
    """
    Clean NIFTY OHLCV data before feature engineering.

    Expected columns:
        timestamp, open, high, low, close, volume
    """

    original_len = len(df)

    # Ensure timestamp is datetime
    df["timestamp"] = pd.to_datetime(df["timestamp"])

    # Sort chronologically
    df = df.sort_values("timestamp").reset_index(drop=True)

    # ── Layer 1: Drop zero/negative prices ───────────────────
    price_cols = ["open", "high", "low", "close"]

    zero_mask = (df[price_cols] <= 0).any(axis=1)
    print(f"Dropped {zero_mask.sum()} rows with zero/negative prices")
    df = df.loc[~zero_mask]

    # ── Layer 2: Validate OHLC relationships ────────────────
    invalid_ohlc = (
        (df["high"] < df["low"])
        | (df["high"] < df["open"])
        | (df["high"] < df["close"])
        | (df["low"] > df["open"])
        | (df["low"] > df["close"])
    )

    print(f"Dropped {invalid_ohlc.sum()} rows with invalid OHLC relationships")
    df = df.loc[~invalid_ohlc]

    # ── Layer 3: Remove duplicate timestamps ────────────────
    # duplicate_mask = df.duplicated(subset=["timestamp"], keep="last")
    # print(f"Dropped {duplicate_mask.sum()} duplicate timestamps")
    # df = df.loc[~duplicate_mask]

    # ── Layer 4: Remove rows with missing OHLC values ───────
    nan_mask = df[["open", "high", "low", "close"]].isna().any(axis=1)
    print(f"Dropped {nan_mask.sum()} rows with missing OHLC values")
    df = df.loc[~nan_mask]

    df = df.reset_index(drop=True)

    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, ["volume", "open", "high", "low", "close"]])


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


def start_save_feature_store():
    ohlcv_df = pd.read_parquet(INPUT_PATH)
    clean_df = clean_nifty(ohlcv_df)
    identify_bad_df(clean_df)
    nifty_feature_store = build_vix_store_features(clean_df)
    nifty_feature_store = pd.concat([nifty_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
    nifty_feature_store.to_parquet(OUTPUT_PATH, index=False)

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


if __name__ == "__main__":
    start_save_feature_store()