import pandas as pd
import os
from hmmlearn.hmm import GaussianHMM
import numpy as np

FEATURE_DIR = "/opt/airflow/data/features"
OUTPUT_DIR = "/opt/airflow/data/regimes"

def classify_market_regime(asset: str):

    file_path = os.path.join(FEATURE_DIR, f"{asset}_features.parquet")
    df = pd.read_parquet(file_path)

    # --- Добавляем новые фичи ---
    df["trend_50"] = df["close"] / df["close"].rolling(50).mean() - 1
    df["trend_200"] = df["close"] / df["close"].rolling(200).mean() - 1
    df["momentum"] = df["close"].pct_change(20)

    feature_cols = [
        "return",
        "volatility",
        "trend_50",
        "trend_200",
        "momentum"
    ]

    df = df.dropna().copy()

    X = df[feature_cols].values

    # --- HMM ---
    model = GaussianHMM(
        n_components=3,
        covariance_type="full",
        n_iter=1000,
        random_state=42
    )

    model.fit(X)

    hidden_states = model.predict(X)
    df["market_regime"] = hidden_states

    # --- Интерпретация состояний ---
    state_summary = (
        df.groupby("market_regime")[["return", "volatility"]]
        .mean()
        .sort_values("return", ascending=False)
    )

    sorted_states = state_summary.index.tolist()

    regime_mapping = {
        sorted_states[0]: "Bull",
        sorted_states[1]: "Sideways",
        sorted_states[2]: "Bear",
    }

    df["market_regime_label"] = df["market_regime"].map(regime_mapping)

    os.makedirs(OUTPUT_DIR, exist_ok=True)
    output_file = os.path.join(OUTPUT_DIR, f"{asset}_regimes.parquet")

    df.to_parquet(output_file, index=False)

    print(f"Market regimes saved to {output_file}")


if __name__ == "__main__":
    classify_market_regime("BTC-USD")