trading/graph-engine/main.py

"""
Graph Engine (port 9090, 172.30.0.25)
한국 종목 그래프 신경망 (GAT) — 다음날 수익률 예측 → stock_scores.graph_score

노드: 한국 활성종목 (dart_corps.is_active=true)
피처(12): 1d/5d/20d 수익률, vol_ratio, rsi, tech_score,
         roe, operating_margin, debt_ratio, news_7d, us_overnight, log_mcap
엣지: ① 가격 60일 상관 |corr|>0.4   ② 같은 sector   ③ 뉴스 공기 ≥3회

학습: 매주 일요일 06:00 (6mo rolling window)
추론: 매일 08:30 → stock_scores.graph_score
"""
import os
import asyncio
import json
import math
from datetime import date, datetime, timedelta
from typing import Optional, List, Tuple

import asyncpg
import structlog
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from fastapi import FastAPI, Query, BackgroundTasks
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger
from pytz import timezone

# ─────────────────────────────────────────────────────────────
# 설정
# ─────────────────────────────────────────────────────────────
PG = {
    "host": os.getenv("POSTGRES_HOST", "postgres"),
    "port": int(os.getenv("POSTGRES_PORT", 5432)),
    "database": os.getenv("POSTGRES_DB", "trading_ai"),
    "user": os.getenv("POSTGRES_USER", "kyu"),
    "password": os.getenv("POSTGRES_PASSWORD", ""),
}
KST = timezone("Asia/Seoul")
MODEL_DIR = os.getenv("GRAPH_MODEL_DIR", "/mnt/nas/models/graph")
os.makedirs(MODEL_DIR, exist_ok=True)
MODEL_PATH = os.path.join(MODEL_DIR, "gat_latest.pt")

FEATURE_DIM = 12
HIDDEN_DIM = 32
ATT_HEADS = 4
DROPOUT = 0.3
CORR_LOOKBACK = 60
CORR_THRESHOLD = 0.65
TOP_K_NEIGHBORS = 20      # 노드당 가격 상관 엣지 상한 (메모리/속도 캡)
NEWS_LOOKBACK = 14
NEWS_COOC_MIN = 3
TRAIN_WINDOW_DAYS = 180
SAMPLE_STRIDE_DAYS = 5    # 6개월 / 5일 = ~36 학습 샘플

logger = structlog.get_logger()
app = FastAPI(title="Graph Engine (GAT)")
pg_pool: Optional[asyncpg.Pool] = None
scheduler = AsyncIOScheduler(timezone=KST)
device = torch.device("cpu")

DDL = """
ALTER TABLE stock_scores
  ADD COLUMN IF NOT EXISTS graph_score DOUBLE PRECISION;
CREATE TABLE IF NOT EXISTS graph_model_meta (
  model_date    DATE PRIMARY KEY,
  train_samples INT,
  val_samples   INT,
  val_loss      DOUBLE PRECISION,
  edge_count    INT,
  node_count    INT,
  hidden_dim    INT,
  notes         TEXT,
  created_at    TIMESTAMP DEFAULT NOW()
);
CREATE TABLE IF NOT EXISTS graph_predictions (
  stock_code   VARCHAR(10),
  predict_date DATE,
  pred_return  DOUBLE PRECISION,
  created_at   TIMESTAMP DEFAULT NOW(),
  PRIMARY KEY (stock_code, predict_date)
);
"""


