"""
Competition / submission agent for 51.512 Paddle Duel RL Arena.

Identifier: 1011241_Siddhant

The tabular policy tracks incoming balls. The default pixel policy uses
predictive geometry. A residual Double DQN or a DAgger + Double DQN checkpoint
can replace that image policy after a held-out gate against the predictive
baseline.

Usage (evaluation / competition)::

    from student_agent import Agent
    agent = Agent(side="left")
    agent.reset()
    action = agent.select_action(obs, info=info, training=False)
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, Optional

import numpy as np

from paddle_duel.agents.controllers import PredictivePixelAgent, TabularTrackingAgent
from paddle_duel.paths import CHECKPOINTS, METRICS
from paddle_duel_env import BaseAgent


def _gate_promoted(filename: str) -> bool:
    path = METRICS / filename
    if not path.exists():
        return False
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return False
    return bool(payload.get("promote"))


def _preferred_imitation_checkpoint() -> Path:
    for name in ("level5_imitation_dqn.pt", "level4_imitation_dqn.pt"):
        path = CHECKPOINTS / name
        if path.exists():
            return path
    return CHECKPOINTS / "level4_imitation_dqn.pt"


class Agent(BaseAgent):
    """Paddle Duel agent."""

    def __init__(
        self,
        side: str = "left",
        name: str = "1011241_Siddhant",
        seed: Optional[int] = None,
        tabular_path: Optional[str | Path] = None,
        pixel_path: Optional[str | Path] = None,
        residual_path: Optional[str | Path] = None,
        imitation_path: Optional[str | Path] = None,
        force_residual: bool = False,
        force_imitation: bool = False,
    ):
        assert side in ("left", "right")
        self.side = side
        self.name = name
        self.rng = np.random.default_rng(seed)
        self._tabular = TabularTrackingAgent(name=f"{name}_tabular")
        self._pixel = PredictivePixelAgent(name=f"{name}_pixel")
        self._learned_tabular = None
        self._learned_pixel = None
        self._residual = None
        self._imitation = None
        self._load_checkpoints(
            tabular_path,
            pixel_path,
            residual_path,
            imitation_path,
            force_residual,
            force_imitation,
        )

    def _load_checkpoints(
        self,
        tabular_path,
        pixel_path,
        residual_path,
        imitation_path,
        force_residual: bool,
        force_imitation: bool,
    ) -> None:
        if tabular_path is not None:
            from paddle_duel.agents.tabular import TabularTDAgent

            self._learned_tabular = TabularTDAgent.load(
                Path(tabular_path),
                name=f"{self.name}_learned_tabular",
            )
            self._learned_tabular.epsilon = 0.0

        if pixel_path is not None:
            from paddle_duel.agents.dqn import FeatureDQNAgent

            self._learned_pixel = FeatureDQNAgent.load(
                Path(pixel_path),
                name=f"{self.name}_learned_pixel",
            )
            self._learned_pixel.epsilon = 0.0

        residual_ckpt = Path(residual_path) if residual_path is not None else (
            CHECKPOINTS / "level4_residual_dqn.pt"
        )
        if (force_residual or _gate_promoted("residual_promoted.json")) and residual_ckpt.exists():
            from paddle_duel.agents.residual_dqn import PredictiveResidualDQNAgent

            self._residual = PredictiveResidualDQNAgent.load(
                residual_ckpt,
                name=f"{self.name}_residual",
            )
            self._residual.epsilon = 0.0

        imitation_ckpt = (
            Path(imitation_path) if imitation_path is not None else _preferred_imitation_checkpoint()
        )
        if (force_imitation or _gate_promoted("imitation_promoted.json")) and imitation_ckpt.exists():
            from paddle_duel.agents.imitation import ImitationDQNAgent

            self._imitation = ImitationDQNAgent.load(
                imitation_ckpt,
                name=f"{self.name}_imitation",
            )
            self._imitation.epsilon = 0.0

    def reset(self) -> None:
        self._tabular.reset()
        self._pixel.reset()
        if self._learned_tabular is not None:
            self._learned_tabular.reset()
        if self._learned_pixel is not None:
            self._learned_pixel.reset()
        if self._residual is not None:
            self._residual.reset()
        if self._imitation is not None:
            self._imitation.reset()

    def select_action(
        self,
        observation,
        info: Optional[Dict[str, Any]] = None,
        training: bool = False,
    ) -> int:
        if isinstance(observation, tuple):
            if self._learned_tabular is not None:
                return int(
                    self._learned_tabular.select_action(
                        observation,
                        info=info,
                        training=False,
                    )
                )
            return int(
                self._tabular.select_action(
                    observation,
                    info=info,
                    training=False,
                )
            )

        if self._residual is not None:
            return int(
                self._residual.select_action(
                    observation,
                    info=info,
                    training=False,
                )
            )
        if self._imitation is not None:
            return int(
                self._imitation.select_action(
                    observation,
                    info=info,
                    training=False,
                )
            )
        if self._learned_pixel is not None:
            return int(
                self._learned_pixel.select_action(
                    observation,
                    info=info,
                    training=False,
                )
            )
        return int(
            self._pixel.select_action(
                observation,
                info=info,
                training=False,
            )
        )


StudentAgent = Agent
MyAgent = Agent