import pandas as pd
import numpy as np
from functools import reduce
from pathlib import Path
from typing import Literal
from functools import reduce

# Tickers mapped to your requested identifier names
TICKER_MAP = {
    "BZ=F": "brent_crude_oil",
    "GC=F": "gold",
    "INR=X": "usd_inr",
}


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 build_macro_features(df, prefix) -> pd.Series:
    """
    Generic — call for gold, brent, usd_inr with their prefix.
    df must have a 'close' column.
    """
    print(f"feature store {prefix} has length of {len(df)}")
    df = df.sort_index().copy()
    p = prefix  # e.g. "gold", "brent", "usd_inr"

    # ── Returns (keep) ────────────────────────────────────────
    df[f"{p}_ret_1d"] = df["close"].pct_change(1)
    df[f"{p}_ret_5d"] = df["close"].pct_change(5)
    df[f"{p}_ret_20d"] = df["close"].pct_change(20)

    # ── Trend (keep, fix momentum to be % not absolute) ───────
    df[f"{p}_sma_20"] = df["close"].rolling(20).mean()
    df[f"{p}_sma_50"] = df["close"].rolling(50).mean()
    df[f"{p}_sma_200"] = df["close"].rolling(200).mean()

    df[f"{p}_sma20_ratio"] = df["close"] / (df[f"{p}_sma_20"] + 1e-9) - 1
    df[f"{p}_sma200_ratio"] = df["close"] / (df[f"{p}_sma_200"] + 1e-9) - 1

    # Momentum as % not absolute (absolute is not cross-sectionally comparable)
    df[f"{p}_momentum_10"] = df["close"].pct_change(10)
    df[f"{p}_momentum_20"] = df["close"].pct_change(20)

    # ── Volatility (keep) ─────────────────────────────────────
    df[f"{p}_volatility_20d"] = df[f"{p}_ret_1d"].rolling(20).std()

    # ── ADD: Trend direction ──────────────────────────────────
    # Is asset in uptrend? Model knows level but not direction
    df[f"{p}_trend"] = (df[f"{p}_sma_20"] > df[f"{p}_sma_50"]).astype("int8")

    # ── ADD: Z-score (regime context) ────────────────────────
    df[f"{p}_zscore"] = compute_zscore(df["close"])

    # ── ADD: Volatility regime ────────────────────────────────
    df[f"{p}_vol_zscore"] = compute_zscore(df[f"{p}_volatility_20d"])
    df[f"{p}_high_vol"] = (df[f"{p}_vol_zscore"] > 1).astype("int8")

    return pd.concat([df])


MergeHow = Literal["left", "right", "outer", "inner", "cross"]


def merge_feature_frames_asof(
    dfs: dict[str, pd.DataFrame],
    timestamp_col: str = "timestamp",
    tolerance: str = "2D",
    prefix_columns: bool = True,
) -> pd.DataFrame:
    """
    Merge feature frames by NEAREST timestamp within `tolerance`, not exact match.
    Use this when sources may label the same trading day under slightly
    different timestamps (common when mixing futures and FX data).
    """
    prepped = {}
    for name, df in dfs.items():
        df = df.copy()
        df[timestamp_col] = pd.to_datetime(df[timestamp_col], utc=True).dt.tz_convert(
            None
        )
        df = df.sort_values(timestamp_col).reset_index(drop=True)

        dupes = df[timestamp_col].duplicated().sum()
        if dupes:
            print(f"WARNING: {name} has {dupes} duplicate timestamps")

        if prefix_columns:
            df = df.rename(
                columns={c: f"{name}_{c}" for c in df.columns if c != timestamp_col}
            )

        prepped[name] = df
        print(
            f"{name}: {df[timestamp_col].min().date()} -> {df[timestamp_col].max().date()} ({len(df)} rows)"
        )

    def _merge(left, right):
        return pd.merge_asof(
            left,
            right,
            on=timestamp_col,
            direction="nearest",
            tolerance=pd.Timedelta(tolerance),
        )

    merged = reduce(_merge, prepped.values())
    merged = merged.sort_values(timestamp_col).reset_index(drop=True)

    overlap_pct = merged.drop(columns=[timestamp_col]).notna().all(axis=1).mean() * 100
    print(
        f"\nmerged: {merged[timestamp_col].min().date()} -> {merged[timestamp_col].max().date()} "
        f"({len(merged)} rows, {overlap_pct:.1f}% fully-populated rows)"
    )

    return merged


def normalize_timestamp(df: pd.DataFrame, col: str = "timestamp") -> pd.DataFrame:
    df = df.copy()
    ts = pd.to_datetime(df[col], utc=True)
    df[col] = ts.dt.tz_convert(None).dt.normalize()
    return df


