import copy

import gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch_geometric.data import Data
from torch_geometric.utils import dropout_edge
from torch_scatter import scatter_mean


def get_scatter_reduce(name: str):
    name = name.lower()
    if name == "mean":
        from torch_scatter import scatter_mean
        return scatter_mean
    if name == "sum":
        from torch_scatter import scatter_add
        return scatter_add
    if name == "max":
        from torch_scatter import scatter_max
        return lambda *a, **kw: scatter_max(*a, **kw)[0]
    if name == "min":
        from torch_scatter import scatter_min
        return lambda *a, **kw: scatter_min(*a, **kw)[0]
    if name == "std":
        from torch_scatter import scatter_std
        return scatter_std
    raise ValueError(f"Unknown scatter reduce '{name}'")


# ──
#  1. LatentMLP — GNN 内部使用的 MLP（保持隐层维度不变）
# ──
class LatentMLP(nn.Module):
    """
    MLP that operates entirely in latent space (dim in == dim out == latent_dim).
    Used inside EdgeModule and NodeModule.
    """

    def __init__(self, in_features: int, latent_dim: int, config: dict):
        super().__init__()
        num_layers = config.get("num_layers", 2)
        activation = config.get("activation_function", "leakyrelu").lower()
        add_output = config.get("add_output_layer", False)

        layers = []
        prev_dim = in_features
        for i in range(num_layers):
            layers.append(nn.Linear(prev_dim, latent_dim))
            layers.append(_get_activation(activation))
            prev_dim = latent_dim

        if add_output:
            layers.append(nn.Linear(prev_dim, latent_dim))

        self.mlp = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.mlp(x)


def _get_activation(name: str) -> nn.Module:
    name = name.lower()
    if name == "relu":
        return nn.ReLU()
    elif name == "leakyrelu":
        return nn.LeakyReLU()
    elif name == "elu":
        return nn.ELU()
    elif name in ("swish", "silu"):
        return nn.SiLU()
    elif name == "mish":
        return nn.Mish()
    elif name == "gelu":
        return nn.GELU()
    elif name == "tanh":
        return nn.Tanh()
    raise ValueError(f"Unknown activation '{name}'")


# ──
#  2. EdgeModule — 边更新：MLP([src_node | dst_node | edge_attr])
# ──
class EdgeModule(nn.Module):
    """Update edge features from sender node, receiver node, and existing edge features."""

    def __init__(self, latent_dim: int, mlp_config: dict):
        super().__init__()
        in_features = 3 * latent_dim  # [src_node, dst_node, edge_attr]
        self.mlp = LatentMLP(in_features, latent_dim, mlp_config)

    def forward(self, graph: Data):
        src, dst = graph.edge_index
        agg = torch.cat([graph.x[src], graph.x[dst], graph.edge_attr], dim=-1)
        graph.edge_attr = self.mlp(agg)


# ──
#  4. NodeModule — 节点更新：MLP([node | scatter(入边)])
# ──
class NodeModule(nn.Module):
    """Update node features from own features and aggregated incoming edge features."""

    def __init__(self, latent_dim: int, mlp_config: dict, scatter_reducer):
        super().__init__()
        in_features = 2 * latent_dim  # [node, aggregated_edges]
        self.mlp = LatentMLP(in_features, latent_dim, mlp_config)
        self.scatter = scatter_reducer

    def forward(self, graph: Data):
        _, dst = graph.edge_index
        agg_edges = self.scatter(
            graph.edge_attr, dst, dim=0, dim_size=graph.x.shape[0]
        )
        agg = torch.cat([graph.x, agg_edges], dim=-1)
        graph.x = self.mlp(agg)


# ──
#  5. MessagePassingStep — 单步消息传递
# ──
class MessagePassingStep(nn.Module):
    """
    One full message-passing step:
      1. Edge update
      2. Edge inner residual + LayerNorm
      3. Node update
      4. Node inner residual + LayerNorm
    """

    def __init__(self, latent_dim: int, stack_config: dict, scatter_reducer):
        super().__init__()
        mlp_config = stack_config["mlp"]

        self.edge_module = EdgeModule(latent_dim, mlp_config)
        self.node_module = NodeModule(latent_dim, mlp_config, scatter_reducer)

        self.node_ln = nn.LayerNorm(latent_dim)
        self.edge_ln = nn.LayerNorm(latent_dim)

    def forward(self, graph: Data):
        old_x = graph.x
        old_edge = graph.edge_attr

        # Edge update
        self.edge_module(graph)
        graph.edge_attr = self.edge_ln(graph.edge_attr + old_edge)

        # Node update
        self.node_module(graph)
        graph.x = self.node_ln(graph.x + old_x)


