"""
피처 엔지니어링 모듈

KIS API 차트 데이터(OHLCV)에서 ML 모델용 피처를 추출합니다.
- 기본 피처 12개: RSI, MACD, 스토캐스틱, 볼린저, MA비율, 거래량비율
- 파생 피처 8개: 수익률, 변동성, 모멘텀, 크로스, ADX, 이격도
"""

import numpy as np
import pandas as pd
from typing import Optional, Dict, List, Tuple


def _ema(values: np.ndarray, period: int) -> np.ndarray:
    """지수이동평균 계산"""
    result = np.full_like(values, np.nan, dtype=float)
    if len(values) < period:
        return result
    result[period - 1] = np.mean(values[:period])
    multiplier = 2.0 / (period + 1)
    for i in range(period, len(values)):
        result[i] = (values[i] - result[i - 1]) * multiplier + result[i - 1]
    return result


def chart_data_to_dataframe(chart_data: list) -> pd.DataFrame:
    """KIS API 차트 데이터(최신순)를 DataFrame(시간순)으로 변환"""
    rows = []
    for item in chart_data:
        rows.append({
            "date": item.get("stck_bsop_date", ""),
            "open": float(item.get("stck_oprc", 0) or 0),
            "high": float(item.get("stck_hgpr", 0) or 0),
            "low": float(item.get("stck_lwpr", 0) or 0),
            "close": float(item.get("stck_clpr", 0) or 0),
            "volume": float(item.get("acml_vol", 0) or 0),
        })
    df = pd.DataFrame(rows)
    # KIS API는 최신순 → 시간순으로 뒤집기
    df = df.iloc[::-1].reset_index(drop=True)
    return df


def chart_data_to_dataframe_v2(chart_data: list) -> pd.DataFrame:
    """KIS 형식 차트 데이터를 **입력 순서와 무관하게** 시간순 DataFrame으로 변환.

    chart_data_to_dataframe() 은 "입력이 최신순"이라는 가정 아래 무조건 뒤집는다.
    그런데 market_analyzer 는 oldest-first 로 정규화한 배열을 넘기므로
    (analyze_stock: `list(reversed(kis_output2))`, naver: oldest-first)
    그 함수를 그대로 쓰면 **시계열이 역방향으로 뒤집힌 상태**로 피처가 계산된다.
    학습 데이터(db_data_to_dataframe → trade_date ASC)는 정방향이므로
    추론이 학습과 정반대 방향의 시계열을 보게 되는 결함이었다. (2026-09-28)

    stck_bsop_date 를 기준으로 오름차순 정렬하여 방향 의존성을 제거한다.
    날짜가 없으면 기존 동작(최신순 가정 → reverse)으로 폴백한다.
    """
    if not chart_data:
        return chart_data_to_dataframe(chart_data)

    dates = [str(item.get("stck_bsop_date", "") or "") for item in chart_data]
    if not all(len(d) == 8 and d.isdigit() for d in dates):
        # 날짜 누락/형식 이상 → 기존 규약(최신순 입력) 유지
        return chart_data_to_dataframe(chart_data)

    order = sorted(range(len(chart_data)), key=lambda i: dates[i])
    ordered = [chart_data[i] for i in order]
    # chart_data_to_dataframe 은 뒤집으므로, 정렬 결과를 한 번 더 뒤집어 넘긴다
    return chart_data_to_dataframe(list(reversed(ordered)))


def db_data_to_dataframe(rows: list) -> pd.DataFrame:
    """DB StockDailyPrice 레코드를 DataFrame(시간순)으로 변환.

    Args:
        rows: StockDailyPrice 객체 리스트 또는 dict 리스트
              (trade_date, open_price, high_price, low_price, close_price, volume)

    Returns:
        chart_data_to_dataframe()과 동일한 형태의 DataFrame
        columns: date, open, high, low, close, volume
    """
    records = []
    for row in rows:
        if isinstance(row, dict):
            records.append({
                "date": str(row["trade_date"]),
                "open": float(row["open_price"]),
                "high": float(row["high_price"]),
                "low": float(row["low_price"]),
                "close": float(row["close_price"]),
                "volume": float(row["volume"]),
            })
        else:
            records.append({
                "date": str(row.trade_date),
                "open": float(row.open_price),
                "high": float(row.high_price),
                "low": float(row.low_price),
                "close": float(row.close_price),
                "volume": float(row.volume),
            })
    df = pd.DataFrame(records)
    # trade_date 기준 오름차순 정렬
    df = df.sort_values("date").reset_index(drop=True)
    return df