# ─────────────────────────────────────────────────────────────
# GAT 모델 (torch-geometric 없이 순수 torch)
# ─────────────────────────────────────────────────────────────
class GATLayer(nn.Module):
    def __init__(self, in_dim, out_dim, heads, dropout):
        super().__init__()
        assert out_dim % heads == 0
        self.heads = heads
        self.head_dim = out_dim // heads
        self.W = nn.Linear(in_dim, out_dim, bias=False)
        a_src = torch.empty(1, heads, self.head_dim)
        a_dst = torch.empty(1, heads, self.head_dim)
        nn.init.xavier_uniform_(a_src)
        nn.init.xavier_uniform_(a_dst)
        self.a_src = nn.Parameter(a_src.squeeze(0))
        self.a_dst = nn.Parameter(a_dst.squeeze(0))
        nn.init.xavier_uniform_(self.W.weight)
        self.dropout = dropout

    def forward(self, x, edge_index, edge_weight=None):
        N = x.size(0)
        h = self.W(x).view(N, self.heads, self.head_dim)
        src, dst = edge_index[0], edge_index[1]
        # 어텐션 logits per (edge, head)
        e_src = (h[src] * self.a_src.unsqueeze(0)).sum(-1)
        e_dst = (h[dst] * self.a_dst.unsqueeze(0)).sum(-1)
        e = F.leaky_relu(e_src + e_dst, 0.2)
        if edge_weight is not None:
            e = e + edge_weight.unsqueeze(-1).log().clamp(min=-10)
        # dst별 softmax: max-subtract for stability
        e_max = torch.full((N, self.heads), -1e9, device=x.device)
        e_max = e_max.scatter_reduce(0, dst.unsqueeze(-1).expand(-1, self.heads),
                                      e, reduce="amax", include_self=True)
        e_exp = torch.exp(e - e_max[dst])
        denom = torch.zeros(N, self.heads, device=x.device).index_add_(0, dst, e_exp)
        alpha = e_exp / (denom[dst] + 1e-12)
        alpha = F.dropout(alpha, p=self.dropout, training=self.training)
        # 메시지 집계
        m = h[src] * alpha.unsqueeze(-1)
        out = torch.zeros(N, self.heads, self.head_dim, device=x.device)
        out = out.index_add_(0, dst, m)
        return out.reshape(N, -1)


class GraphNet(nn.Module):
    def __init__(self, in_dim=FEATURE_DIM, hidden=HIDDEN_DIM,
                 heads=ATT_HEADS, dropout=DROPOUT):
        super().__init__()
        self.gat1 = GATLayer(in_dim, hidden, heads, dropout)
        self.gat2 = GATLayer(hidden, hidden, heads, dropout)
        self.head = nn.Linear(hidden, 1)
        self.dropout = dropout

    def forward(self, x, edge_index, edge_weight=None):
        x = F.elu(self.gat1(x, edge_index, edge_weight))
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = F.elu(self.gat2(x, edge_index, edge_weight))
        return self.head(x).squeeze(-1)


# ─────────────────────────────────────────────────────────────
# 데이터 로더
# ─────────────────────────────────────────────────────────────
async def load_active_codes(conn) -> List[str]:
    rows = await conn.fetch(
        "SELECT stock_code FROM dart_corps WHERE is_active=true ORDER BY stock_code")
    return [r["stock_code"] for r in rows]


