"""
HMM(Hidden Markov Model) 기반 시장 국면 감지

관측 변수 3개: [일간 수익률, 변동성(20일 rolling std), 거래량 변화율]
은닉 상태 3개: BULL / BEAR / SIDEWAYS

학습 데이터: KODEX 200 (069500) 최근 2년 일봉 (네이버 파이낸스 또는 DB)
모델 캐싱: ~/.cache/auto-stock/hmm_model.pkl (TTL 7일, 주 1회 재학습)

hmmlearn 미설치 시 ImportError 를 잡아 graceful degradation —
기존 규칙 기반(MA200+ADX)만 사용.
"""
import asyncio
import logging
import os
import pickle
import time
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np

logger = logging.getLogger(__name__)

# ── hmmlearn optional import ─────────────────────────────────────────────────
_HMM_AVAILABLE = False
try:
    from hmmlearn.hmm import GaussianHMM  # type: ignore
    _HMM_AVAILABLE = True
except ImportError:
    logger.info("hmmlearn 미설치 — HMM 국면 감지 비활성화 (규칙 기반만 사용)")

# ── 상수 ──────────────────────────────────────────────────────────────────────
_CACHE_DIR = Path.home() / ".cache" / "auto-stock"
_MODEL_PATH = _CACHE_DIR / "hmm_model.pkl"
_CACHE_TTL_SECONDS = 7 * 24 * 3600  # 7일

_N_STATES = 3  # BULL=0, BEAR=1, SIDEWAYS=2  (학습 후 재매핑)
_ROLLING_VOL_WINDOW = 20
_MIN_TRAIN_DAYS = 120  # 최소 학습 데이터 (거래일 기준)

# 앙상블 가중치 (초기값)
HMM_WEIGHT = 0.4
RULE_WEIGHT = 0.6

# ── 관측 데이터 구성 ──────────────────────────────────────────────────────────

def _build_observations(
    closes: List[float],
    volumes: List[float],
) -> Optional[np.ndarray]:
    """일봉 종가·거래량 → HMM 관측 행렬 (N×3).

    Features:
        0: 일간 수익률    log(close_t / close_{t-1})
        1: 변동성         20일 rolling std of 일간 수익률
        2: 거래량 변화율   log(vol_t / vol_{t-1}), vol=0 → NaN 제거

    Returns:
        np.ndarray shape (T, 3) or None if 데이터 부족
    """
    if len(closes) < _ROLLING_VOL_WINDOW + 2 or len(volumes) < _ROLLING_VOL_WINDOW + 2:
        return None

    c = np.array(closes, dtype=np.float64)
    v = np.array(volumes, dtype=np.float64)

    # 일간 수익률 (log return)
    log_ret = np.diff(np.log(np.maximum(c, 1e-8)))

    # 20일 rolling std
    rolling_std = np.full_like(log_ret, np.nan)
    for i in range(_ROLLING_VOL_WINDOW - 1, len(log_ret)):
        window = log_ret[i - _ROLLING_VOL_WINDOW + 1 : i + 1]
        rolling_std[i] = np.std(window, ddof=1)

    # 거래량 변화율 (log scale)
    v_safe = np.maximum(v[1:], 1.0)  # 0 방지
    v_prev = np.maximum(v[:-1], 1.0)
    vol_change = np.log(v_safe / v_prev)

    # rolling_std 가 NaN 인 초기 구간 제거
    valid_start = _ROLLING_VOL_WINDOW - 1
    log_ret = log_ret[valid_start:]
    rolling_std = rolling_std[valid_start:]
    vol_change = vol_change[valid_start:]

    if len(log_ret) < _MIN_TRAIN_DAYS:
        return None

    obs = np.column_stack([log_ret, rolling_std, vol_change])

    # NaN / Inf 제거
    mask = np.all(np.isfinite(obs), axis=1)
    obs = obs[mask]

    if len(obs) < _MIN_TRAIN_DAYS:
        return None

    return obs


# ── 상태 매핑 ──────────────────────────────────────────────────────────────────

