# Synthesized wrapper model file (inspect and adapt before use)
import torch
import torch.nn as nn

# --- extracted class 1 ---
class LossWeights:
    lambda_task: float = 1.0
    lambda_res: float = 0.5
    lambda_ent: float = 0.2


# --- extracted class 2 ---
class RRF_Ultra_CNN(nn.Module):
    def __init__(self, input_dim=1, output_dim=1):
        super(RRF_Ultra_CNN, self).__init__()
        self.conv1 = nn.Conv1d(input_dim, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(128*160, 256)
        self.fc2 = nn.Linear(256, output_dim)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        return torch.sigmoid(self.fc2(x))


# --- extracted class 3 ---
class SavantRRF_Gauge(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SavantRRF_Gauge, self).__init__()
        self.conv1 = nn.Conv1d(input_dim, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv1d(128, 256, kernel_size=3, padding=1)
        self.dropout = nn.Dropout(0.25)
        # The input size to fc1 is based on the output size of conv3.
        # Assuming input sequence length is 160, after 3 conv layers with kernel_size 3 and padding 1,
        # the sequence length remains 160. 256 channels * 160 length = 40960.
        self.fc1 = nn.Linear(256*160, 512) # Corrected input size based on sequence_length=160
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, output_dim)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return torch.sigmoid(self.fc3(x))


# --- extracted class 4 ---
class DiracGraphConv(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, alpha: float = 1.0, bias: bool = True):
        super().__init__()
        self.lin = nn.Linear(in_dim, out_dim, bias=bias)
        self.alpha = nn.Parameter(torch.tensor(alpha, dtype=torch.float32))
        self.bias_edge = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))

    @staticmethod
    def cosine_corr(z_i: torch.Tensor, z_j: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
        num = (z_i * z_j).sum(dim=-1)
        den = torch.clamp(z_i.norm(dim=-1) * z_j.norm(dim=-1), min=eps)
        return num / den

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        N = x.size(0)
        row, col = edge_index
        corr = self.cosine_corr(z[row], z[col])
        logits = self.alpha * corr + self.bias_edge
        device = x.device
        E = row.size(0)
        ones = torch.ones(E, device=device)
        max_per_row = torch.full((N,), -1e9, device=device)
        max_per_row = max_per_row.index_put((row,), logits, accumulate=False).scatter_reduce_(0, row, logits, reduce="amax")
        logits_centered = logits - max_per_row[row]
        exp_logits = torch.exp(logits_centered)
        denom = torch.zeros(N, device=device).index_add_(0, row, exp_logits)
        attn = exp_logits / (denom[row] + 1e-9)
        deg = torch.zeros(N, device=device).index_add_(0, row, ones)
        norm = 1.0 / torch.clamp(deg[row], min=1.0)
        msgs = norm.unsqueeze(-1) * attn.unsqueeze(-1) * x[col]
        out = torch.zeros_like(x).index_add_(0, row, msgs)
        return self.lin(out)


# --- extracted class 5 ---
class GNNDiracRRF(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int,
                 alpha_attn: float = 1.0, dropout: float = 0.1):
        super().__init__()
        self.z_dim = z_dim
        self.layers = nn.ModuleList()
        self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn))
        for _ in range(num_layers - 2):
            self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn))
        self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        h = x
        for i, layer in enumerate(self.layers):
            h = layer(h, edge_index, z)
            if i < len(self.layers) - 1:
                h = F.gelu(h)
                h = self.dropout(h)
        return h


# --- extracted class 6 ---
class LossWeights:
    lambda_task: float = 1.0
    lambda_res: float = 0.5
    lambda_ent: float = 0.2


# --- extracted class 7 ---
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


