"""HMM regime detector — 3-state Gaussian HMM on realized_vol.

Fits on the FULL dataset (training + OOS) so the state sequence is globally
consistent across all periods. States are sorted by mean realized vol (ascending):
    0 = calm,   1 = stressed,   2 = crisis

Output: data/processed/eurusd_regime.parquet
Columns: datetime (or date), regime (int: 0/1/2)

Deterministic: fixed random_state=42 throughout.
Cached: if the parquet already exists, it is not re-computed.

Usage:
    python scripts/prepare_regime.py [--hourly] [--daily] [--force]

jepa-fx-risk#13
"""

import argparse
import os
from pathlib import Path

import numpy as np
import pandas as pd
from hmmlearn import hmm

DATA_DIR = Path(__file__).parent.parent / "data" / "processed"
HOURLY_PATH = DATA_DIR / "eurusd_hourly.parquet"
DAILY_PATH = DATA_DIR / "eurusd_daily.parquet"
OUTPUT_PATH = DATA_DIR / "eurusd_regime.parquet"

N_STATES = 3
RANDOM_STATE = 42


def fit_regime_hmm(realized_vol: np.ndarray, n_states: int = 3, random_state: int = 42) -> np.ndarray:
    """Fit a Gaussian HMM on realized_vol and return state labels (0=calm → n_states-1=crisis).

    States are sorted by mean realized vol ascending so label 0 is always calm,
    label n_states-1 is always crisis. This makes the labelling deterministic
    across datasets with different vol levels.

    Args:
        realized_vol:  1-D array of realized vol values
        n_states:      number of HMM hidden states (default 3)
        random_state:  random seed for reproducibility

    Returns:
        Integer label array of shape (len(realized_vol),), dtype int64
    """
    X = realized_vol.reshape(-1, 1).astype(np.float64)
    model = hmm.GaussianHMM(
        n_components=n_states,
        covariance_type="diag",
        min_covar=1e-6,
        n_iter=100,
        random_state=random_state,
        tol=1e-4,
    )
    model.fit(X)
    raw_labels = model.predict(X)

    # Sort states by mean realized vol (ascending: calm=0, crisis=n_states-1)
    state_means = np.array([X[raw_labels == s].mean() if (raw_labels == s).any() else 0.0
                            for s in range(n_states)])
    rank = np.argsort(state_means)          # rank[0] = original state id of the calmest cluster
    remap = np.empty(n_states, dtype=np.int64)
    for new_label, old_label in enumerate(rank):
        remap[old_label] = new_label
    return remap[raw_labels].astype(np.int64)


def prepare_regime_df(parquet_path: str, freq: str = "hourly") -> pd.DataFrame:
    """Load parquet, fit HMM, return DataFrame with timestamp + regime columns.

    Args:
        parquet_path: path to input parquet (hourly or daily)
        freq:         "hourly" | "daily" — determines timestamp column name

    Returns:
        DataFrame with columns: (datetime|date), regime
    """
    df = pd.read_parquet(parquet_path)
    if freq == "hourly":
        ts = pd.to_datetime(df["datetime"])
    else:
        ts = pd.to_datetime(df["date"])

    rv = df["realized_vol"].to_numpy(np.float32)
    labels = fit_regime_hmm(rv, n_states=N_STATES, random_state=RANDOM_STATE)
    return pd.DataFrame({"datetime": ts.values, "regime": labels})


def main():
    parser = argparse.ArgumentParser(description="Fit HMM regime detector")
    parser.add_argument("--hourly", action="store_true", default=True,
                        help="use hourly parquet (default)")
    parser.add_argument("--daily", action="store_true", default=False,
                        help="use daily parquet instead of hourly")
    parser.add_argument("--force", action="store_true", default=False,
                        help="overwrite existing output")
    parser.add_argument("--out", default=str(OUTPUT_PATH),
                        help="output parquet path")
    args = parser.parse_args()

    out_path = Path(args.out)
    if out_path.exists() and not args.force:
        print("regime parquet already exists:", out_path, "(use --force to recompute)")
        return

    if args.daily and DAILY_PATH.exists():
        src, freq = str(DAILY_PATH), "daily"
    elif HOURLY_PATH.exists():
        src, freq = str(HOURLY_PATH), "hourly"
    elif DAILY_PATH.exists():
        src, freq = str(DAILY_PATH), "daily"
    else:
        raise FileNotFoundError("no parquet found in data/processed/")

    print(f"fitting HMM ({N_STATES} states) on {src} ...")
    df = prepare_regime_df(src, freq=freq)

    counts = df["regime"].value_counts().sort_index()
    print("regime distribution:")
    for state, count in counts.items():
        label = {0: "calm", 1: "stressed", 2: "crisis"}.get(state, f"state{state}")
        print(f"  {state} ({label}): {count} ({100*count/len(df):.1f}%)")

    df.to_parquet(out_path, index=False)
    print("wrote:", out_path)


if __name__ == "__main__":
    main()