def _map_states(model: "GaussianHMM") -> Dict[int, str]:
    """학습된 HMM 은닉 상태 → BULL/BEAR/SIDEWAYS 매핑.

    기준: 각 상태의 평균 수익률(means_[state][0])
      - 수익률 가장 높은 상태 → BULL
      - 수익률 가장 낮은 상태 → BEAR
      - 나머지 → SIDEWAYS
    """
    mean_returns = model.means_[:, 0]  # 각 상태의 평균 일간 수익률
    sorted_idx = np.argsort(mean_returns)  # 오름차순: [가장작은, 중간, 가장큰]

    mapping = {}
    mapping[int(sorted_idx[0])] = "bear"
    mapping[int(sorted_idx[1])] = "sideways"
    mapping[int(sorted_idx[2])] = "bull"
    return mapping


# ── 모델 학습 ──────────────────────────────────────────────────────────────────

async def _load_chart_data_for_training() -> Optional[Tuple[List[float], List[float]]]:
    """KODEX 200(069500) 최근 2년(≈500거래일) 일봉 로드.

    1차: 네이버 파이낸스 (최대 600일)
    2차: DB stock_daily_prices 폴백

    Returns:
        (closes, volumes) oldest-first, or None
    """
    # 1차: 네이버 파이낸스
    try:
        from app.services.naver_finance import get_naver_daily_chart
        chart = await get_naver_daily_chart("069500", days=600)
        if chart and len(chart) >= _MIN_TRAIN_DAYS:
            closes = [float(item.get("stck_clpr", 0)) for item in chart]
            volumes = [float(item.get("acml_vol", 0)) for item in chart]
            if all(c > 0 for c in closes[-_MIN_TRAIN_DAYS:]):
                logger.info("HMM 학습 데이터: 네이버 파이낸스 %d일", len(closes))
                return closes, volumes
    except Exception as e:
        logger.warning("HMM: 네이버 차트 로드 실패 — %s", e)

    # 2차: DB 폴백
    try:
        from app.core.database import AsyncSessionLocal
        from app.models.models import StockDailyPrice
        from sqlalchemy import select
        from datetime import date

        two_years_ago = date.today() - timedelta(days=730)
        async with AsyncSessionLocal() as session:
            stmt = (
                select(StockDailyPrice)
                .where(StockDailyPrice.stock_code == "069500")
                .where(StockDailyPrice.trade_date >= two_years_ago)
                .order_by(StockDailyPrice.trade_date.asc())
            )
            result = await session.execute(stmt)
            rows = result.scalars().all()

        if rows and len(rows) >= _MIN_TRAIN_DAYS:
            closes = [float(r.close_price) for r in rows if r.close_price]
            volumes = [float(r.volume or 0) for r in rows]
            logger.info("HMM 학습 데이터: DB stock_daily_prices %d일", len(closes))
            return closes, volumes
    except Exception as e:
        logger.warning("HMM: DB 폴백 실패 — %s", e)

    return None


def _train_hmm(obs: np.ndarray) -> Optional["GaussianHMM"]:
    """GaussianHMM 학습 (3 states, diagonal covariance).

    hmmlearn 의 ConvergenceWarning 을 무시하고,
    최대 200회 반복, random seed 고정으로 재현성 확보.
    """
    if not _HMM_AVAILABLE:
        return None

    import warnings
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        try:
            model = GaussianHMM(
                n_components=_N_STATES,
                covariance_type="diag",
                n_iter=200,
                random_state=42,
                tol=1e-4,
            )
            model.fit(obs)
            return model
        except Exception as e:
            logger.error("HMM 학습 실패: %s", e)
            return None


def _save_model(model: "GaussianHMM", state_map: Dict[int, str]) -> bool:
    """학습된 모델 + 상태 매핑을 pkl 캐시에 저장."""
    try:
        _CACHE_DIR.mkdir(parents=True, exist_ok=True)
        payload = {
            "model": model,
            "state_map": state_map,
            "trained_at": time.time(),
        }
        with open(_MODEL_PATH, "wb") as f:
            pickle.dump(payload, f)
        logger.info("HMM 모델 저장 완료: %s", _MODEL_PATH)
        return True
    except Exception as e:
        logger.error("HMM 모델 저장 실패: %s", e)
        return False


