"""Q-Table — 이산 (state, action) 가치표 + 방문횟수 (우리 자체 numpy 구현).

클린룸: Q-learning(off-policy TD) 과 UCB 는 표준 알고리즘(아이디어)이고, 아래 구현은 우리 작성이다.
state_space_size x action_space_size 밀집 행렬. 차원은 TenantConfig 에서 산출된 값을 주입.
"""

from typing import List, Tuple

import numpy as np


class QTable:
    def __init__(self, state_space_size: int, action_space_size: int,
                 learning_rate: float = 0.1, discount_factor: float = 0.95):
        self.state_space_size = state_space_size
        self.action_space_size = action_space_size
        self.lr = learning_rate
        self.gamma = discount_factor
        self.q = np.zeros((state_space_size, action_space_size), dtype=float)
        self.visits = np.zeros((state_space_size, action_space_size), dtype=np.int64)

    # ---- 접근 ----------------------------------------------------------
    def row(self, state_index: int) -> np.ndarray:
        return self.q[state_index]

    def visit_row(self, state_index: int) -> np.ndarray:
        return self.visits[state_index]

    def state_visits(self, state_index: int) -> int:
        return int(self.visits[state_index].sum())

    def best_action(self, state_index: int) -> int:
        return int(np.argmax(self.q[state_index]))

    # ---- 갱신 ----------------------------------------------------------
    def mark_visit(self, state_index: int, action_id: int):
        self.visits[state_index, action_id] += 1

    def update(self, state_index: int, action_id: int, reward: float,
               next_state_index: int = None, done: bool = False) -> float:
        """Q-learning 1-스텝 갱신. 반환: 갱신 후 Q[s,a].

        target = reward + (0 if done/next 없음 else gamma * max_a' Q[s',a'])
        Q[s,a] += lr * (target - Q[s,a])
        """
        bootstrap = 0.0
        if not done and next_state_index is not None:
            bootstrap = self.gamma * float(np.max(self.q[next_state_index]))
        td_target = reward + bootstrap
        self.q[state_index, action_id] += self.lr * (td_target - self.q[state_index, action_id])
        return float(self.q[state_index, action_id])

    # ---- 직렬화 (희소: 0 아닌 셀만) ------------------------------------
    def nonzero_cells(self) -> List[Tuple[int, int, float, int]]:
        """(state_index, action_id, q_value, visit_count) — q 또는 visit 가 0 이 아닌 셀."""
        out = []
        nz = np.argwhere((self.q != 0) | (self.visits != 0))
        for s, a in nz:
            out.append((int(s), int(a), float(self.q[s, a]), int(self.visits[s, a])))
        return out

    def load_cells(self, cells: List[Tuple[int, int, float, int]]):
        for s, a, q, v in cells:
            if 0 <= s < self.state_space_size and 0 <= a < self.action_space_size:
                self.q[s, a] = q
                self.visits[s, a] = v