"""
---------------------------------------------------------------------------
Air Force Research Laboratory (AFRL) Autonomous Capabilities Team (ACT3)
Reinforcement Learning (RL) Core.

This is a US Government Work not subject to copyright protection in the US.

The use, dissemination or disclosure of data in this file is subject to
limitation or restriction. See accompanying README and LICENSE for details.
---------------------------------------------------------------------------
"""
from __future__ import annotations

import dataclasses
import typing
from collections import OrderedDict
from datetime import datetime

import numpy as np
from gymnasium import spaces
from ray.rllib.models.preprocessors import Preprocessor
from ray.rllib.utils.filter import Filter

from corl.dones.done_func_base import DoneStatusCodes
from corl.evaluation.runners.section_factories.plugins.platform_serializer import PlatformSerializer
from corl.libraries.functor import Functor


@dataclasses.dataclass
class EpisodeArtifact:
    """Artifact generated by an episode"""

    @dataclasses.dataclass
    class PolicyArtifact:
        """Postprocessed new observation, value function and value targets
        for one agent.
        """

        preprocessor: Preprocessor
        filters: Filter

    @dataclasses.dataclass
    class AgentStep:
        """Data collected from perspective a an agent at each step"""

        observations: OrderedDict
        actions: OrderedDict
        rewards: dict[str, float] | None
        total_reward: float
        observation_tensors: np.ndarray | None = None
        value_function: float | None = None

    @dataclasses.dataclass
    class Step:
        """Data collected at each timestep"""

        agents: dict[str, EpisodeArtifact.AgentStep]
        """Dictionary of data for each agent
        """

        platforms: list[dict]
        """Serialized platform data
        """

        environment_state: dict[str, typing.Any]
        """Dictionary of environment specific things
        The exact details of this field depends on the environment
        """

    @dataclasses.dataclass
    class SpaceDefinitions:
        """Original and Normalized space definitions"""

        action_space: spaces.Dict
        normalized_action_space: spaces.Dict
        observation_space: spaces.Dict
        normalized_observation_space: spaces.Dict

    @dataclasses.dataclass
    class DoneConfig:
        """Configurations that define dones"""

        task: dict[str, dict | list[Functor]]
        world: list[typing.Any] | None
        agent_dones: dict[str, list]

    test_case: int | None
    env_config: dict[str, typing.Any]
    worker_index: int
    params: dict[str, typing.Any]
    parameter_values: dict[str, typing.Any]  # full parameters set for episode (environment's + agent's local_variable_stores)
    artifacts_filenames: dict[str, str]
    start_time: datetime
    duration_sec: float
    frame_rate: float | None
    steps: list[EpisodeArtifact.Step]
    walltime_sec: list[float]
    dones: dict[str, dict[str, bool]]
    episode_state: dict[str, dict[str, DoneStatusCodes]]
    simulator_info: list[dict[str, str]]
    observation_units: dict[str, typing.Any]
    platform_to_agents: dict[str, typing.Any]
    agent_to_platforms: dict[str, typing.Any]
    done_config: EpisodeArtifact.DoneConfig
    policy_artifact: dict[str, EpisodeArtifact.PolicyArtifact]
    initial_state: OrderedDict
    platform_serializer: PlatformSerializer
    space_definitions: EpisodeArtifact.SpaceDefinitions | None = None
    algorithm_iteration: int | None = None

    def __str__(self) -> str:
        keys = ["test_case", "worker_index", "dones", "algorithm_iteration"]

        info = ", ".join([f"{key}={getattr(self, key)}" for key in keys])
        return f"<{type(self)}[{info}]>"

    def __repr__(self) -> str:
        return str(self)
