"""feature_dqn 체크포인트(.pt) → 서빙 번들(dqn_serving.npz) export.

서빙 컨테이너에 PyTorch 를 넣지 않기 위해 ScoreNet(3층 MLP) 가중치와 카드 특징
(임베딩384 + 전략 one-hot4 + 톤 one-hot4 = 392)을 numpy 번들 하나로 묶는다.
추론은 negotiation.policy.dqn_store 의 numpy forward 가 수행한다.

실행(호스트, torch 필요):  APP_ENV=local python -m tools.export_dqn_serving
산출: agent/artifacts/dqn_serving.npz  (.dockerignore 미제외 → 이미지에 포함)
"""

import os

import numpy as np
import torch

from tools.train_feature_dqn import load_cards

_HERE = os.path.dirname(os.path.abspath(__file__))
CKPT_PATH = os.path.join(_HERE, "..", "artifacts", "feature_dqn_ktcommerce.pt")
OUT_PATH = os.path.join(_HERE, "..", "artifacts", "dqn_serving.npz")

STATE_DIM = 14   # build_state_features(9) + build_tenant_features(5)
CARD_DIM = 392


def _np_forward(x, W0, b0, W1, b1, W2, b2):
    h = np.maximum(x @ W0.T + b0, 0.0)
    h = np.maximum(h @ W1.T + b1, 0.0)
    return h @ W2.T + b2


def export_bundle(sd, out_path: str) -> str:
    """state_dict → 서빙 번들 npz (원자적 교체: .tmp 작성 후 replace). 반환: 절대경로.

    retrain_from_logs 재학습 배포도 이 함수를 쓴다 — 검증(torch/numpy 일치)은 main() 전용.
    """
    W0, b0 = sd["net.0.weight"].numpy(), sd["net.0.bias"].numpy()
    W1, b1 = sd["net.2.weight"].numpy(), sd["net.2.bias"].numpy()
    W2, b2 = sd["net.4.weight"].numpy(), sd["net.4.bias"].numpy()
    assert W0.shape[1] == STATE_DIM + CARD_DIM, f"입력 차원 불일치: {W0.shape[1]}"
    numbers, feat, _ = load_cards()
    card_feats = np.stack([feat[n] for n in numbers]).astype(np.float32)
    tmp = out_path + ".tmp"
    with open(tmp, "wb") as f:
        np.savez(
            f,
            W0=W0, b0=b0, W1=W1, b1=b1, W2=W2, b2=b2,
            card_numbers=np.array(numbers), card_feats=card_feats,
            state_dim=STATE_DIM, card_dim=CARD_DIM,
        )
    if os.path.exists(out_path):
        os.replace(out_path, out_path + ".prev")   # 직전 번들 백업(롤백용)
    os.replace(tmp, out_path)
    return os.path.abspath(out_path)


def main():
    sd = torch.load(CKPT_PATH, map_location="cpu")
    # 정합성 검증: torch forward == numpy forward
    from negotiation.policies.feature_dqn_policy import ScoreNet
    net = ScoreNet(STATE_DIM, CARD_DIM)
    net.load_state_dict(sd)
    net.eval()
    x = np.random.default_rng(0).normal(size=(8, STATE_DIM + CARD_DIM)).astype(np.float32)
    with torch.no_grad():
        ref = net(torch.tensor(x)).numpy()
    W0, b0 = sd["net.0.weight"].numpy(), sd["net.0.bias"].numpy()
    W1, b1 = sd["net.2.weight"].numpy(), sd["net.2.bias"].numpy()
    W2, b2 = sd["net.4.weight"].numpy(), sd["net.4.bias"].numpy()
    out = _np_forward(x, W0, b0, W1, b1, W2, b2).squeeze(-1)
    diff = float(np.abs(ref - out).max())
    assert diff < 1e-4, f"numpy/torch forward 불일치: {diff}"

    path = export_bundle(sd, OUT_PATH)
    print(f"[저장] {path}  forward 오차 {diff:.2e}")


if __name__ == "__main__":
    main()