# --- extracted class 8 ---
class RRF_Ultra_CNN(nn.Module):
    def __init__(self, input_dim=1, output_dim=1):
        super(RRF_Ultra_CNN, self).__init__()
        self.conv1 = nn.Conv1d(input_dim, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(128*160, 256)
        self.fc2 = nn.Linear(256, output_dim)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        return torch.sigmoid(self.fc2(x))


# --- extracted class 9 ---
class SavantRRF_Gauge(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SavantRRF_Gauge, self).__init__()
        self.conv1 = nn.Conv1d(input_dim, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv1d(128, 256, kernel_size=3, padding=1)
        self.dropout = nn.Dropout(0.25)
        # The input size to fc1 is based on the output size of conv3.
        # Assuming input sequence length is 160, after 3 conv layers with kernel_size 3 and padding 1,
        # the sequence length remains 160. 256 channels * 160 length = 40960.
        self.fc1 = nn.Linear(256*160, 512) # Corrected input size based on sequence_length=160
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, output_dim)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return torch.sigmoid(self.fc3(x))


# --- extracted class 10 ---
class DiracGraphConv(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, alpha: float = 1.0, bias: bool = True):
        super().__init__()
        self.lin = nn.Linear(in_dim, out_dim, bias=bias)
        self.alpha = nn.Parameter(torch.tensor(alpha, dtype=torch.float32))
        self.bias_edge = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))

    @staticmethod
    def cosine_corr(z_i: torch.Tensor, z_j: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
        num = (z_i * z_j).sum(dim=-1)
        den = torch.clamp(z_i.norm(dim=-1) * z_j.norm(dim=-1), min=eps)
        return num / den

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        N = x.size(0)
        row, col = edge_index
        corr = self.cosine_corr(z[row], z[col])
        logits = self.alpha * corr + self.bias_edge
        device = x.device
        E = row.size(0)
        ones = torch.ones(E, device=device)
        max_per_row = torch.full((N,), -1e9, device=device)
        max_per_row = max_per_row.index_put((row,), logits, accumulate=False).scatter_reduce_(0, row, logits, reduce="amax")
        logits_centered = logits - max_per_row[row]
        exp_logits = torch.exp(logits_centered)
        denom = torch.zeros(N, device=device).index_add_(0, row, exp_logits)
        attn = exp_logits / (denom[row] + 1e-9)
        deg = torch.zeros(N, device=device).index_add_(0, row, ones)
        norm = 1.0 / torch.clamp(deg[row], min=1.0)
        msgs = norm.unsqueeze(-1) * attn.unsqueeze(-1) * x[col]
        out = torch.zeros_like(x).index_add_(0, row, msgs)
        return self.lin(out)


# --- extracted class 11 ---
class GNNDiracRRF(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int,
                 alpha_attn: float = 1.0, dropout: float = 0.1):
        super().__init__()
        self.z_dim = z_dim
        self.layers = nn.ModuleList()
        self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn))
        for _ in range(num_layers - 2):
            self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn))
        self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        h = x
        for i, layer in enumerate(self.layers):
            h = layer(h, edge_index, z)
            if i < len(self.layers) - 1:
                h = F.gelu(h)
                h = self.dropout(h)
        return h


# --- extracted class 12 ---
class LossWeights:
    lambda_task: float = 1.0
    lambda_res: float = 0.5
    lambda_ent: float = 0.2


# --- extracted class 13 ---
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


# --- extracted class 14 ---
class LossWeights:
    lambda_task: float = 1.0
    lambda_res: float = 0.5
    lambda_ent: float = 0.2


