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

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 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"]


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

    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)

    # 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)

    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
    """
    if not chart_data or len(chart_data) < 61:
        return None, None

    df = chart_data_to_dataframe(chart_data)
    features_df = extract_features(df)

    # 마지막 행의 탭뷸러 피처
    last_row = features_df[FEATURE_COLUMNS].iloc[-1].values
    if np.any(np.isnan(last_row)):
        return None, None

    # 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) < 60:
        return None, None

    sequence = seq_raw[-60:]  # (60, 5)

    return last_row, sequence