async def load_features(conn, codes: List[str], target_date: date) -> np.ndarray:
    """노드 피처 행렬 (N, 12) — 결측은 0으로."""
    N = len(codes)
    F_ = FEATURE_DIM
    out = np.zeros((N, F_), dtype=np.float32)
    idx = {c: i for i, c in enumerate(codes)}

    # ── 가격 모멘텀 (1d/5d/20d 수익률), vol_ratio ──
    rows = await conn.fetch("""
        SELECT stock_code, dt, close_price, volume
        FROM stock_ohlcv
        WHERE dt BETWEEN $1 AND $2
        ORDER BY stock_code, dt
    """, target_date - timedelta(days=35), target_date)
    df = pd.DataFrame(rows, columns=["stock_code", "dt", "close", "volume"])
    if not df.empty:
        df["close"] = df["close"].astype(float)
        df["volume"] = df["volume"].astype(float)
        for code, g in df.groupby("stock_code"):
            if code not in idx:
                continue
            g = g.sort_values("dt")
            closes = g["close"].values
            vols = g["volume"].values
            if len(closes) >= 2:
                out[idx[code], 0] = (closes[-1] / closes[-2] - 1) * 100
            if len(closes) >= 6:
                out[idx[code], 1] = (closes[-1] / closes[-6] - 1) * 100
            if len(closes) >= 21:
                out[idx[code], 2] = (closes[-1] / closes[-21] - 1) * 100
            if len(vols) >= 20:
                v20 = vols[-20:].mean()
                out[idx[code], 3] = (vols[-1] / v20) if v20 > 0 else 1.0

    # ── 기술적 (RSI, tech_score) ──
    rows = await conn.fetch("""
        SELECT DISTINCT ON (stock_code) stock_code, rsi, tech_score
        FROM stock_technical
        WHERE analyzed_at::date <= $1
        ORDER BY stock_code, analyzed_at DESC
    """, target_date)
    for r in rows:
        if r["stock_code"] in idx:
            out[idx[r["stock_code"]], 4] = float(r["rsi"] or 50)
            out[idx[r["stock_code"]], 5] = float(r["tech_score"] or 0)

    # ── 펀더멘털 (ROE, op_margin, debt_ratio) ──
    rows = await conn.fetch("""
        SELECT DISTINCT ON (stock_code) stock_code, roe, operating_margin, debt_ratio
        FROM dart_financials
        WHERE reprt_code='11011'
        ORDER BY stock_code, bsns_year DESC
    """)
    for r in rows:
        if r["stock_code"] in idx:
            out[idx[r["stock_code"]], 6] = float(r["roe"] or 0)
            out[idx[r["stock_code"]], 7] = float(r["operating_margin"] or 0)
            out[idx[r["stock_code"]], 8] = float(r["debt_ratio"] or 0)

    # ── 뉴스 감성 7일 ──
    rows = await conn.fetch("""
        SELECT primary_stock,
               SUM(CASE sentiment WHEN '호재' THEN intensity
                                  WHEN '악재' THEN -intensity ELSE 0 END)::float AS s
        FROM news_analysis
        WHERE analyzed_at BETWEEN $1 AND $2
          AND primary_stock IS NOT NULL
        GROUP BY primary_stock
    """, target_date - timedelta(days=7), target_date + timedelta(days=1))
    for r in rows:
        if r["primary_stock"] in idx:
            out[idx[r["primary_stock"]], 9] = float(r["s"])

    # ── us_overnight_adj (latest) ──
    rows = await conn.fetch("""
        SELECT DISTINCT ON (stock_code) stock_code, us_overnight_adj
        FROM stock_scores
        WHERE score_date <= $1 AND us_overnight_adj IS NOT NULL
        ORDER BY stock_code, score_date DESC
    """, target_date)
    for r in rows:
        if r["stock_code"] in idx:
            out[idx[r["stock_code"]], 10] = float(r["us_overnight_adj"])

    # ── log market cap ──
    rows = await conn.fetch("""
        SELECT DISTINCT ON (stock_code) stock_code, market_cap
        FROM stock_prices
        WHERE market_cap IS NOT NULL
        ORDER BY stock_code, collected_at DESC
    """)
    for r in rows:
        if r["stock_code"] in idx and r["market_cap"]:
            out[idx[r["stock_code"]], 11] = math.log10(float(r["market_cap"]) + 1)

    # 표준화 (피처별 z-score, 클리핑)
    mu = out.mean(axis=0)
    sd = out.std(axis=0) + 1e-6
    out = np.clip((out - mu) / sd, -5, 5)
    return out.astype(np.float32)