def compute_rsi(close: np.ndarray, period: int = 14) -> np.ndarray:
    """RSI 계산"""
    result = np.full_like(close, np.nan, dtype=float)
    if len(close) < period + 1:
        return result
    delta = np.diff(close)
    gain = np.where(delta > 0, delta, 0.0)
    loss = np.where(delta < 0, -delta, 0.0)

    avg_gain = np.mean(gain[:period])
    avg_loss = np.mean(loss[:period])

    for i in range(period, len(delta)):
        avg_gain = (avg_gain * (period - 1) + gain[i]) / period
        avg_loss = (avg_loss * (period - 1) + loss[i]) / period
        if avg_loss == 0:
            result[i + 1] = 100.0
        else:
            rs = avg_gain / avg_loss
            result[i + 1] = 100.0 - (100.0 / (1.0 + rs))
    # 첫 번째 유효 값
    if avg_loss == 0:
        result[period] = 100.0
    else:
        result[period] = 100.0 - (100.0 / (1.0 + np.mean(gain[:period]) / max(np.mean(loss[:period]), 1e-10)))
    return result


def extract_features(df: pd.DataFrame) -> pd.DataFrame:
    """DataFrame(OHLCV 시간순)에서 전체 피처를 계산하여 반환.

    반환되는 DataFrame은 입력과 같은 인덱스를 갖되, 각 행마다
    해당 시점까지의 정보만을 사용하여 계산된 20개 피처 컬럼을 포함합니다.
    초기 데이터 부족 구간은 NaN이 됩니다.
    """
    close = df["close"].values.astype(float)
    high = df["high"].values.astype(float)
    low = df["low"].values.astype(float)
    volume = df["volume"].values.astype(float)
    n = len(close)

    features: Dict[str, np.ndarray] = {}

    # --- 기본 피처 (12개) ---

    # 1-2. RSI (14일, 7일)
    features["rsi_14"] = compute_rsi(close, 14)
    features["rsi_7"] = compute_rsi(close, 7)

    # 3-5. MACD (12, 26, signal=9)
    ema12 = _ema(close, 12)
    ema26 = _ema(close, 26)
    macd_line = ema12 - ema26
    macd_signal = _ema(macd_line[~np.isnan(macd_line)], 9)
    # macd_signal을 원래 길이에 맞추기
    full_signal = np.full(n, np.nan, dtype=float)
    valid_start = np.argmax(~np.isnan(macd_line))
    if len(macd_signal) > 0:
        full_signal[valid_start:valid_start + len(macd_signal)] = macd_signal
    features["macd"] = macd_line
    features["macd_signal"] = full_signal
    features["macd_hist"] = macd_line - full_signal

    # 6-7. 스토캐스틱 %K, %D (14일)
    stoch_k = np.full(n, np.nan, dtype=float)
    stoch_d = np.full(n, np.nan, dtype=float)
    for i in range(13, n):
        h14 = np.max(high[i - 13:i + 1])
        l14 = np.min(low[i - 13:i + 1])
        if h14 != l14:
            stoch_k[i] = (close[i] - l14) / (h14 - l14) * 100.0
    # %D = %K의 3일 SMA
    for i in range(15, n):
        vals = stoch_k[i - 2:i + 1]
        if not np.any(np.isnan(vals)):
            stoch_d[i] = np.mean(vals)
    features["stochastic_k"] = stoch_k
    features["stochastic_d"] = stoch_d

    # 8. 볼린저 밴드 위치 (0~1)
    bb_pos = np.full(n, np.nan, dtype=float)
    for i in range(19, n):
        window = close[i - 19:i + 1]
        mid = np.mean(window)
        std = np.std(window, ddof=0)
        if std > 0:
            upper = mid + 2 * std
            lower = mid - 2 * std
            bb_pos[i] = (close[i] - lower) / (upper - lower)
    features["bb_position"] = bb_pos

    # 9-11. MA 비율 (현재가 / MA)
    for period, name in [(5, "ma5_ratio"), (20, "ma20_ratio"), (60, "ma60_ratio")]:
        ratio = np.full(n, np.nan, dtype=float)
        for i in range(period - 1, n):
            ma = np.mean(close[i - period + 1:i + 1])
            if ma > 0:
                ratio[i] = close[i] / ma
        features[name] = ratio

    # 12. 거래량 비율 (현재 / 20일 평균)
    vol_ratio = np.full(n, np.nan, dtype=float)
    for i in range(19, n):
        avg_vol = np.mean(volume[i - 19:i + 1])
        if avg_vol > 0:
            vol_ratio[i] = volume[i] / avg_vol
    features["volume_ratio"] = vol_ratio

    # --- 파생 피처 (8개) ---

    # 13-15. 단기 수익률
    for days, name in [(1, "change_rate_1d"), (3, "change_rate_3d"), (5, "change_rate_5d")]:
        cr = np.full(n, np.nan, dtype=float)
        for i in range(days, n):
            if close[i - days] > 0:
                cr[i] = (close[i] - close[i - days]) / close[i - days]
        features[name] = cr

    # 16. 20일 변동성 (수익률의 표준편차)
    vol20 = np.full(n, np.nan, dtype=float)
    returns = np.full(n, np.nan, dtype=float)
    for i in range(1, n):
        if close[i - 1] > 0:
            returns[i] = (close[i] - close[i - 1]) / close[i - 1]
    for i in range(20, n):
        window = returns[i - 19:i + 1]
        valid = window[~np.isnan(window)]
        if len(valid) >= 10:
            vol20[i] = np.std(valid, ddof=1)
    features["volatility_20d"] = vol20

    # 17. 5일 모멘텀 (현재가 / 5일전가)
    mom = np.full(n, np.nan, dtype=float)
    for i in range(5, n):
        if close[i - 5] > 0:
            mom[i] = close[i] / close[i - 5]
    features["price_momentum"] = mom

    # 18. MA 크로스 상태 (-1, 0, 1)
    ma_cross = np.full(n, 0.0, dtype=float)
    for i in range(59, n):
        ma5 = np.mean(close[i - 4:i + 1])
        ma20 = np.mean(close[i - 19:i + 1])
        ma60 = np.mean(close[i - 59:i + 1])
        if ma5 > ma20 > ma60:
            ma_cross[i] = 1.0  # 정배열
        elif ma5 < ma20 < ma60:
            ma_cross[i] = -1.0  # 역배열
    features["ma_cross"] = ma_cross

    # 19. 추세 강도 (간이 ADX, True Range 기반)
    adx_val = np.full(n, np.nan, dtype=float)
    tr = np.zeros(n, dtype=float)
    for i in range(1, n):
        tr[i] = max(high[i] - low[i], abs(high[i] - close[i - 1]), abs(low[i] - close[i - 1]))
    for i in range(13, n):
        avg_tr = np.mean(tr[i - 13:i + 1])
        current_tr = tr[i]
        if avg_tr > 0:
            adx_val[i] = current_tr / avg_tr
    features["trend_strength"] = adx_val

    # 20. 이격도 (현재가 / 20일 이평선 * 100)
    disp = np.full(n, np.nan, dtype=float)
    for i in range(19, n):
        ma20 = np.mean(close[i - 19:i + 1])
        if ma20 > 0:
            disp[i] = close[i] / ma20 * 100.0
    features["disparity_20"] = disp

    feat_df = pd.DataFrame(features, index=df.index)
    return feat_df


