"""Visit (N) table for UCB-based policies."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict

import numpy as np


@dataclass
class VisitTable:
    """Tracks state-action visit counts for UCB selection."""

    state_space_size: int
    action_space_size: int

    def __post_init__(self) -> None:
        """
        Args:
            state_space_size: 상태 공간 크기
            action_space_size: 액션 공간 크기
        """
        self._counts = np.zeros(
            (self.state_space_size, self.action_space_size), dtype=np.int64
        )

    def increment(self, state_index: int, action_id: int, count: int = 1) -> None:
        """
        특정 (state, action)의 방문 횟수 증가

        Args:
            state_index: State 인덱스
            action_id: Action ID
            count: 증가시킬 횟수 (기본 1)
        """
        if state_index < 0 or state_index >= self.state_space_size:
            raise ValueError(
                f"state_index {state_index} out of range [0, {self.state_space_size})"
            )
        if action_id < 0 or action_id >= self.action_space_size:
            raise ValueError(
                f"action_id {action_id} out of range [0, {self.action_space_size})"
            )
        if count < 0:
            raise ValueError("count must be non-negative")
        self._counts[state_index, action_id] += count

    def get_action_count(self, state_index: int, action_id: int) -> int:
        """
        특정 (state, action)의 방문 횟수 조회

        Args:
            state_index: State 인덱스
            action_id: Action ID

        Returns:
            방문 횟수
        """
        self._validate_indices(state_index, action_id)
        return int(self._counts[state_index, action_id])

    def get_state_counts(self, state_index: int) -> np.ndarray:
        """
        특정 state의 모든 action에 대한 방문 횟수 조회

        Args:
            state_index: State 인덱스

        Returns:
            shape (action_space_size,)의 방문 횟수 배열
        """
        if state_index < 0 or state_index >= self.state_space_size:
            raise ValueError(
                f"state_index {state_index} out of range [0, {self.state_space_size})"
            )
        return self._counts[state_index, :].copy()

    def total_visits(self, state_index: int) -> int:
        """
        특정 state의 총 방문 횟수 (모든 action 합계)

        Args:
            state_index: State 인덱스

        Returns:
            총 방문 횟수
        """
        if state_index < 0 or state_index >= self.state_space_size:
            raise ValueError(
                f"state_index {state_index} out of range [0, {self.state_space_size})"
            )
        return int(self._counts[state_index, :].sum())

    def reset_state(self, state_index: int) -> None:
        """
        특정 state의 방문 기록 초기화

        Args:
            state_index: State 인덱스
        """
        if state_index < 0 or state_index >= self.state_space_size:
            raise ValueError(
                f"state_index {state_index} out of range [0, {self.state_space_size})"
            )
        self._counts[state_index, :] = 0

    def to_dict(self) -> Dict:
        """
        VisitTable을 딕셔너리로 직렬화

        Returns:
            직렬화된 VisitTable 데이터
        """
        return {
            "state_space_size": self.state_space_size,
            "action_space_size": self.action_space_size,
            "counts": self._counts.tolist(),
        }

    @classmethod
    def from_dict(cls, data: Dict) -> "VisitTable":
        """
        딕셔너리에서 VisitTable 복원

        Args:
            data: 직렬화된 VisitTable 데이터

        Returns:
            복원된 VisitTable 인스턴스
        """
        table = cls(
            state_space_size=data["state_space_size"],
            action_space_size=data["action_space_size"],
        )
        table._counts = np.array(data["counts"], dtype=np.int64)
        return table

    def _validate_indices(self, state_index: int, action_id: int) -> None:
        """인덱스 유효성 검증"""
        if state_index < 0 or state_index >= self.state_space_size:
            raise ValueError(
                f"state_index {state_index} out of range [0, {self.state_space_size})"
            )
        if action_id < 0 or action_id >= self.action_space_size:
            raise ValueError(
                f"action_id {action_id} out of range [0, {self.action_space_size})"
            )