def feature_building_macro(macro_df: pd.Series):
    gold_df = (
        macro_df[macro_df["feature_type"] == TICKER_MAP["GC=F"]]
        .sort_values("timestamp")
        .reset_index(drop=True)
    )
    brent_df = (
        macro_df[macro_df["feature_type"] == TICKER_MAP["BZ=F"]]
        .sort_values("timestamp")
        .reset_index(drop=True)
    )
    usdinr_df = (
        macro_df[macro_df["feature_type"] == TICKER_MAP["INR=X"]]
        .sort_values("timestamp")
        .reset_index(drop=True)
    )

    gold_features = build_macro_features(gold_df, "gold")
    brent_features = build_macro_features(brent_df, "brent")
    usdinr_features = build_macro_features(usdinr_df, "usd_inr")

    gold_features = normalize_timestamp(gold_features)
    brent_features = normalize_timestamp(brent_features)
    usdinr_features = normalize_timestamp(usdinr_features)

    print(gold_features["timestamp"].dtype)
    print(gold_features["timestamp"].head(3).tolist())

    print(brent_features["timestamp"].dtype)
    print(brent_features["timestamp"].head(3).tolist())

    print(usdinr_features["timestamp"].dtype)
    print(usdinr_features["timestamp"].head(3).tolist())

    gold_dates = set(gold_features["timestamp"])
    brent_dates = set(brent_features["timestamp"])
    usdinr_dates = set(usdinr_features["timestamp"])

    print("gold ∩ brent ∩ usdinr:", len(gold_dates & brent_dates & usdinr_dates))
    print("gold ∩ brent:", len(gold_dates & brent_dates))
    print("gold ∩ usdinr:", len(gold_dates & usdinr_dates))
    print("brent ∩ usdinr:", len(brent_dates & usdinr_dates))

    for name, df in [
        ("gold", gold_features),
        ("brent", brent_features),
        ("usdinr", usdinr_features),
    ]:
        ts = pd.to_datetime(df["timestamp"], utc=True).dt.tz_convert(None)
        diffs = ts.sort_values().diff().dropna()
        print(f"\n{name}:")
        print(f"  rows: {len(df)}")
        print(f"  date range: {ts.min().date()} -> {ts.max().date()}")
        print(f"  gap distribution (days between consecutive rows):")
        print(diffs.dt.days.value_counts().sort_index().head(10))

    # USD/INR only — currency regime matters differently
    usdinr_features["usd_inr_appreciation"] = (
        usdinr_features["usd_inr_ret_5d"] > 0
    ).astype(
        "int8"
    )  # rupee weakening = capital outflow risk for equities

    # Brent only — supply shock detection
    brent_features["brent_spike"] = (
        brent_features["brent_ret_1d"].abs() > 0.03
    ).astype(
        "int8"
    )  # >3% single day move = supply shock event

    # merged = pd.concat(
    #     [gold_features, brent_features, usdinr_features], ignore_index=True
    # )
    # print(f"before gold has following columns {gold_features.columns}")
    gold_features = gold_features.drop(
        columns=["open", "high", "low", "close", "volume", "feature_type"]
    )
    # print(f"after gold has following columns {gold_features.columns}")

    # print(f"before brent has following columns {brent_features.columns}")
    brent_features = brent_features.drop(
        columns=["open", "high", "low", "close", "volume", "feature_type"]
    )
    # print(f"after brent has following columns {brent_features.columns}")

    # print(f"before usdinr has following columns {usdinr_features.columns}")
    usdinr_features = usdinr_features.drop(
        columns=["open", "high", "low", "close", "volume", "feature_type"]
    )
    # print(f"after usdinr has following columns {usdinr_features.columns}")

    # usage with your actual date ranges
    merged = merge_feature_frames_asof(
        {
            "gold": gold_features,  # 2000-08-30 -> 2026-06-29
            "brent": brent_features,  # 2007-07-31 -> 2026-06-30
            "usdinr": usdinr_features,  # 2003-12-02 -> 2026-07-01
        },
        "timestamp",
        "2D",
        False,
    )

    merged = merged.sort_values("timestamp").reset_index(drop=True)

    return pd.concat([merged])


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


def start_save_feature_store():
    from OHLCV_to_store import get_start_end_date_for_df
    from feature_building_vix import clean_nifty, identify_bad_df

    ohlcv_df = pd.read_parquet(INPUT_PATH)
    clean_df = clean_nifty(ohlcv_df)
    identify_bad_df(clean_df)
    macro_feature_store = feature_building_macro(clean_df)
    macro_feature_store = pd.concat([macro_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
    macro_feature_store.to_parquet(OUTPUT_PATH, index=False)

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


if __name__ == "__main__":
    start_save_feature_store()