# 피처 컬럼 순서 (모델 학습/추론 시 일관성 유지)
FEATURE_COLUMNS = [
    "rsi_14", "rsi_7",
    "macd", "macd_signal", "macd_hist",
    "stochastic_k", "stochastic_d",
    "bb_position",
    "ma5_ratio", "ma20_ratio", "ma60_ratio",
    "volume_ratio",
    "change_rate_1d", "change_rate_3d", "change_rate_5d",
    "volatility_20d",
    "price_momentum",
    "ma_cross",
    "trend_strength",
    "disparity_20",
]

# LSTM 시퀀스에 사용할 피처 (5개)
LSTM_SEQUENCE_FEATURES = ["close", "volume", "rsi_14", "ma20_ratio", "volatility_20d"]

# LSTM 시퀀스 길이
LSTM_SEQUENCE_LENGTH = 60

# ─── 워밍업 상수 (extract_features() 구현과 1:1 대응, 실측값) ───
# 각 피처가 처음으로 non-NaN 이 되는 인덱스:
#   rsi_14=14, macd=25, macd_signal/hist=33, ma60_ratio=59,
#   ma20_ratio/bb_position/volume_ratio/disparity_20=19, volatility_20d=20
# → 탭뷸러 마지막 행이 전부 유효하려면 인덱스 59 이상 (봉 60개)
TABULAR_WARMUP_BARS = 60
# → LSTM 시퀀스(마지막 60행)가 전부 유효하려면 시퀀스 시작 인덱스 >= 20
#   (volatility_20d 의 첫 유효 인덱스)
SEQUENCE_FEATURE_WARMUP_BARS = 20

