"""UCBQTablePolicy — UCB 탐색 기반 Q-Table 정책 (NegotiationPolicy 구현, 우리 자체 구현).

select: 가용 액션 중 UCB 점수 최대 선택 (중복방지 마스킹 + propensity 산출).
update: Q-learning 1-스텝.

propensity(계획서 G): UCB 는 결정론적이라 그대로면 IPS 지지(support)가 0 이 된다.
과거 로그를 ε-greedy 근사로 본다 — 선택(greedy) 액션에 (1-ε)+ε/n, 나머지 ε/n.
이렇게 로깅된 propensity 가 OPE(IPS/DR/SNIPS)의 입력이 된다.
"""

import math
from typing import List

import numpy as np

from negotiation.policies.base import ActionDecision, NegotiationPolicy, PolicyContext, Transition
from negotiation.qtable.domain.model.q_table import QTable


class UCBQTablePolicy(NegotiationPolicy):
    name = "qtable_ucb"

    def __init__(self, qtable: QTable, exploration_constant: float = math.sqrt(2.0),
                 epsilon: float = 0.1, mark_visits: bool = True):
        self.qtable = qtable
        self.c = exploration_constant
        self.epsilon = epsilon       # propensity 근사용 ε (로깅 전용, 선택 자체는 결정론적 UCB)
        self.mark_visits = mark_visits

    # ---- 선택 ----------------------------------------------------------
    def _available(self, ctx: PolicyContext) -> List[int]:
        if ctx.available_mask is not None:
            avail = [a for a in range(ctx.action_space_size) if ctx.available_mask[a]]
        else:
            used = ctx.episode.used_action_ids if ctx.episode else set()
            avail = [a for a in range(ctx.action_space_size) if a not in used]
        return avail or list(range(ctx.action_space_size))  # 다 썼으면 전체 허용

    def _ucb_scores(self, state_index: int, available: List[int],
                    prior: "np.ndarray | None" = None) -> np.ndarray:
        q = self.qtable.row(state_index)
        visits = self.qtable.visit_row(state_index)
        total = self.qtable.state_visits(state_index)
        ln = math.log(total + 1.0)
        scores = np.full(self.qtable.action_space_size, -np.inf)
        for a in available:
            bonus = self.c * math.sqrt(ln / (visits[a] + 1e-6))
            scores[a] = q[a] + bonus
            if prior is not None:
                # 의도층 prior — 방문수 감쇠: 콜드 스타트 동점(전부 Q=0·bonus=0)일 때만 순서를
                # 결정하고, 학습이 쌓이면 1/(1+visits) 로 사라진다.
                scores[a] += float(prior[a]) / (1.0 + visits[a])
        return scores

    def select(self, ctx: PolicyContext) -> ActionDecision:
        available = self._available(ctx)
        scores = self._ucb_scores(ctx.state_index, available, prior=ctx.prior_bonus)
        action_id = int(np.argmax(scores))
        n = len(available)
        # ε-greedy 근사 propensity (greedy 액션)
        propensity = (1.0 - self.epsilon) + self.epsilon / n
        if self.mark_visits:
            self.qtable.mark_visit(ctx.state_index, action_id)
            if ctx.episode:
                ctx.episode.mark_used(action_id)
        return ActionDecision(
            action_id=action_id,
            propensity=propensity,
            q_value=float(self.qtable.row(ctx.state_index)[action_id]),
            ucb_score=float(scores[action_id]),
            available_actions=available,
        )

    # ---- 학습 ----------------------------------------------------------
    def update(self, transition: Transition) -> None:
        self.qtable.update(
            transition.state_index, transition.action_id, transition.reward,
            next_state_index=transition.next_state_index, done=transition.done,
        )

    def predict_action_dist(self, ctx: PolicyContext) -> np.ndarray:
        """ε-greedy 근사 분포 (OPE/시뮬레이터용)."""
        available = self._available(ctx)
        scores = self._ucb_scores(ctx.state_index, available)
        greedy = int(np.argmax(scores))
        n = len(available)
        dist = np.zeros(ctx.action_space_size)
        for a in available:
            dist[a] = self.epsilon / n
        dist[greedy] += (1.0 - self.epsilon)
        return dist

    # ---- warm-start / 직렬화 ------------------------------------------
    def warm_start(self, other: "UCBQTablePolicy") -> None:
        if (other.qtable.state_space_size != self.qtable.state_space_size
                or other.qtable.action_space_size != self.qtable.action_space_size):
            raise ValueError("dimension mismatch — warm-start 불가 (휴리스틱 init 폴백 필요)")
        self.qtable.q = other.qtable.q.copy()
        # 탐색 여지를 위해 visit 은 감쇠 복제 (계획서 D)
        self.qtable.visits = (other.qtable.visits * 0.5).astype(np.int64)

    def snapshot(self) -> dict:
        return {"cells": self.qtable.nonzero_cells(),
                "state_space_size": self.qtable.state_space_size,
                "action_space_size": self.qtable.action_space_size}

    def load_snapshot(self, data: dict) -> None:
        self.qtable.load_cells(data.get("cells", []))