async def load_price_corr_edges(conn, codes: List[str], target_date: date,
                                 lookback: int = CORR_LOOKBACK,
                                 threshold: float = CORR_THRESHOLD
                                 ) -> Tuple[np.ndarray, np.ndarray]:
    """가격 상관 엣지. 반환: edge_index (2, E), edge_weight (E,)."""
    idx = {c: i for i, c in enumerate(codes)}
    rows = await conn.fetch("""
        SELECT stock_code, dt, close_price FROM stock_ohlcv
        WHERE dt BETWEEN $1 AND $2 AND stock_code = ANY($3)
        ORDER BY stock_code, dt
    """, target_date - timedelta(days=int(lookback * 1.6)),
        target_date, codes)
    if not rows:
        return np.zeros((2, 0), dtype=np.int64), np.zeros(0, dtype=np.float32)
    df = pd.DataFrame(rows, columns=["code", "dt", "close"])
    pivot = df.pivot(index="dt", columns="code", values="close").astype(float)
    pivot = pivot.tail(lookback)
    rets = pivot.pct_change().dropna(how="all")
    # 결측치 많은 종목 제거 (절반 이상 결측)
    rets = rets.dropna(axis=1, thresh=int(len(rets) * 0.5))
    if rets.shape[1] < 2:
        return np.zeros((2, 0), dtype=np.int64), np.zeros(0, dtype=np.float32)
    rets = rets.fillna(0)
    M = rets.values  # (T, K)
    M = (M - M.mean(axis=0)) / (M.std(axis=0) + 1e-9)
    K = M.shape[1]
    corr = (M.T @ M) / max(M.shape[0] - 1, 1)
    np.fill_diagonal(corr, 0)
    abs_corr = np.abs(corr)
    # 임계값 필터 + 노드당 top-K 이웃만 유지 (메모리 캡)
    mask = abs_corr > threshold
    src_set: List[int] = []
    dst_set: List[int] = []
    w_set: List[float] = []
    code_arr = list(rets.columns)
    for i in range(K):
        cands = np.where(mask[i])[0]
        if len(cands) == 0:
            continue
        if len(cands) > TOP_K_NEIGHBORS:
            top_idx = np.argpartition(-abs_corr[i, cands], TOP_K_NEIGHBORS)[:TOP_K_NEIGHBORS]
            cands = cands[top_idx]
        for j in cands:
            src_set.append(idx[code_arr[i]])
            dst_set.append(idx[code_arr[j]])
            w_set.append(float(abs_corr[i, j]))
    if not src_set:
        return np.zeros((2, 0), dtype=np.int64), np.zeros(0, dtype=np.float32)
    return (np.stack([np.array(src_set, dtype=np.int64),
                       np.array(dst_set, dtype=np.int64)]),
            np.array(w_set, dtype=np.float32))


async def load_sector_edges(conn, codes: List[str]) -> np.ndarray:
    """같은 sector 노드 간 양방향 엣지."""
    idx = {c: i for i, c in enumerate(codes)}
    rows = await conn.fetch("""
        SELECT stock_code, sector FROM dart_corps
        WHERE is_active=true AND sector IS NOT NULL AND stock_code = ANY($1)
    """, codes)
    by_sec = {}
    for r in rows:
        by_sec.setdefault(r["sector"], []).append(r["stock_code"])
    src, dst = [], []
    for sec, lst in by_sec.items():
        # 섹터당 종목 수가 100개 초과면 비대해지므로 무작위 샘플
        if len(lst) > 80:
            import random
            random.seed(hash(sec) & 0xffff)
            lst = random.sample(lst, 80)
        for i, a in enumerate(lst):
            for b in lst[i+1:]:
                if a in idx and b in idx:
                    src += [idx[a], idx[b]]
                    dst += [idx[b], idx[a]]
    if not src:
        return np.zeros((2, 0), dtype=np.int64)
    return np.stack([np.array(src, dtype=np.int64), np.array(dst, dtype=np.int64)])