# predict() 에 필요한 최소 봉수 = 시퀀스 60 + 시퀀스 피처 워밍업 20 = 80
# (61봉이면 탭뷸러는 통과하지만 시퀀스 앞쪽 20행이 NaN → 모델이 NaN 을 전파해
#  `int(prob_up*100)` 에서 ValueError. 2026-09-23~28 ML 전면 미작동 원인)
MIN_BARS_FOR_PREDICT = LSTM_SEQUENCE_LENGTH + SEQUENCE_FEATURE_WARMUP_BARS
# 권장 봉수: RSI(Wilder)·EMA 등 재귀 지표의 번인을 감안한 여유치
RECOMMENDED_BARS_FOR_PREDICT = 120


def prepare_training_data(
    df: pd.DataFrame,
    features_df: pd.DataFrame,
    forecast_days: int = 3,
    up_threshold: float = 0.015,
    down_threshold: float = -0.015,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """학습 데이터 준비.

    Returns:
        X_tabular: (N, 20) 탭뷸러 피처
        X_seq: (N, 60, 5) LSTM 시퀀스 피처
        y: (N,) 레이블 (0=DOWN, 1=HOLD, 2=UP)
        valid_indices: 유효한 행 인덱스
    """
    close = df["close"].values.astype(float)
    n = len(close)

    # 레이블 생성: forecast_days 후 수익률
    labels = np.full(n, np.nan, dtype=float)
    for i in range(n - forecast_days):
        if close[i] > 0:
            future_return = (close[i + forecast_days] - close[i]) / close[i]
            if future_return >= up_threshold:
                labels[i] = 2  # UP
            elif future_return <= down_threshold:
                labels[i] = 0  # DOWN
            else:
                labels[i] = 1  # HOLD

    # 유효한 행 찾기 (피처와 레이블 모두 NaN 아닌 행)
    tabular = features_df[FEATURE_COLUMNS].values
    seq_window = 60

    valid_mask = np.ones(n, dtype=bool)
    valid_mask[:seq_window] = False  # 시퀀스 윈도우 확보
    valid_mask[np.isnan(labels)] = False
    # 탭뷸러 피처 NaN 체크
    for i in range(n):
        if valid_mask[i] and np.any(np.isnan(tabular[i])):
            valid_mask[i] = False

    # LSTM 시퀀스 피처 준비
    # 시퀀스에 사용할 raw 피처 결합
    seq_cols = []
    for col in LSTM_SEQUENCE_FEATURES:
        if col in df.columns:
            seq_cols.append(df[col].values.astype(float))
        elif col in features_df.columns:
            seq_cols.append(features_df[col].values.astype(float))
        else:
            seq_cols.append(np.zeros(n, dtype=float))
    seq_raw = np.column_stack(seq_cols)  # (n, 5)

    # [2026-09-28] 시퀀스 윈도우 NaN 체크 — 기존에는 `valid_mask[:60]=False` 만 있어
    # 인덱스 60~79 샘플의 시퀀스(seq_raw[idx-60:idx])에 워밍업 NaN 이 섞여 들어갔다.
    # X_seq 에 NaN 이 하나라도 있으면 `X_seq.mean(axis=(0,1))` = NaN 이 되어
    # seq_mean/seq_std 가 NaN → LSTM 이 학습·추론 전체에서 NaN 을 뱉는다.
    seq_row_finite = np.all(np.isfinite(seq_raw), axis=1)
    # 누적합으로 각 윈도우의 유효 행 수를 O(n) 에 계산
    finite_cumsum = np.concatenate(([0], np.cumsum(seq_row_finite.astype(int))))
    for i in range(n):
        if not valid_mask[i]:
            continue
        start = i - seq_window
        if start < 0 or (finite_cumsum[i] - finite_cumsum[start]) < seq_window:
            valid_mask[i] = False

    valid_indices = np.where(valid_mask)[0]

    if len(valid_indices) == 0:
        return np.array([]), np.array([]), np.array([]), np.array([])

    # 탭뷸러 피처
    X_tabular = tabular[valid_indices]

    # 레이블
    y = labels[valid_indices].astype(int)

    X_seq = np.zeros((len(valid_indices), seq_window, len(LSTM_SEQUENCE_FEATURES)), dtype=float)
    for j, idx in enumerate(valid_indices):
        X_seq[j] = seq_raw[idx - seq_window:idx]

    return X_tabular, X_seq, y, valid_indices


def extract_single_features(chart_data: list) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
    """단일 종목 차트 데이터에서 최신 시점의 피처를 추출.

    Returns:
        tabular_features: (20,) 탭뷸러 피처 또는 None
        sequence_features: (60, 5) LSTM 시퀀스 또는 None

    실패 사유가 필요하면 extract_single_features_v2() 를 사용하세요.
    """
    tabular, sequence, _reason = extract_single_features_v2(chart_data)
    return tabular, sequence


def extract_single_features_v2(
    chart_data: list,
) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[str]]:
    """extract_single_features() + 실패 사유 반환 (2026-09-28 신설).

    기존 버전의 결함 두 가지를 여기서 바로잡는다.

    1) **최소 봉수 61 → 80**.
       61봉이면 탭뷸러 마지막 행은 유효하지만(워밍업 60) LSTM 시퀴스
       `seq_raw[-60:]` 의 앞쪽 행에 NaN 이 남는다
       (volatility_20d 첫 유효 인덱스 20, ma20_ratio 19, rsi_14 14).
       그 NaN 은 예외를 던지지 않고 LSTM → 가중평균 확률로 **전파**되어
       `int(prob_up * 100)` 에서 `ValueError: cannot convert float NaN to integer`
       를 유발했다. 봉수를 워밍업에 맞춰 올리는 것이 근본 수정이다.

    2) **입력 순서 의존 제거**. chart_data_to_dataframe() 은 입력이 최신순이라
       가정해 무조건 뒤집는데, 호출부(market_analyzer)는 oldest-first 를 넘긴다.
       → 시계열이 역방향으로 뒤집힌 상태로 피처가 계산되고 있었다.
       chart_data_to_dataframe_v2() 로 날짜 기준 정렬한다.

    Returns:
        (tabular(20,), sequence(60,5), None) 또는 (None, None, 실패사유)
    """
    if not chart_data:
        return None, None, "차트 데이터 없음"
    if len(chart_data) < MIN_BARS_FOR_PREDICT:
        return None, None, (
            f"봉수 부족 {len(chart_data)}/{MIN_BARS_FOR_PREDICT}"
            f"(시퀀스{LSTM_SEQUENCE_LENGTH}+워밍업{SEQUENCE_FEATURE_WARMUP_BARS})"
        )

    df = chart_data_to_dataframe_v2(chart_data)
    if len(df) < MIN_BARS_FOR_PREDICT:
        return None, None, f"변환 후 봉수 부족 {len(df)}/{MIN_BARS_FOR_PREDICT}"
    features_df = extract_features(df)

    # 마지막 행의 탭뷸러 피처
    last_row = features_df[FEATURE_COLUMNS].iloc[-1].values.astype(float)
    if not np.all(np.isfinite(last_row)):
        bad = [c for c, v in zip(FEATURE_COLUMNS, last_row) if not np.isfinite(v)]
        return None, None, f"탭뷸러 피처 무효: {','.join(bad)}"

    # LSTM 시퀀스 (최근 60일)
    seq_cols = []
    for col in LSTM_SEQUENCE_FEATURES:
        if col in df.columns:
            seq_cols.append(df[col].values.astype(float))
        elif col in features_df.columns:
            seq_cols.append(features_df[col].values.astype(float))
        else:
            seq_cols.append(np.zeros(len(df), dtype=float))
    seq_raw = np.column_stack(seq_cols)

    if len(seq_raw) < LSTM_SEQUENCE_LENGTH:
        return None, None, f"시퀀스 길이 부족 {len(seq_raw)}/{LSTM_SEQUENCE_LENGTH}"

    sequence = seq_raw[-LSTM_SEQUENCE_LENGTH:].astype(float)  # (60, 5)

    # 최소 봉수를 만족해도 데이터에 결측/0가격이 섞이면 NaN 이 남을 수 있다.
    # 여기서 반드시 잡아야 모델이 NaN 을 전파하지 않는다.
    if not np.all(np.isfinite(sequence)):
        bad_idx = np.where(~np.isfinite(sequence).all(axis=0))[0]
        bad = [LSTM_SEQUENCE_FEATURES[i] for i in bad_idx]
        return None, None, f"시퀀스 피처 무효: {','.join(bad)}"

    return last_row, sequence, None