from abc import ABC, abstractmethod
from typing import Optional

import numpy as np

from ..model.visit_table import VisitTable


class Policy(ABC):
    @abstractmethod
    def select_action(
        self,
        state_index: int,
        q_values: np.ndarray,
        available_mask: Optional[np.ndarray] = None,
    ) -> Optional[int]:
        """
        주어진 상태와 Q-value를 기반으로 액션 선택

        Args:
            state_index: 현재 상태 인덱스
            q_values: 해당 상태의 Q-value 배열
            available_mask: 선택 가능한 액션 마스크 (True=선택가능)

        Returns:
            선택된 액션 ID (선택 불가 시 None)
        """
        raise NotImplementedError

    def reset_episode(self) -> None:
        """에피소드 시작 시 상태 초기화"""
        raise NotImplementedError

    def get_action_mask(self) -> np.ndarray:
        """
        현재 정책 상태에 따른 액션 마스크 반환

        Returns:
            액션 마스크 배열 (1=선택가능, 0=선택불가)
        """
        raise NotImplementedError


class UCBPolicy(Policy):
    """Upper Confidence Bound policy with per-episode action masking."""

    def __init__(
        self,
        visit_table: VisitTable,
        exploration_constant: float = np.sqrt(2.0),
        rng: Optional[np.random.Generator] = None,
    ) -> None:
        self.visit_table = visit_table
        self.action_space_size = visit_table.action_space_size
        self.exploration_constant = exploration_constant
        self._rng = rng or np.random.default_rng()
        self._episode_actions: set[int] = set()

    def select_action(
        self,
        state_index: int,
        q_values: np.ndarray,
        available_mask: Optional[np.ndarray] = None,
    ) -> Optional[int]:
        """
        UCB 알고리즘을 사용하여 액션 선택

        UCB = Q(s,a) + c * sqrt(ln(N(s)) / N(s,a))
        """
        mask = self._prepare_mask(available_mask)
        if not mask.any():
            self.reset_episode()
            mask = self._prepare_mask(available_mask)
            if not mask.any():
                return None

        counts = self.visit_table.get_state_counts(state_index)
        masked_counts = np.where(mask, counts, 0)

        zero_visit_candidates = np.where((counts == 0) & mask)[0]
        if zero_visit_candidates.size > 0:
            best = zero_visit_candidates[np.argmax(q_values[zero_visit_candidates])]
            action = int(best)
        else:
            total = masked_counts.sum()
            # Avoid division by zero while keeping exploration pressure.
            denom = counts.astype(float) + 1e-9
            bonus = self.exploration_constant * np.sqrt(np.log(total + 1.0) / denom)
            scores = q_values + bonus
            scores[~mask] = -np.inf
            action = int(np.argmax(scores))

        self.visit_table.increment(state_index, action)
        self._episode_actions.add(action)
        return action

    def reset_episode(self) -> None:
        """에피소드 내 사용된 액션 기록 초기화"""
        self._episode_actions.clear()

    def get_action_mask(self) -> np.ndarray:
        """
        이미 사용된 액션을 제외한 마스크 반환

        Returns:
            마스크 배열 (1=선택가능, 0=선택불가)
        """
        mask = np.ones(self.action_space_size, dtype=int)
        if self._episode_actions:
            for action in self._episode_actions:
                mask[action] = 0
        return mask

    def _prepare_mask(self, available_mask: Optional[np.ndarray]) -> np.ndarray:
        """입력 마스크와 이미 사용된 액션을 결합하여 최종 마스크 생성"""
        if available_mask is None:
            mask = np.ones(self.action_space_size, dtype=bool)
        else:
            mask = np.asarray(available_mask, dtype=bool).copy()
        if self._episode_actions:
            for action in self._episode_actions:
                mask[action] = False
        return mask