def _load_model() -> Optional[Dict]:
    """캐시에서 모델 로드. TTL 7일 초과 시 None 반환."""
    if not _MODEL_PATH.exists():
        return None
    try:
        with open(_MODEL_PATH, "rb") as f:
            payload = pickle.load(f)
        trained_at = payload.get("trained_at", 0)
        if time.time() - trained_at > _CACHE_TTL_SECONDS:
            logger.info("HMM 모델 캐시 만료 (%.1f일 경과)", (time.time() - trained_at) / 86400)
            return None
        return payload
    except Exception as e:
        logger.warning("HMM 모델 로드 실패: %s", e)
        return None


# ── 공개 API ──────────────────────────────────────────────────────────────────

async def fit_regime_model(force: bool = False) -> bool:
    """HMM 시장 국면 모델 학습.

    Args:
        force: True 이면 캐시 TTL 무시하고 재학습

    Returns:
        True 학습 성공, False 실패
    """
    if not _HMM_AVAILABLE:
        logger.info("hmmlearn 미설치 — HMM 학습 스킵")
        return False

    # 캐시 TTL 확인
    if not force:
        cached = _load_model()
        if cached is not None:
            logger.info("HMM 모델 캐시 유효 — 재학습 불필요")
            return True

    # 데이터 로드
    data = await _load_chart_data_for_training()
    if data is None:
        logger.warning("HMM: 학습 데이터 부족")
        return False

    closes, volumes = data

    # 관측 행렬 구성
    obs = _build_observations(closes, volumes)
    if obs is None:
        logger.warning("HMM: 관측 데이터 구성 실패")
        return False

    # 학습 (CPU-bound → thread pool)
    model = await asyncio.to_thread(_train_hmm, obs)
    if model is None:
        return False

    # 상태 매핑 결정
    state_map = _map_states(model)
    logger.info(
        "HMM 상태 매핑: %s  (평균 수익률: %s)",
        state_map,
        {k: f"{model.means_[k, 0]:.6f}" for k in state_map},
    )

    # 저장
    return await asyncio.to_thread(_save_model, model, state_map)


async def predict_current_regime(
    closes: Optional[List[float]] = None,
    volumes: Optional[List[float]] = None,
) -> Optional[Dict]:
    """현재 시장 국면 HMM 예측.

    Args:
        closes: 최근 일봉 종가 (oldest first, 최소 40일)
        volumes: 최근 일봉 거래량 (oldest first, 같은 길이)
        미제공 시 네이버/DB 에서 자동 로드

    Returns:
        {
            "regime": "bull" | "bear" | "sideways",
            "confidence": float (0.0~1.0),
            "transition_prob": {"bull": float, "bear": float, "sideways": float},
            "state_sequence_last5": [str, ...],
        }
        or None (모델 없음/예측 실패)
    """
    if not _HMM_AVAILABLE:
        return None

    # 모델 로드
    cached = _load_model()
    if cached is None:
        # 자동 학습 시도
        success = await fit_regime_model()
        if not success:
            return None
        cached = _load_model()
        if cached is None:
            return None

    model: "GaussianHMM" = cached["model"]
    state_map: Dict[int, str] = cached["state_map"]

    # 관측 데이터 준비
    if closes is None or volumes is None:
        data = await _load_chart_data_for_training()
        if data is None:
            return None
        closes, volumes = data

    obs = _build_observations(closes, volumes)
    if obs is None:
        return None

    try:
        # Viterbi 디코딩으로 최적 상태 시퀀스
        _, state_seq = model.decode(obs, algorithm="viterbi")

        # 현재 상태
        current_state = int(state_seq[-1])
        current_regime = state_map.get(current_state, "sideways")

        # 사후 확률 (마지막 관측에 대한 각 상태 확률)
        posteriors = model.predict_proba(obs)
        last_prob = posteriors[-1]  # shape: (n_states,)

        # confidence = 현재 상태의 사후 확률
        confidence = float(last_prob[current_state])

        # 전이 확률: 현재 상태에서 다음 상태로의 전이 확률
        trans_row = model.transmat_[current_state]
        transition_prob = {}
        for state_idx, regime_name in state_map.items():
            transition_prob[regime_name] = round(float(trans_row[state_idx]), 4)

        # 최근 5일 상태 시퀀스
        last_5 = [state_map.get(int(s), "sideways") for s in state_seq[-5:]]

        return {
            "regime": current_regime,
            "confidence": round(confidence, 4),
            "transition_prob": transition_prob,
            "state_sequence_last5": last_5,
        }
    except Exception as e:
        logger.error("HMM 예측 실패: %s", e)
        return None


