# Auto-extracted class source (static)

class IcosahedralRRF(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, gnn_num_layers=2, gnn_z_dim=16, gnn_alpha_attn=1.0, gnn_dropout=0.1):
        super(IcosahedralRRF, self).__init__()
        # 12 nodos gauge
        self.nodes = nn.ModuleList([
            SavantRRF_Gauge(input_dim, hidden_dim, output_dim) for _ in range(12)
        ])
        # Núcleo ético
        # The input to ethical_core is the concatenation of the outputs of the 12 gauge nodes.
        # Each gauge node outputs a tensor of shape [batch_size, output_dim].
        # Concatenating these along dim=1 results in a shape [batch_size, 12 * output_dim].
        self.ethical_core = nn.Linear(12 * output_dim, output_dim)

        # Subconsciente (dodecaedro) using GNNDiracRRF
        # The input dimension (in_dim) for the GNN should match the feature dimension of its input nodes.
        # There's ambiguity in the original code about what the GNN's nodes and features are.
        # Interpretation 1 (based on original code passing 'regulated'): GNN operates on 'batch_size' nodes, with 'output_dim' features. in_dim = output_dim.
        # Interpretation 2 (more conventional for graph on icosahedron/dodecahedron): GNN operates on 12 or 20 nodes, with features derived from gauge outputs.
        # Let's assume interpretation 2, where the GNN operates on the 12 gauge nodes.
        # The features for each of these 12 nodes would be the output of the corresponding gauge node, shape [batch_size, output_dim].
        # For a GNN layer expecting [num_nodes, in_channels], the input should be [12, output_dim] per batch item.
        # This means the GNN's in_dim should be output_dim. This matches the current GNN init below.
        # The GNN's out_dim should match the desired output feature dimension per node (e.g., output_dim).
        # The number of nodes for the GNN is 12 (for icosahedral).

        # Let's define the memory_map GNN assuming it operates on the 12 gauge nodes.
        # The input features to the GNN will be the outputs of the 12 gauge nodes.
        # Each gauge node outputs a tensor of shape [batch_size, output_dim].
        # We will treat output_dim as the feature dimension for the GNN nodes (the 12 gauge nodes).
        # So, in_dim for GNN = output_dim.
        # The GNN will output features for each of the 12 nodes. Let's assume out_dim for GNN is also output_dim.
        self.memory_map = GNNDiracRRF(in_dim=output_dim, # Feature dimension for GNN nodes (output_dim of gauge nodes)
                                      hidden_dim=hidden_dim,
                                      out_dim=output_dim, # Output feature dimension per GNN node
                                      num_layers=gnn_num_layers,
                                      z_dim=gnn_z_dim,
                                      alpha_attn=gnn_alpha_attn,
                                      dropout=gnn_dropout)


    def forward(self, x, edge_index=None, z=None):
        # x is the input to the gauge nodes, shape [batch_size, input_dim, sequence_length]
        outputs = [node(x) for node in self.nodes]
        # outputs is a list of 12 tensors, each [batch_size, output_dim]

        # Concatenate outputs for the ethical core
        concat = torch.cat(outputs, dim=1) # [batch_size, 12 * output_dim]
        regulated = torch.sigmoid(self.ethical_core(concat)) # [batch_size, output_dim]

        # GNN operation on the 12 gauge nodes
        if edge_index is not None and z is not None:
            # Prepare input for the GNN: Features for the 12 nodes (the gauge node outputs).
            # Stack the outputs to get [batch_size, 12, output_dim]
            stacked_outputs = torch.stack(outputs, dim=1) # [batch_size, 12, output_dim]

            # Reshape for GNN input: [num_nodes, in_channels] = [12, output_dim] per batch item.
            # Need to process batch items. Simplest is to iterate.
            # A more efficient way is to use torch_geometric.data.Batch

            gnn_outputs_list = []
            for i in range(stacked_outputs.size(0)):
                # GNN input features for this batch item: [12, output_dim]
                gnn_input_features_i = stacked_outputs[i]

                # Ensure edge_index and z are on the correct device
                edge_index_i = edge_index.to(x.device)
                z_i = z.to(x.device)

                # GNN forward pass for one batch item
                gnn_output_i = self.memory_map(gnn_input_features_i, edge_index_i, z_i) # [12, output_dim]
                gnn_outputs_list.append(gnn_output_i)

            # Stack GNN outputs back into a batch tensor: [batch_size, 12, output_dim]
            gnn_outputs_stacked = torch.stack(gnn_outputs_list, dim=0)

            # Now, how to combine the GNN output [batch_size, 12, output_dim] with the 'regulated' output [batch_size, output_dim]?
            # The original model returned just 'regulated'.
            # A simple approach is to maybe combine them, e.g., add, concatenate, or use the GNN output as a modulation.
            # Let's stick to returning the aggregated GNN output as the final output when GNN is used.
            # This changes the model's behavior compared to the original.

            # Alternative: The GNN output modulates the 'regulated' output.
            # E.g., regulated * sigmoid(aggregated_gnn_output) or similar.
            # Let's stick to returning the aggregated GNN output when edge_index and z are provided,
            # and the original 'regulated' output otherwise. This seems the most direct path based on the conditional in the original forward.

            # Aggregate the 12 nodes' outputs from the GNN
            aggregated_gnn_output = gnn_outputs_stacked.mean(dim=1) # [batch_size, output_dim]

            return aggregated_gnn_output # [batch_size, output_dim]

        else:
             # If edge_index and z are not provided, return the output of the ethical core as before.
             return regulated