async def load_news_edges(conn, codes: List[str], target_date: date,
                           lookback_days: int = NEWS_LOOKBACK,
                           min_count: int = NEWS_COOC_MIN) -> np.ndarray:
    """뉴스 공기관계: 같은 뉴스의 affected_stocks 페어 카운트 ≥ min_count."""
    idx = {c: i for i, c in enumerate(codes)}
    rows = await conn.fetch("""
        SELECT affected_stocks FROM news_analysis
        WHERE analyzed_at BETWEEN $1 AND $2
          AND jsonb_array_length(affected_stocks) > 1
    """, target_date - timedelta(days=lookback_days),
        target_date + timedelta(days=1))
    from collections import Counter
    counter: Counter = Counter()
    for r in rows:
        codes_in_news = r["affected_stocks"]
        if isinstance(codes_in_news, str):
            codes_in_news = json.loads(codes_in_news)
        codes_in_news = [c for c in codes_in_news if c in idx]
        for i, a in enumerate(codes_in_news):
            for b in codes_in_news[i+1:]:
                key = (a, b) if a < b else (b, a)
                counter[key] += 1
    src, dst = [], []
    for (a, b), cnt in counter.items():
        if cnt >= min_count:
            src += [idx[a], idx[b]]
            dst += [idx[b], idx[a]]
    if not src:
        return np.zeros((2, 0), dtype=np.int64)
    return np.stack([np.array(src, dtype=np.int64), np.array(dst, dtype=np.int64)])


async def build_graph(conn, target_date: date):
    """전체 그래프 구성. 반환: (codes, x, edge_index, edge_weight)."""
    codes = await load_active_codes(conn)
    x = await load_features(conn, codes, target_date)
    ei_corr, ew_corr = await load_price_corr_edges(conn, codes, target_date)
    ei_sec = await load_sector_edges(conn, codes)
    ei_news = await load_news_edges(conn, codes, target_date)
    # 결합 (sector/news는 weight=1.0)
    ei_all = np.concatenate([ei_corr, ei_sec, ei_news], axis=1)
    ew_all = np.concatenate([
        ew_corr,
        np.full(ei_sec.shape[1], 0.5, dtype=np.float32),
        np.full(ei_news.shape[1], 0.7, dtype=np.float32),
    ])
    logger.info("graph.built", nodes=len(codes), edges=int(ei_all.shape[1]),
                price=int(ei_corr.shape[1]), sector=int(ei_sec.shape[1]),
                news=int(ei_news.shape[1]))
    return codes, x, ei_all.astype(np.int64), ew_all.astype(np.float32)


async def load_labels(conn, codes: List[str], target_date: date) -> np.ndarray:
    """target_date 종가 → target_date+1 영업일 종가의 변화율 (%)."""
    idx = {c: i for i, c in enumerate(codes)}
    labels = np.zeros(len(codes), dtype=np.float32)
    mask = np.zeros(len(codes), dtype=bool)
    rows = await conn.fetch("""
        SELECT stock_code, dt, close_price FROM stock_ohlcv
        WHERE dt BETWEEN $1 AND $2 AND stock_code = ANY($3)
        ORDER BY stock_code, dt
    """, target_date, target_date + timedelta(days=7), codes)
    by_code = {}
    for r in rows:
        by_code.setdefault(r["stock_code"], []).append(
            (r["dt"], float(r["close_price"])))
    for code, lst in by_code.items():
        lst.sort()
        if code in idx and len(lst) >= 2:
            t0, c0 = lst[0]
            t1, c1 = lst[1]
            if c0 > 0:
                labels[idx[code]] = (c1 / c0 - 1) * 100
                mask[idx[code]] = True
    return labels, mask


