from gymnasium.spaces import Box
import numpy as np
from typing import List

from ray.rllib.examples.policy.random_policy import RandomPolicy
from ray.rllib.policy.policy import Policy
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.policy.view_requirement import ViewRequirement
from ray.rllib.utils.annotations import override
from ray.rllib.utils.typing import TensorType


class EpisodeEnvAwareLSTMPolicy(RandomPolicy):
    """A Policy that always knows the current EpisodeID and EnvID and
    returns these in its actions."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.state_space = Box(-1.0, 1.0, (1,))

        class _fake_model:
            def __init__(self, state_space, action_space):
                self.view_requirements = {
                    SampleBatch.AGENT_INDEX: ViewRequirement(),
                    SampleBatch.EPS_ID: ViewRequirement(),
                    "env_id": ViewRequirement(),
                    "t": ViewRequirement(),
                    SampleBatch.OBS: ViewRequirement(),
                    SampleBatch.PREV_ACTIONS: ViewRequirement(
                        SampleBatch.ACTIONS, space=action_space, shift=-1
                    ),
                    SampleBatch.PREV_REWARDS: ViewRequirement(
                        SampleBatch.REWARDS, shift=-1
                    ),
                }
                self.time_major = True
                for i in range(2):
                    self.view_requirements["state_in_{}".format(i)] = ViewRequirement(
                        "state_out_{}".format(i), shift=-1, space=state_space
                    )
                    self.view_requirements["state_out_{}".format(i)] = ViewRequirement(
                        space=state_space
                    )

        self.model = _fake_model(self.state_space, self.action_space)

        self.view_requirements = dict(
            **{
                SampleBatch.NEXT_OBS: ViewRequirement(
                    SampleBatch.OBS,
                    shift=1,
                    used_for_compute_actions=False,
                ),
                SampleBatch.ACTIONS: ViewRequirement(space=self.action_space),
                SampleBatch.REWARDS: ViewRequirement(),
                SampleBatch.TERMINATEDS: ViewRequirement(),
                SampleBatch.TRUNCATEDS: ViewRequirement(),
                SampleBatch.UNROLL_ID: ViewRequirement(),
            },
            **self.model.view_requirements
        )

    @override(Policy)
    def is_recurrent(self):
        return True

    @override(Policy)
    def get_initial_state(self) -> List[TensorType]:
        return [
            np.zeros(self.view_requirements["state_in_0"].space.shape),
            np.zeros(self.view_requirements["state_in_1"].space.shape),
        ]

    @override(Policy)
    def compute_actions_from_input_dict(
        self, input_dict, explore=None, timestep=None, **kwargs
    ):
        ts = input_dict["t"]
        print(ts)
        # Always return [episodeID, envID] as actions.
        actions = np.array(
            [
                [
                    input_dict[SampleBatch.AGENT_INDEX][i],
                    input_dict[SampleBatch.EPS_ID][i],
                    input_dict["env_id"][i],
                ]
                for i, _ in enumerate(input_dict["obs"])
            ]
        )
        states = [
            np.array([[ts[i]] for i in range(len(input_dict["obs"]))]) for _ in range(2)
        ]
        return actions, states, {}

    @override(Policy)
    def postprocess_trajectory(
        self, sample_batch, other_agent_batches=None, episode=None
    ):
        sample_batch["2xobs"] = sample_batch["obs"] * 2.0
        return sample_batch


class EpisodeEnvAwareAttentionPolicy(RandomPolicy):
    """A Policy that always knows the current EpisodeID and EnvID and
    returns these in its actions."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.state_space = Box(-1.0, 1.0, (1,))
        self.config["model"] = {"max_seq_len": 50}

        class _fake_model:
            def __init__(self, state_space, config):
                self.state_space = state_space
                self.view_requirements = {
                    SampleBatch.AGENT_INDEX: ViewRequirement(),
                    SampleBatch.EPS_ID: ViewRequirement(),
                    "env_id": ViewRequirement(),
                    "t": ViewRequirement(),
                    SampleBatch.OBS: ViewRequirement(),
                    "state_in_0": ViewRequirement(
                        "state_out_0",
                        # Provide state outs -50 to -1 as "state-in".
                        shift="-50:-1",
                        # Repeat the incoming state every n time steps (usually max seq
                        # len).
                        batch_repeat_value=config["model"]["max_seq_len"],
                        space=state_space,
                    ),
                    "state_out_0": ViewRequirement(
                        space=state_space, used_for_compute_actions=False
                    ),
                }

            # We need to provide at least an initial state if we want agent collector
            # to build episodes with state_in_ and state_out_
            def get_initial_state(self):
                return [self.state_space.sample()]

        self.model = _fake_model(self.state_space, self.config)

        self.view_requirements = dict(
            super()._get_default_view_requirements(), **self.model.view_requirements
        )

    @override(Policy)
    def get_initial_state(self):
        return self.model.get_initial_state()

    @override(Policy)
    def is_recurrent(self):
        return True

    @override(Policy)
    def compute_actions_from_input_dict(
        self, input_dict, explore=None, timestep=None, **kwargs
    ):
        ts = input_dict["t"]
        print(ts)
        # Always return [episodeID, envID] as actions.
        actions = np.array(
            [
                [
                    input_dict[SampleBatch.AGENT_INDEX][i],
                    input_dict[SampleBatch.EPS_ID][i],
                    input_dict["env_id"][i],
                ]
                for i, _ in enumerate(input_dict["obs"])
            ]
        )
        states = [np.array([[ts[i]] for i in range(len(input_dict["obs"]))])]
        self.global_timestep += 1
        return actions, states, {}

    @override(Policy)
    def postprocess_trajectory(
        self, sample_batch, other_agent_batches=None, episode=None
    ):
        sample_batch["3xobs"] = sample_batch["obs"] * 3.0
        return sample_batch