def ensemble_regime(
    hmm_result: Optional[Dict],
    rule_regime: str,
) -> Dict:
    """HMM 결과와 규칙 기반 결과를 앙상블.

    가중 평균 방식:
      - HMM weight 0.4, Rule weight 0.6 (초기값)
      - 두 모델 일치 시 confidence 높음
      - 불일치 시 보수적 판정 (SIDEWAYS)

    Args:
        hmm_result: predict_current_regime() 반환값 (None 가능)
        rule_regime: 기존 detect_market_regime() 결과
                     ("bull", "bear", "sideways")

    Returns:
        {
            "regime": str,
            "confidence": float,
            "hmm_regime": str | None,
            "rule_regime": str,
            "method": "ensemble" | "rule_only",
        }
    """
    rule_lower = rule_regime.lower()

    # HMM 결과 없으면 규칙 기반만 사용
    if hmm_result is None:
        return {
            "regime": rule_lower,
            "confidence": 0.6,  # 규칙 기반 기본 신뢰도
            "hmm_regime": None,
            "rule_regime": rule_lower,
            "method": "rule_only",
        }

    hmm_regime = hmm_result["regime"]
    hmm_conf = hmm_result["confidence"]

    # 일치 여부 판정
    if hmm_regime == rule_lower:
        # 두 모델 일치 → 높은 confidence
        combined_conf = min(1.0, RULE_WEIGHT * 0.9 + HMM_WEIGHT * hmm_conf)
        return {
            "regime": rule_lower,
            "confidence": round(combined_conf, 4),
            "hmm_regime": hmm_regime,
            "rule_regime": rule_lower,
            "method": "ensemble",
        }

    # 불일치: 보수적 판정
    # CASE 1: 한쪽이 sideways → sideways
    if hmm_regime == "sideways" or rule_lower == "sideways":
        return {
            "regime": "sideways",
            "confidence": round(0.5 * hmm_conf, 4),
            "hmm_regime": hmm_regime,
            "rule_regime": rule_lower,
            "method": "ensemble",
        }

    # CASE 2: bull vs bear → HMM confidence 높으면 HMM, 아니면 규칙 기반
    # 단, confidence 낮으면 sideways 로 보수적 처리
    if hmm_conf >= 0.7:
        # HMM 이 확신 높으면 HMM 에 가중치 부여
        # 그래도 RULE_WEIGHT 가 0.6 이므로 규칙 기반 우선
        final_regime = rule_lower
        combined_conf = round(RULE_WEIGHT * 0.5, 4)  # 낮은 신뢰도
    else:
        # 불확실 → sideways (보수적)
        final_regime = "sideways"
        combined_conf = round(0.3 + 0.2 * hmm_conf, 4)

    return {
        "regime": final_regime,
        "confidence": combined_conf,
        "hmm_regime": hmm_regime,
        "rule_regime": rule_lower,
        "method": "ensemble",
    }


# ── 모듈 상태 확인 ────────────────────────────────────────────────────────────

def is_hmm_available() -> bool:
    """hmmlearn 라이브러리 설치 여부."""
    return _HMM_AVAILABLE


def get_model_info() -> Dict:
    """캐시된 모델 메타 정보 반환."""
    if not _MODEL_PATH.exists():
        return {"status": "no_model", "available": _HMM_AVAILABLE}

    cached = _load_model()
    if cached is None:
        return {"status": "expired", "available": _HMM_AVAILABLE, "path": str(_MODEL_PATH)}

    trained_at = cached.get("trained_at", 0)
    age_days = (time.time() - trained_at) / 86400
    return {
        "status": "ready",
        "available": _HMM_AVAILABLE,
        "path": str(_MODEL_PATH),
        "trained_at": datetime.fromtimestamp(trained_at).isoformat(),
        "age_days": round(age_days, 1),
        "ttl_remaining_days": round(7 - age_days, 1),
        "state_map": cached.get("state_map", {}),
    }