# ─────────────────────────────────────────────────────────────
# 학습
# ─────────────────────────────────────────────────────────────
async def train_model(window_days: int = TRAIN_WINDOW_DAYS,
                       stride: int = SAMPLE_STRIDE_DAYS,
                       epochs: int = 30, lr: float = 1e-3):
    """rolling 윈도우 학습. 일별이 아닌 stride 간격으로 샘플링."""
    today = date.today()
    sample_dates = [today - timedelta(days=window_days - i)
                    for i in range(0, window_days, stride)]
    sample_dates = [d for d in sample_dates if d.weekday() < 5]
    if len(sample_dates) < 8:
        return {"err": "too few samples", "samples": len(sample_dates)}

    val_split = max(1, len(sample_dates) // 5)
    train_dates = sample_dates[:-val_split]
    val_dates = sample_dates[-val_split:]
    logger.info("train.split", train=len(train_dates), val=len(val_dates))

    async with pg_pool.acquire() as conn:
        # 모든 샘플 그래프 미리 구성
        samples = []
        for d in sample_dates:
            try:
                codes, x, ei, ew = await build_graph(conn, d)
                y, m = await load_labels(conn, codes, d)
                if m.sum() < 10:
                    continue
                samples.append({
                    "date": d,
                    "x": torch.tensor(x),
                    "ei": torch.tensor(ei),
                    "ew": torch.tensor(ew),
                    "y": torch.tensor(y),
                    "m": torch.tensor(m),
                })
            except Exception as e:
                logger.warning("train.sample_err", date=str(d), err=str(e))

    if len(samples) < 4:
        return {"err": "no valid samples", "got": len(samples)}

    val_count = max(1, len(samples) // 5)
    train_set = samples[:-val_count]
    val_set = samples[-val_count:]

    model = GraphNet().to(device)
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)
    best_val = float("inf")

    for epoch in range(epochs):
        model.train()
        tot = 0.0
        for s in train_set:
            opt.zero_grad()
            pred = model(s["x"].to(device), s["ei"].to(device),
                         s["ew"].to(device))
            mask = s["m"].to(device)
            loss = F.huber_loss(pred[mask], s["y"].to(device)[mask], delta=2.0)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step()
            tot += float(loss)
        tot /= len(train_set)

        model.eval()
        v = 0.0
        with torch.no_grad():
            for s in val_set:
                pred = model(s["x"].to(device), s["ei"].to(device),
                             s["ew"].to(device))
                mask = s["m"].to(device)
                v += float(F.huber_loss(pred[mask], s["y"].to(device)[mask],
                                         delta=2.0))
        v /= len(val_set)
        logger.info("train.epoch", epoch=epoch, train_loss=tot, val_loss=v)
        if v < best_val:
            best_val = v
            torch.save({"state": model.state_dict(),
                        "feature_dim": FEATURE_DIM,
                        "hidden": HIDDEN_DIM,
                        "heads": ATT_HEADS}, MODEL_PATH)

    async with pg_pool.acquire() as conn:
        await conn.execute("""
            INSERT INTO graph_model_meta
              (model_date, train_samples, val_samples, val_loss,
               edge_count, node_count, hidden_dim, notes)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
            ON CONFLICT (model_date) DO UPDATE
              SET train_samples=$2, val_samples=$3, val_loss=$4,
                  edge_count=$5, node_count=$6, hidden_dim=$7, notes=$8
        """, today, len(train_set), len(val_set), best_val,
            int(samples[-1]["ei"].shape[1]), int(samples[-1]["x"].shape[0]),
            HIDDEN_DIM, f"epochs={epochs} lr={lr} stride={stride}")
    return {"trained": True, "train_samples": len(train_set),
            "val_samples": len(val_set), "best_val_loss": best_val}


# ─────────────────────────────────────────────────────────────
# 추론
# ─────────────────────────────────────────────────────────────
async def predict_today() -> dict:
    if not os.path.exists(MODEL_PATH):
        return {"err": "no trained model"}
    today = date.today()
    async with pg_pool.acquire() as conn:
        codes, x, ei, ew = await build_graph(conn, today)
    model = GraphNet().to(device)
    ckpt = torch.load(MODEL_PATH, map_location=device, weights_only=True)
    model.load_state_dict(ckpt["state"])
    model.eval()
    with torch.no_grad():
        pred = model(torch.tensor(x).to(device),
                     torch.tensor(ei).to(device),
                     torch.tensor(ew).to(device))
    pred_np = pred.cpu().numpy()
    # graph_predictions만 저장 — score-engine이 calculate_daily_scores 시 읽어 ensemble에 적용.
    async with pg_pool.acquire() as conn:
        async with conn.transaction():
            for i, code in enumerate(codes):
                p = float(pred_np[i])
                await conn.execute("""
                    INSERT INTO graph_predictions (stock_code, predict_date, pred_return)
                    VALUES ($1, $2, $3)
                    ON CONFLICT (stock_code, predict_date) DO UPDATE
                      SET pred_return=$3, created_at=NOW()
                """, code, today, p)
    return {"predicted": len(codes), "date": str(today)}


# ─────────────────────────────────────────────────────────────
# FastAPI
# ─────────────────────────────────────────────────────────────
@app.on_event("startup")
async def on_start():
    global pg_pool
    pg_pool = await asyncpg.create_pool(**PG, min_size=2, max_size=5)
    async with pg_pool.acquire() as conn:
        await conn.execute(DDL)
    # 16:25 KST: score-engine 16:30 calculate_daily_scores 직전 추론
    scheduler.add_job(predict_today, CronTrigger(hour=16, minute=25,
                                                  day_of_week="mon-fri"),
                      id="graph_predict", replace_existing=True)
    scheduler.add_job(train_model, CronTrigger(day_of_week="sun", hour=6,
                                                minute=0),
                      id="graph_train", replace_existing=True)
    scheduler.start()
    logger.info("graph-engine.started")


@app.on_event("shutdown")
async def on_stop():
    if scheduler.running:
        scheduler.shutdown()
    if pg_pool:
        await pg_pool.close()


@app.get("/health")
async def health():
    return {"status": "ok",
            "model_exists": os.path.exists(MODEL_PATH)}


@app.post("/graph/build")
async def manual_build(target: Optional[str] = None):
    d = date.fromisoformat(target) if target else date.today()
    async with pg_pool.acquire() as conn:
        codes, x, ei, ew = await build_graph(conn, d)
    return {"nodes": len(codes), "edges": int(ei.shape[1]),
            "feature_shape": list(x.shape), "date": str(d)}


@app.post("/train")
async def manual_train(epochs: int = Query(default=30, ge=1, le=200),
                        window_days: int = Query(default=180, ge=30, le=365),
                        stride: int = Query(default=5, ge=1, le=30),
                        bg: BackgroundTasks = None):
    if bg:
        bg.add_task(train_model, window_days, stride, epochs)
        return {"status": "queued", "epochs": epochs}
    return await train_model(window_days, stride, epochs)


@app.post("/predict")
async def manual_predict():
    return await predict_today()


@app.get("/predict/{code}")
async def predict_stock(code: str, limit: int = Query(default=30, ge=1, le=180)):
    async with pg_pool.acquire() as conn:
        rows = await conn.fetch("""
            SELECT predict_date, pred_return
            FROM graph_predictions
            WHERE stock_code=$1
            ORDER BY predict_date DESC LIMIT $2
        """, code, limit)
    return {"code": code, "history": [dict(r) for r in rows]}


@app.get("/status")
async def status():
    async with pg_pool.acquire() as conn:
        meta = await conn.fetchrow("""
            SELECT * FROM graph_model_meta ORDER BY model_date DESC LIMIT 1
        """)
        latest_pred = await conn.fetchrow("""
            SELECT predict_date, COUNT(*) AS n
            FROM graph_predictions
            GROUP BY predict_date
            ORDER BY predict_date DESC LIMIT 1
        """)
    return {
        "model": dict(meta) if meta else None,
        "latest_prediction": dict(latest_pred) if latest_pred else None,
        "model_exists": os.path.exists(MODEL_PATH),
    }