from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

from ai_hydra.server.SnakeMgr import SnakeMgr
from ai_hydra.constants.DGame import DGameField

from ai_hydra.nnet.Policy.HydraPolicy import HydraPolicy

# from ai_hydra.nnet.UNUSED_RNNTrainer import RNNTrainer
from ai_hydra.nnet.LinearTrainer import LinearTrainer
from ai_hydra.nnet.ReplayMemory import ReplayMemory
from ai_hydra.nnet.Transition import Transition


@dataclass(frozen=True)
class EpisodeStats:
    episode_id: int
    steps: int
    score: int
    reward_total: int
    losses: int
    loss_avg: float


@dataclass(frozen=True)
class TrainRunStats:
    episodes: int
    avg_score: float
    avg_steps: float
    avg_reward_total: float
    avg_loss: float


class TrainMgr:
    """
    Training orchestrator.

    Owns:
      - policy (possibly epsilon-greedy wrapper)
      - trainer (torch + optimizer)
      - replay memory

    Uses:
      - SnakeMgr as the environment/session API

    Does NOT:
      - talk to MQ
      - render
    """

    def __init__(
        self,
        *,
        snake_mgr: SnakeMgr,
        policy: HydraPolicy,
        trainer,
        replay: ReplayMemory,
        client_id: str = "TrainMgr",
    ) -> None:
        self.snake_mgr = snake_mgr
        self.policy = policy
        self.trainer = trainer
        self.replay = replay
        self.client_id = client_id