# --- extracted class 15 ---
class IcosahedralRRFDataset(InMemoryDataset):
    def __init__(self, num_graphs: int = 64, k_modes: int = 16, feat_dim: int = 8,
                 task_type: str = 'classification', split: str = 'train', transform=None, pre_transform=None):
        super().__init__('.', transform, pre_transform)
        self.task_type = task_type
        self.num_graphs = num_graphs
        self.k_modes = k_modes
        self.feat_dim = feat_dim

        # Generate graphs and process them
        data_list = []
        rng = np.random.default_rng(42 if split == 'train' else (43 if split == 'val' else 44))

        for i in range(num_graphs):
            G = nx.icosahedral_graph()
            n_nodes = G.number_of_nodes()

            # Build Dirac operator and compute spectral modes
            D = build_dirac_operator(G, normalize=True)
            # Use the modified dirac_eigendecomp that uses np.linalg.eigh
            vals, vecs = dirac_eigendecomp(D, k=k_modes)
            Z = node_spectral_coords_from_dirac(vecs, n_nodes)  # N x k

            # Get edge index
            edge_list = list(G.edges())
            edge_index = torch.tensor(edge_list, dtype=torch.long).t().contiguous()
            # Add reverse edges for undirected graph
            row, col = edge_index
            edge_index = torch.cat([edge_index, torch.stack([col, row], dim=0)], dim=1)

            # Generate synthetic node features (x) and labels (y)
            # Features: [n_nodes, feat_dim]
            x = torch.randn(n_nodes, feat_dim, dtype=torch.float32)

            # Labels: based on task_type
            if task_type == 'classification':
                # Example: Binary classification based on a simple rule, e.g., sum of features > threshold
                threshold = 0.0 # Example threshold
                y = (x.sum(dim=-1) > threshold).long() # [n_nodes]
            elif task_type == 'regression':
                 # Example: Regression target based on sum of features
                 y = x.sum(dim=-1) # [n_nodes]
            else:
                raise ValueError("task_type must be 'classification' or 'regression'")


            # Create Data object
            # Note: The IcosahedralRRF model expects input 'x' as [batch_size, input_dim, sequence_length],
            # edge_index [2, num_edges], and z [num_nodes, z_dim].
            # The IcosahedralRRFDataset provides batch.x [num_nodes, feat_dim], batch.edge_index [2, num_edges], and batch.U [num_nodes, k_modes].
            # There is a mismatch in the expected input format for the IcosahedralRRF model's forward pass when using the DataLoader.
            # The IcosahedralRRF expects a single batch tensor `x` for the gauge nodes, and graph data (edge_index, z) for the GNN part which operates on gauge outputs.
            # The IcosahedralRRFDataset provides node features `batch.x` that are intended as features *for the graph nodes themselves*, not as input to the gauge nodes.
            # The current IcosahedralRRF forward pass processes a single input `x` [batch_size, input_dim, sequence_length] through all gauge nodes.
            # The GNN then operates on the *outputs* of these gauge nodes, using the provided edge_index and z.

            # To use the IcosahedralRRFDataset with the current IcosahedralRRF model structure,
            # we need to map the dataset's structure to the model's expectations.
            # The dataset provides graphs, each with nodes (typically 12 for icosahedral), node features (batch.x), edge_index, and spectral coords (batch.U).
            # The IcosahedralRRF model has 12 gauge nodes, each designed to process a sequence [input_dim, sequence_length].
            # It seems there is a conceptual mismatch in how the IcosahedralRRFDataset is structured (graph-centric with node features)
            # and how the IcosahedralRRF model processes input (sequence-centric through gauge nodes first).

            # Alternative Interpretation: The IcosahedralRRFDataset is meant to provide data where each *graph* is a sample in the batch.
            # batch.x would be the concatenated node features for all graphs in the batch: [total_num_nodes_in_batch, feat_dim].
            # batch.edge_index would be the block-diagonal edge indices for all graphs: [2, total_num_edges_in_batch].
            # batch.U would be the concatenated spectral coordinates for all nodes: [total_num_nodes_in_batch, k_modes].
            # In this case, the input to the IcosahedralRRF model's forward pass is still expected to be a single tensor `x` for the gauge nodes.
            # The IcosahedralRRFDataset does *not* provide this `x` input directly in the expected format.

            # There is a fundamental incompatibility in how the IcosahedralRRFDataset provides data (graph-batching)
            # and how the IcosahedralRRF model expects input (single batch of sequences + graph data for GNN).

            # To make this cell runnable, we need to either:
            # 1. Modify the IcosahedralRRF model's forward pass to handle graph batches from DataLoader.
            # 2. Modify the IcosahedralRRFDataset or create a custom Dataset/DataLoader that provides data in the format expected by the IcosahedralRRF model.
            # 3. Use a simplified evaluation approach that aligns with the synthetic data generation method used in the training loop (single batch).

            # Given the current structure, the simplest approach to get the cell running is to align the evaluation data generation
            # with the training data generation (single synthetic batch) and evaluate on that.
            # This bypasses the DataLoader incompatibility but doesn't fully test with graph batching.

            # Let's revert to generating a single synthetic batch for evaluation, similar to training.
            # This requires defining x_val and y_val outside the DataLoader loop.

            # Reverting the evaluation loop to use the single synthetic batch approach:

            # Check if x_val and y_val are defined (from previous code cell)
            if 'x_val' not in locals() or 'y_val' not in locals():
                 # Generate synthetic validation data if not already defined
                 val_batch_size = 16 # Example validation batch size
                 x_val = torch.randn(val_batch_size, input_dim, sequence_length, dtype=torch.float32).to(device)
                 y_val = torch.randint(0, 2, (val_batch_size,), dtype=torch.long).to(device) # Binary labels
                 print("Generated synthetic validation data for evaluation.")

        # Ensure z and edge_index are on the correct device
        if 'z' in locals() and 'edge_index' in locals():
            z = z.to(device)
            edge_index = edge_index.to(device)
        else:
            print("⚠️ Warning: Graph data (z, edge_index) not found. Skipping evaluation.")
            # Exit the evaluation block if graph data is missing
            # break # This will exit the with torch.no_grad(): block - REMOVED/COMMENTED OUT DUE TO SyntaxError
            pass # Use pass instead of break to avoid SyntaxError outside a loop

        # Forward pass on validation data using the single batch
        # Pass the validation input features (x_val), edge index, and spectral coordinates (z) through the model
        val_outputs = hybrid_model(x_val, edge_index, z) # Shape: [val_batch_size, output_dim]

        # Calculate the validation loss (using BCEWithLogitsLoss as corrected in training)
        val_loss = F.binary_cross_entropy_with_logits(val_outputs.squeeze(-1), y_val.float())

        # Calculate evaluation metrics (e.g., accuracy for binary classification)
        # Convert logits to predicted class (0 or 1)
        predicted_classes = (torch.sigmoid(val_outputs.squeeze(-1)) > 0.5).long()

        # Calculate accuracy
        correct_predictions = (predicted_classes == y_val).sum().item()
        accuracy = correct_predictions / val_batch_size

        print(f'Validation Loss: {val_loss.item():.4f}, Validation Accuracy: {accuracy:.4f}')