# ──
#  6. MessagePassingStack — 堆叠 N 个 Step
# ──
class MessagePassingStack(nn.Module):
    """Stack of multiple MessagePassingSteps with optional step repeats."""

    def __init__(self, latent_dim: int, stack_config: dict, scatter_reducer):
        super().__init__()
        num_steps = stack_config.get("num_steps", 2)
        self.num_step_repeats = stack_config.get("num_step_repeats", 1)
        self.steps = nn.ModuleList(
            [
                MessagePassingStep(latent_dim, stack_config, scatter_reducer)
                for _ in range(num_steps)
            ]
        )

    def forward(self, graph: Data):
        for step in self.steps:
            for _ in range(self.num_step_repeats):
                step(graph)


# ──
#  7. MessagePassingBase — GNN 基座
# ──
class MessagePassingBase(nn.Module):
    """
    Full GNN base: Linear → Stack → unpacked output.
    Returns (node_features_dict, edge_features, None, batch_dict)
    """

    def __init__(
        self,
        in_node_features: int,
        in_edge_features: int,
        latent_dim: int,
        base_config: dict,
        device=None,
    ):
        super().__init__()
        self.edge_dropout = base_config.get("edge_dropout", 0.0)
        self.create_copy = base_config.get("create_graph_copy", True)

        scatter_name = base_config.get("scatter_reduce", "mean")
        self.scatter_reducer = get_scatter_reduce(scatter_name)

        self.node_embedding = nn.Linear(in_node_features, latent_dim)
        self.edge_embedding = nn.Linear(in_edge_features, latent_dim)

        # Stack
        stack_config = base_config.get("stack", {})
        self.stack = MessagePassingStack(latent_dim, stack_config, self.scatter_reducer)

        if device is not None:
            self.to(device)

    def forward(self, graph: Data):
        if self.create_copy:
            graph = copy.deepcopy(graph)

        # Edge dropout (training only)
        if self.edge_dropout > 0 and self.training:
            graph.edge_index, mask = dropout_edge(
                graph.edge_index, p=self.edge_dropout, training=True
            )
            graph.edge_attr = graph.edge_attr[mask]

        # Embed
        graph.x = self.node_embedding(graph.x)
        graph.edge_attr = self.edge_embedding(graph.edge_attr)

        # Message passing
        self.stack(graph)

        # Unpack
        node_name = "element"  # homogeneous graph node type for mesh refinement
        batch = (
            graph.batch
            if hasattr(graph, "batch") and graph.batch is not None
            else torch.zeros(graph.x.shape[0], dtype=torch.long, device=graph.x.device)
        )

        edge_key = f"{node_name}2{node_name}"
        return (
            {node_name: graph.x},
            {
                edge_key: {
                    "edge_index": graph.edge_index.long(),
                    "edge_attr": graph.edge_attr,
                }
            },
            None,
            {node_name: batch},
        )


# ──
#  8. MLP — Actor/Critic 头使用的 MLP
# ──
class MLP(nn.Module):
    """Feedforward MLP for actor/critic heads."""

    def __init__(
        self,
        in_features: int,
        config: dict,
        latent_dim: int = None,
        out_features: int = None,
        device=None,
    ):
        super().__init__()
        activation = config.get("activation_function", "tanh").lower()
        num_layers = config.get("num_layers", 2)

        layers = []
        prev = in_features
        dim = latent_dim or 64
        for _ in range(num_layers):
            layers.append(nn.Linear(prev, dim))
            layers.append(_get_activation(activation))
            prev = dim

        if out_features is not None:
            layers.append(nn.Linear(prev, out_features))
            self._out_features = out_features
        else:
            self._out_features = prev

        self.net = nn.Sequential(*layers)
        if device is not None:
            self.to(device)

    @property
    def out_features(self) -> int:
        return self._out_features

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


# ──
#  9. ActorCritic — PPO Actor-Critic 网络
# ──
def create_model(env, network_config: dict, ppo_config: dict, device=None):
    """Factory function: create Actor-Critic model from environment and configs."""
    return ActorCritic(
        environment=env,
        network_config=network_config,
        ppo_config=ppo_config,
        device=device,
    )


