"""
Q-Table 모델
동적 action_space_size를 지원하여 협상 카드 추가/삭제에 대응
"""

import numpy as np
from typing import Optional, Dict


class QTable:
    """
    Q-Learning을 위한 Q-Table

    Q-Table은 추상적인 action_id만 다루며,
    실제 협상 카드(card_id)는 ActionCardMapper에서 매핑
    """

    def __init__(
        self,
        state_space_size: int,
        action_space_size: int,
        learning_rate: float = 0.1,
        discount_factor: float = 0.95,
    ):
        """
        Args:
            state_space_size: 상태 공간 크기
                             (Scenario=4 x PriceZone=3 x AcceptanceRate=3 = 36)
            action_space_size: 액션 공간 크기 (협상 카드 개수, 현재 21개)
            learning_rate: 학습률 (alpha)
            discount_factor: 할인 인자 (gamma)
        """
        self.state_space_size = state_space_size
        self.action_space_size = action_space_size
        self.learning_rate = learning_rate
        self.discount_factor = discount_factor

        # Q-Table 초기화: (state_size, action_size) 형태
        self.q_values = np.zeros((state_space_size, action_space_size))

    def get_q_value(self, state_index: int, action_id: int) -> float:
        """
        특정 (state, action)의 Q-value 조회

        Args:
            state_index: State의 인덱스 (0 ~ state_space_size-1)
            action_id: Action ID (0 ~ action_space_size-1)

        Returns:
            Q-value
        """
        self._validate_indices(state_index, action_id)
        return float(self.q_values[state_index, action_id])

    def get_q_values_for_state(self, state_index: int) -> np.ndarray:
        """
        특정 state의 모든 action에 대한 Q-values

        Args:
            state_index: State의 인덱스

        Returns:
            shape (action_space_size,)의 Q-values 배열
        """
        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.q_values[state_index, :].copy()

    def set_q_value(self, state_index: int, action_id: int, value: float):
        """
        특정 (state, action)의 Q-value 직접 설정

        Args:
            state_index: State의 인덱스
            action_id: Action ID
            value: 설정할 Q-value
        """
        self._validate_indices(state_index, action_id)
        self.q_values[state_index, action_id] = value

    def update(
        self,
        state_index: int,
        action_id: int,
        reward: float,
        next_state_index: Optional[int] = None,
        done: bool = False,
    ):
        """
        Q-Learning 업데이트

        Q(s,a) ← Q(s,a) + α[r + γ·max_a'Q(s',a') - Q(s,a)]

        Args:
            state_index: 현재 상태 인덱스
            action_id: 선택한 액션 ID
            reward: 받은 보상
            next_state_index: 다음 상태 인덱스 (종료 시 None)
            done: 에피소드 종료 여부
        """
        self._validate_indices(state_index, action_id)

        current_q = self.q_values[state_index, action_id]

        if done or next_state_index is None:
            # 종료 상태: target = reward만 사용
            target = reward
        else:
            # 비종료 상태: target = reward + γ·max Q(s',a')
            if next_state_index < 0 or next_state_index >= self.state_space_size:
                raise ValueError(f"next_state_index {next_state_index} out of range")
            max_next_q = np.max(self.q_values[next_state_index, :])
            target = reward + self.discount_factor * max_next_q

        # Q-Learning 업데이트
        self.q_values[state_index, action_id] += self.learning_rate * (
            target - current_q
        )

    def get_best_action(self, state_index: int) -> int:
        """
        해당 state에서 최고 Q-value를 가진 action_id 반환

        Args:
            state_index: State의 인덱스

        Returns:
            action_id (Q-value가 최대인 액션)
        """
        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(np.argmax(self.q_values[state_index, :]))

    def get_best_actions_with_ties(self, state_index: int) -> list:
        """
        동일한 최대 Q-value를 가진 모든 action_id 반환

        Args:
            state_index: State의 인덱스

        Returns:
            최대 Q-value를 가진 action_id들의 리스트
        """
        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})"
            )

        max_q = np.max(self.q_values[state_index, :])
        max_actions = np.where(self.q_values[state_index, :] == max_q)[0]
        return max_actions.tolist()

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

        Returns:
            직렬화된 Q-Table 데이터
        """
        return {
            "state_space_size": self.state_space_size,
            "action_space_size": self.action_space_size,
            "learning_rate": self.learning_rate,
            "discount_factor": self.discount_factor,
            "q_values": self.q_values.tolist(),
        }

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

        Args:
            data: 직렬화된 Q-Table 데이터

        Returns:
            복원된 QTable 인스턴스
        """
        q_table = cls(
            state_space_size=data["state_space_size"],
            action_space_size=data["action_space_size"],
            learning_rate=data["learning_rate"],
            discount_factor=data["discount_factor"],
        )
        q_table.q_values = np.array(data["q_values"])
        return q_table

    def _validate_indices(self, state_index: int, action_id: int):
        """인덱스 유효성 검증"""
        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})"
            )

    def __repr__(self) -> str:
        return (
            f"QTable(state_space_size={self.state_space_size}, "
            f"action_space_size={self.action_space_size}, "
            f"learning_rate={self.learning_rate}, "
            f"discount_factor={self.discount_factor})"
        )