# --- extracted class 16 ---
class DiracGraphConv(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, alpha: float = 1.0, bias: bool = True):
        super().__init__()
        self.lin = nn.Linear(in_dim, out_dim, bias=bias)
        self.alpha = nn.Parameter(torch.tensor(alpha, dtype=torch.float32))
        self.bias_edge = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))

    @staticmethod
    def cosine_corr(z_i: torch.Tensor, z_j: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
        num = (z_i * z_j).sum(dim=-1)
        den = torch.clamp(z_i.norm(dim=-1) * z_j.norm(dim=-1), min=eps)
        return num / den

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        N = x.size(0)
        row, col = edge_index
        corr = self.cosine_corr(z[row], z[col])
        logits = self.alpha * corr + self.bias_edge
        device = x.device
        E = row.size(0)
        ones = torch.ones(E, device=device)
        max_per_row = torch.full((N,), -1e9, device=device)
        max_per_row = max_per_row.index_put((row,), logits, accumulate=False).scatter_reduce_(0, row, logits, reduce="amax")
        logits_centered = logits - max_per_row[row]
        exp_logits = torch.exp(logits_centered)
        denom = torch.zeros(N, device=device).index_add_(0, row, exp_logits)
        attn = exp_logits / (denom[row] + 1e-9)
        deg = torch.zeros(N, device=device).index_add_(0, row, ones)
        norm = 1.0 / torch.clamp(deg[row], min=1.0)
        msgs = norm.unsqueeze(-1) * attn.unsqueeze(-1) * x[col]
        out = torch.zeros_like(x).index_add_(0, row, msgs)
        return self.lin(out)


# --- extracted class 17 ---
class GNNDiracRRF(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int,
                 alpha_attn: float = 1.0, dropout: float = 0.1):
        super().__init__()
        self.z_dim = z_dim
        self.layers = nn.ModuleList()
        self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn))
        for _ in range(num_layers - 2):
            self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn))
        self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        h = x
        for i, layer in enumerate(self.layers):
            h = layer(h, edge_index, z)
            if i < len(self.layers) - 1:
                h = F.gelu(h)
                h = self.dropout(h)
        return h


