"""PCFA Policy wrapper.

Provides a policy-style interface around the non-learning PCFA market algorithm
so that it can be plugged into the existing runner architecture.
"""

import numpy as np
from algorithms.algorithm.pcfa import run_pcfa_policy


class PCFAPolicy:
    """Non-learning policy wrapper for PCFA.

    The interface mirrors GAMAPPOPolicy/RMAPPOPolicy where convenient, but the
    underlying algorithm is deterministic and does not require gradients or
    training updates.
    """

    def __init__(self, args, obs_space, cent_obs_space, act_space, device=None):
        self.args = args
        self.obs_space = obs_space
        self.cent_obs_space = cent_obs_space
        self.act_space = act_space
        self.device = device

        self.num_agents = getattr(args, "num_agents", 1)
        self.num_tasks = act_space.n
        self.hidden_size = getattr(args, "hidden_size", 64)
        self.recurrent_N = getattr(args, "recurrent_N", 1)

        # Will be set by runner after env is available.
        self.idle_task_id = 0
        self.explore_task_id = 0
        self.resupply_task_id = 0
        self.delivery_task_num = None
        self.explore_task_start = 0
        self.explore_task_end = 1
        self.use_coalition_delivery = False

    def set_task_ids(self, idle_task_id, explore_task_id, resupply_task_id):
        """Set special task IDs used by the environment."""
        self.idle_task_id = idle_task_id
        self.explore_task_id = explore_task_id
        self.resupply_task_id = resupply_task_id
        self.explore_task_start = explore_task_id
        self.explore_task_end = explore_task_id + 1

    def set_task_layout(
        self,
        *,
        delivery_task_num,
        explore_task_start,
        explore_task_end,
        idle_task_id,
        resupply_task_id,
        use_coalition_delivery,
    ):
        """Record the environment task-slot layout.

        DelivCoal exposes several observer-relative Explore-sector slots.  The
        historical PCFA wrapper only received one Explore ID, which caused the
        remaining sector slots to be interpreted as delivery tasks.
        """
        self.delivery_task_num = int(delivery_task_num)
        self.explore_task_start = int(explore_task_start)
        self.explore_task_end = int(explore_task_end)
        self.explore_task_id = self.explore_task_start
        self.idle_task_id = int(idle_task_id)
        self.resupply_task_id = int(resupply_task_id)
        self.use_coalition_delivery = bool(use_coalition_delivery)

    def _parse_batch_args(self, obs, available_actions, active_masks):
        """Return per-env args and detect batch size.

        Returns:
            batch_size: int
            obs_list: list of obs dicts with no batch dimension
            avail_list: list of available_actions arrays
            active_list: list of active_masks arrays
        """
        batch_size = obs["ego"].shape[0]

        obs_list = []
        for b in range(batch_size):
            obs_list.append({k: v[b] for k, v in obs.items()})

        if available_actions is None:
            avail_list = [None] * batch_size
        else:
            avail_list = [available_actions[b] for b in range(batch_size)]

        if active_masks is None:
            active_list = [np.ones(self.num_agents, dtype=np.float32)] * batch_size
        else:
            active_list = [active_masks[b] for b in range(batch_size)]

        return batch_size, obs_list, avail_list, active_list

    def get_actions(
        self,
        cent_obs,
        obs,
        rnn_states_actor,
        rnn_states_critic,
        masks,
        available_actions=None,
        active_masks=None,
        deterministic=False,
    ):
        """Compute actions for the given observations.

        Args follow the MAPPO policy signature for compatibility, but rnn states,
        masks, and cent_obs are ignored.

        Returns:
            values: zeros (batch, n_agents, 1)
            actions: (batch, n_agents) int
            action_log_probs: zeros (batch, n_agents, 1)
            rnn_states_actor: same shape as input, unchanged
            rnn_states_critic: same shape as input, unchanged
        """
        batch_size, obs_list, avail_list, active_list = self._parse_batch_args(
            obs, available_actions, active_masks
        )

        actions_list = []
        for b in range(batch_size):
            acts = run_pcfa_policy(
                obs_list[b],
                active_list[b],
                avail_list[b],
                self.args,
                self.idle_task_id,
                self.explore_task_id,
                self.resupply_task_id,
                delivery_task_num=self.delivery_task_num,
                explore_task_start=self.explore_task_start,
                explore_task_end=self.explore_task_end,
                use_coalition_delivery=self.use_coalition_delivery,
            )
            actions_list.append(acts)

        actions = np.stack(actions_list, axis=0)  # (batch, n_agents)

        values = np.zeros((batch_size, self.num_agents, 1), dtype=np.float32)
        action_log_probs = np.zeros((batch_size, self.num_agents, 1), dtype=np.float32)

        return values, actions, action_log_probs, rnn_states_actor, rnn_states_critic

    def get_values(self, cent_obs, rnn_states_critic, masks):
        """Return dummy values."""
        batch_size = cent_obs.shape[0] if hasattr(cent_obs, "shape") else 1
        return np.zeros((batch_size, self.num_agents, 1), dtype=np.float32)

    def eval_actions(
        self, cent_obs, obs, rnn_states_actor, rnn_states_critic, action, masks,
        available_actions=None, active_masks=None,
    ):
        """Return dummy evaluation quantities."""
        batch_size = action.shape[0]
        action_log_probs = np.zeros((batch_size, self.num_agents, 1), dtype=np.float32)
        dist_entropy = np.zeros((batch_size, self.num_agents, 1), dtype=np.float32)
        values = np.zeros((batch_size, self.num_agents, 1), dtype=np.float32)
        return values, action_log_probs, dist_entropy

    def prep_rollout(self):
        pass

    def prep_training(self):
        pass

    def train(self, *args, **kwargs):
        pass

    def save(self, save_dir, episode):
        pass

    def restore(self, model_dir):
        pass

    def get_max_std(self):
        return 0.0

    def update_actor(self, *args, **kwargs):
        pass

    def update_critic(self, *args, **kwargs):
        pass

    @property
    def actor(self):
        return self

    @property
    def critic(self):
        return self

    def parameters(self):
        return []

    def named_parameters(self):
        return {}

    def state_dict(self):
        return {}

    def load_state_dict(self, state_dict):
        pass