import numpy as np
from gym import spaces
import torch


POLICY_OBS_KEYS = ("ego", "tasks", "others")
POLICY_OBS_FEATURE_RANKS = {
    "ego": 1,
    "tasks": 2,
    "others": 2,
    "reconfiguration": 1,
    "gat": 1,
}

CENTRAL_STATE_KEYS = ("agents", "tasks", "global", "focal")
CENTRAL_STATE_FEATURE_RANKS = {
    "agents": 2,
    "tasks": 2,
    "global": 1,
    "focal": 2,
}


def _flat_dim_from_shape(shape):
    return int(np.prod(shape))


def flat_obs_dim(obs_space):
    if obs_space.__class__.__name__ == "Dict":
        return sum(_flat_dim_from_shape(obs_space.spaces[k].shape) for k in POLICY_OBS_KEYS)
    return int(np.prod(obs_space.shape))


def flat_obs_space(obs_space):
    if obs_space.__class__.__name__ != "Dict":
        return obs_space
    return spaces.Box(low=-np.inf, high=np.inf, shape=(flat_obs_dim(obs_space),), dtype=np.float32)


def flat_state_space(state_space):
    if state_space.__class__.__name__ != "Dict":
        return state_space
    dim = sum(_flat_dim_from_shape(state_space.spaces[k].shape) for k in CENTRAL_STATE_KEYS)
    return spaces.Box(low=-np.inf, high=np.inf, shape=(dim,), dtype=np.float32)


def flatten_obs_dict(obs):
    if not isinstance(obs, dict):
        return obs
    keys = [k for k in POLICY_OBS_KEYS if k in obs]
    if "gat" in obs:
        keys.append("gat")

    use_torch = any(isinstance(obs[k], torch.Tensor) for k in keys)
    device = next((obs[k].device for k in keys if isinstance(obs[k], torch.Tensor)), None)
    parts = []
    for key in keys:
        value = obs[key]
        if use_torch:
            value = value if isinstance(value, torch.Tensor) else torch.as_tensor(value, dtype=torch.float32, device=device)
            value = value.float()
        else:
            value = np.asarray(value, dtype=np.float32)
        feature_rank = POLICY_OBS_FEATURE_RANKS[key]
        prefix_shape = value.shape[:-feature_rank]
        parts.append(value.reshape(*prefix_shape, -1))
    if use_torch:
        return torch.cat(parts, dim=-1)
    return np.concatenate(parts, axis=-1)


def flatten_state_dict(state):
    if not isinstance(state, dict):
        return state

    keys = [k for k in CENTRAL_STATE_KEYS if k in state]
    use_torch = any(isinstance(state[k], torch.Tensor) for k in keys)
    device = next((state[k].device for k in keys if isinstance(state[k], torch.Tensor)), None)
    parts = []
    for key in keys:
        value = state[key]
        if use_torch:
            value = value if isinstance(value, torch.Tensor) else torch.as_tensor(
                value, dtype=torch.float32, device=device
            )
            value = value.float()
        else:
            value = np.asarray(value, dtype=np.float32)
        prefix_shape = value.shape[:-CENTRAL_STATE_FEATURE_RANKS[key]]
        parts.append(value.reshape(*prefix_shape, -1))
    if use_torch:
        return torch.cat(parts, dim=-1)
    return np.concatenate(parts, axis=-1)