# --- extracted class 18 ---
class DiracGraphConv(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, alpha: float = 1.0, bias: bool = True):
        super().__init__()
        self.lin = nn.Linear(in_dim, out_dim, bias=bias)
        self.alpha = nn.Parameter(torch.tensor(alpha, dtype=torch.float32))
        self.bias_edge = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))

    @staticmethod
    def cosine_corr(z_i: torch.Tensor, z_j: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
        num = (z_i * z_i).sum(dim=-1) # Corrected dot product: z_i * z_j
        den = torch.clamp(z_i.norm(dim=-1) * z_j.norm(dim=-1), min=eps)
        return num / den

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        N = x.size(0)
        row, col = edge_index
        corr = self.cosine_corr(z[row], z[col])
        logits = self.alpha * corr + self.bias_edge
        device = x.device
        E = row.size(0)
        ones = torch.ones(E, device=device)
        max_per_row = torch.full((N,), -1e9, device=device)
        max_per_row = max_per_row.index_put((row,), logits, accumulate=False).scatter_reduce_(0, row, logits, reduce="amax")
        logits_centered = logits - max_per_row[row]
        exp_logits = torch.exp(logits_centered)
        denom = torch.zeros(N, device=device).index_add_(0, row, exp_logits)
        attn = exp_logits / (denom[row] + 1e-9)
        deg = torch.zeros(N, device=device).index_add_(0, row, ones)
        norm = 1.0 / torch.clamp(deg[row], min=1.0)
        msgs = norm.unsqueeze(-1) * attn.unsqueeze(-1) * x[col]
        out = torch.zeros_like(x).index_add_(0, row, msgs)
        return self.lin(out)


# --- extracted class 19 ---
class GNNDiracRRF(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int,
                 alpha_attn: float = 1.0, dropout: float = 0.1):
        super().__init__()
        self.z_dim = z_dim
        self.layers = nn.ModuleList()
        self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn))
        for _ in range(num_layers - 2):
            self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn))
        self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        h = x
        for i, layer in enumerate(self.layers):
            h = layer(h, edge_index, z)
            if i < len(self.layers) - 1:
                h = F.gelu(h)
                h = self.dropout(h)
        return h


# --- extracted class 20 ---
class SavantRRF_Gauge(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SavantRRF_Gauge, self).__init__()
        self.conv1 = nn.Conv1d(input_dim, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv1d(128, 256, kernel_size=3, padding=1)
        self.dropout = nn.Dropout(0.25)
        # Assuming input sequence length is 160
        self.fc1 = nn.Linear(256*160, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, output_dim)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return torch.sigmoid(self.fc3(x))


# --- extracted class 21 ---
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
        self.ethical_core = nn.Linear(12 * output_dim, output_dim)

        # Subconsciente (dodecaedro/icosaedro) using GNNDiracRRF
        # The GNN operates on the 12 gauge node outputs.
        # The input features to the GNN are the outputs of the 12 gauge nodes, shape [batch_size, output_dim].
        # For GNN layer, input is [num_nodes, in_channels] = [12, output_dim] per batch item.
        self.memory_map = GNNDiracRRF(in_dim=output_dim,
                                      hidden_dim=hidden_dim,
                                      out_dim=output_dim,
                                      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).
            stacked_outputs = torch.stack(outputs, dim=1) # [batch_size, 12, output_dim]

            gnn_outputs_list = []
            for i in range(stacked_outputs.size(0)):
                gnn_input_features_i = stacked_outputs[i]
                edge_index_i = edge_index.to(x.device)
                z_i = z.to(x.device)
                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)

            gnn_outputs_stacked = torch.stack(gnn_outputs_list, dim=0)
            aggregated_gnn_output = gnn_outputs_stacked.mean(dim=1) # [batch_size, output_dim]

            return aggregated_gnn_output # [batch_size, output_dim]

        else:
             return regulated


