"""Natrenovana siet ako hrac so standardnym rozhranim guess/play.

Rovnake rozhranie ako rl/players.py, takze funguje v rl/evaluate.py aj ako
boti "mozog" v api/bots.py. Hrac vidi len observaciu + masku z rl/encoding.py
-- z principu nemoze podvadzat (do cudzich ruk sa nedostane).
"""

import torch

from rl.encoding import encode_observation, guess_mask, play_mask
from rl.model import BridzikNet, mask_tensor, masked_categorical, obs_tensor


class NeuralPlayer:
    def __init__(self, net: BridzikNet, greedy: bool = True):
        self.net = net
        self.greedy = greedy  # argmax pri evaluacii; sampling pre pestrost

    def _act(self, rnd, seat: int, use_play_head: bool) -> int:
        obs = obs_tensor(encode_observation(rnd, seat)).unsqueeze(0)
        mask = mask_tensor(
            play_mask(rnd, seat) if use_play_head else guess_mask(rnd)
        ).unsqueeze(0)
        self.net.eval()
        with torch.no_grad():
            guess_logits, play_logits, _ = self.net(obs)
            logits = play_logits if use_play_head else guess_logits
            logits = logits.masked_fill(~mask, float('-inf'))
            if self.greedy:
                return int(logits.argmax(dim=-1).item())
            return int(masked_categorical(logits, mask).sample().item())

    def guess(self, rnd, seat: int) -> int:
        return self._act(rnd, seat, use_play_head=False)

    def play(self, rnd, seat: int) -> int:
        return self._act(rnd, seat, use_play_head=True)


def load_player(checkpoint_path: str, greedy: bool = True) -> NeuralPlayer:
    from rl.train import load_checkpoint
    return NeuralPlayer(load_checkpoint(checkpoint_path), greedy=greedy)