class ActorCritic(nn.Module):

    def __init__(
        self, environment, network_config: dict, ppo_config: dict, device=None
    ):
        super().__init__()
        latent_dim = network_config.get("latent_dimension", 64)
        base_config = network_config.get("base", {})
        train_config = network_config.get("training", {})
        actor_cfg = network_config.get("actor", {}).get("mlp", {})
        critic_cfg = network_config.get("critic", {}).get("mlp", {})

        self.value_function_aggr = ppo_config.get("value_function_aggr", "spatial")
        self.agent_node_type = "element"

        self.base = MessagePassingBase(
            in_node_features=environment.num_node_features,
            in_edge_features=environment.num_edge_features,
            latent_dim=latent_dim,
            base_config=base_config,
            device=device,
        )

        self.policy_mlp = MLP(
            latent_dim, actor_cfg, latent_dim=latent_dim, device=device
        )
        action_dim = environment.action_dimension
        if isinstance(environment._action_space, gym.spaces.Box):
            from stable_baselines3.common.distributions import DiagGaussianDistribution

            self.action_dist = DiagGaussianDistribution(action_dim)
            self.action_out, self.log_std = self.action_dist.proba_distribution_net(
                latent_dim=self.policy_mlp.out_features,
                log_std_init=ppo_config.get("initial_log_std", 0.0),
            )
        else:
            from stable_baselines3.common.distributions import CategoricalDistribution

            self.action_dist = CategoricalDistribution(action_dim)
            self.action_out = self.action_dist.proba_distribution_net(
                latent_dim=self.policy_mlp.out_features
            )
            self.log_std = None

        self.value_mlp = MLP(
            latent_dim, critic_cfg, latent_dim=latent_dim, out_features=1, device=device
        )

        self._setup_optimizer(train_config)

        if device is not None:
            self.to(device)

    def _setup_optimizer(self, train_config: dict):
        lr = train_config.get("learning_rate", 3e-4)
        wd = train_config.get("l2_norm", 0)
        params = list(self.parameters())
        if self.log_std is not None and not any(p is self.log_std for p in params):
            params.append(self.log_std)
        self.optimizer = optim.Adam(params, lr=lr, weight_decay=wd)
        sched_rate = train_config.get("lr_decay", train_config.get("lr_scheduling_rate", 1))
        self.lr_scheduler = (
            optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=sched_rate)
            if sched_rate is not None and sched_rate < 1
            else None
        )

    @property
    def device(self):
        return next(self.parameters()).device

    def _encode(self, observations):
        """Run shared GNN backbone once, return (shared_features, batch_indices)."""
        observations = observations.to(self.device)
        node_feats, _, _, batches = self.base(observations)
        batch = batches[self.agent_node_type]
        feats = node_feats[self.agent_node_type]
        return feats, batch

    def _make_distribution(self, latent_pi):
        mean_actions = self.action_out(latent_pi)
        if self.log_std is not None:
            return self.action_dist.proba_distribution(mean_actions, self.log_std)
        return self.action_dist.proba_distribution(mean_actions)

    def _aggregate_values(self, values, batch):
        if self.value_function_aggr == "mean":
            return scatter_mean(values, batch, dim=0)
        elif self.value_function_aggr == "sum":
            from torch_scatter import scatter_add
            return scatter_add(values, batch, dim=0)
        elif self.value_function_aggr == "max":
            from torch_scatter import scatter_max
            return scatter_max(values, batch, dim=0)[0]
        return values

    def forward(self, observations, deterministic: bool = False):
        shared_feats, batch = self._encode(observations)
        dist = self._make_distribution(self.policy_mlp(shared_feats))
        actions = dist.get_actions(deterministic=deterministic)
        log_probs = dist.log_prob(actions)
        values = self._aggregate_values(self.value_mlp(shared_feats).squeeze(-1), batch)
        return actions, values, log_probs

    def evaluate_actions(self, observations, actions):
        actions = actions.to(self.device)
        shared_feats, batch = self._encode(observations)
        dist = self._make_distribution(self.policy_mlp(shared_feats))
        values = self._aggregate_values(self.value_mlp(shared_feats).squeeze(-1), batch)
        return values, dist.log_prob(actions), dist.entropy()