# --- extracted class 22 ---
class DiracGraphConv(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, alpha: float = 1.0, bias: bool = True):
        super().__init__()
        self.lin = nn.Linear(in_dim, out_dim, bias=bias)
        self.alpha = nn.Parameter(torch.tensor(alpha, dtype=torch.float32))
        self.bias_edge = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))

    @staticmethod
    def cosine_corr(z_i: torch.Tensor, z_j: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
        num = (z_i * z_j).sum(dim=-1)
        den = torch.clamp(z_i.norm(dim=-1) * z_j.norm(dim=-1), min=eps)
        return num / den

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        N = x.size(0)
        row, col = edge_index
        # Ensure z has correct shape for cosine_corr
        # z should have shape [num_nodes, z_dim]
        # x has shape [num_nodes, in_dim]
        # When called from GNNDiracRRF, num_nodes is 12 (for icosahedral)
        # z[row] and z[col] should broadcast correctly with x[col]
        corr = self.cosine_corr(z[row], z[col])
        logits = self.alpha * corr + self.bias_edge
        device = x.device
        E = row.size(0)
        ones = torch.ones(E, device=device)
        # Use scatter_reduce_ to calculate max per row
        max_per_row = torch.full((N,), -1e9, device=device)
        max_per_row = max_per_row.index_put((row,), logits, accumulate=False).scatter_reduce_(0, row, logits, reduce="amax")
        logits_centered = logits - max_per_row[row]
        exp_logits = torch.exp(logits_centered)
        denom = torch.zeros(N, device=device).index_add_(0, row, exp_logits)
        attn = exp_logits / (denom[row] + 1e-9)
        deg = torch.zeros(N, device=device).index_add_(0, row, ones)
        norm = 1.0 / torch.clamp(deg[row], min=1.0)
        msgs = norm.unsqueeze(-1) * attn.unsqueeze(-1) * x[col]
        out = torch.zeros_like(x).index_add_(0, row, msgs)
        return self.lin(out)


# --- extracted class 23 ---
class GNNDiracRRF(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int,
                 alpha_attn: float = 1.0, dropout: float = 0.1):
        super().__init__()
        self.z_dim = z_dim
        self.layers = nn.ModuleList()
        # Ensure DiracGraphConv is defined before this line
        self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn))
        for _ in range(num_layers - 2):
            self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn))
        self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        h = x
        for i, layer in enumerate(self.layers):
            h = layer(h, edge_index, z)
            if i < len(self.layers) - 1:
                h = F.gelu(h)
                h = self.dropout(h)
        return h


# --- extracted class 24 ---
class RRF_Dataset(Dataset):
    def __init__(self, strain, weights, seq_len=160): # Use seq_len=160 to match model input
        self.seq_len = seq_len
        self.strain = strain
        self.weights = weights
        print(f"Debug: RRF_Dataset __init__ - len(strain): {len(strain)}, seq_len: {self.seq_len}") # Debug print
        # Calculate n only if strain is long enough
        if len(strain) >= seq_len:
            self.n = len(strain) // seq_len
        else:
            self.n = 0 # Set n to 0 if strain is too short
        print(f"Debug: RRF_Dataset __init__ - Calculated self.n: {self.n}") # New debug print
        # Add a check to ensure there's at least one sequence
        if self.n == 0:
            raise ValueError(f"Strain data length ({len(strain)}) is less than sequence length ({seq_len}). Cannot create any samples.")


    def __len__(self):
        return self.n

    def __getitem__(self, idx):
        start = idx * self.seq_len
        # Extract the strain sequence x
        x = self.strain[start:start+self.seq_len] # Shape: [seq_len]

        # Use the mean of the provided weights as the global resonance factor w
        w = np.mean(self.weights)  # global resonance factor

        # Define the target label y as the mean of the strain sequence x, scaled by w
        # This creates a regression target derived from the strain data.
        y = np.mean(x) * w  # synthetic label (proxy resonance)

        # Convert x and y to PyTorch tensors with float dtype
        # The model expects input x as [1, seq_len] for a single sample, so add unsqueeze(0)
        return torch.tensor(x).float().unsqueeze(0), torch.tensor(y).float()



def load_model_state(path, model_instance, map_location='cpu'):
    '''Helper: load state_dict from path into model_instance (PyTorch).'''
    state = torch.load(path, map_location=map_location)
    if isinstance(state, dict) and ('state_dict' in state and isinstance(state['state_dict'], dict)):
        state = state['state_dict']
    model_instance.load_state_dict(state)
    return model_instance