def component_name_to_indices(component_name):
    component_order = ['M0', 'M1', 'M2', 'M3', 'M4', 'M5', 'M6', 'M7', 'C0', 'I0', 'V1']
    return component_order.index(component_name) + 7  # Adjust 7 based on your actual structure

class EnhancedGNNModelWithSharedParams(torch.nn.Module):
    def __init__(self, num_node_features, num_edge_features, num_out_features):
        print("EnhancedGNNModelWithSharedParams __init__ arguments:", num_node_features, num_edge_features, num_out_features)
        super(EnhancedGNNModelWithSharedParams, self).__init__()
        self.conv1 = GATConv(num_node_features, 8, heads=8, dropout=0.6)

        nn_model = Sequential(Linear(num_edge_features, 32), ReLU(), Linear(32, 8 * 8))
        self.conv2 = NNConv(8 * 8, num_out_features, nn_model, aggr='mean')

    def forward(self, node_features, edge_index, edge_attr):
        # Apply adjustment of shared parameters integrated within the model class
        modified_node_features = self.adjust_node_features(node_features)

        x = F.dropout(modified_node_features, p=0.6, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index, edge_attr)
        print(f"GNN Model Output Shape: {x.shape}")  # Debug print statement
        return x
    
    def adjust_node_features(self, node_features_tensor):
        # Filters out component nodes based on the ‘device_type’ feature
        component_mask = node_features_tensor[:, 0] == 0.0
        node_features = node_features_tensor[component_mask]
    
        # Pairs for shared ‘w_value’ and ‘l_value’
        pairs = [('M0', 'M1'), ('M2', 'M3'), ('M4', 'M7')]
        
        for pair in pairs:
            # Identify indices for the pair components
            indices = [component_name_to_indices(component) for component in pair]
    
            value_indices = [18, 19]  # Assuming these are the indices for ‘w_value’, ‘l_value’
            for i in indices[1:]:  # Start from the second component of the pair
                first_component_mask = node_features[:, indices[0]] == 1.0
                other_component_mask = node_features[:, i] == 1.0
                if torch.any(first_component_mask) and torch.any(other_component_mask):
                    # Adjust values of the other components within the pair to match the first
                    #other_values = node_features[other_component_mask][:, value_indices]
                    #first_values = node_features[first_component_mask][:, value_indices][0]
                    #other_values[:] = first_values 
                    node_features[other_component_mask][:, value_indices] = node_features[first_component_mask][:, value_indices][0]
      
        return node_features

class Actor(nn.Module):
    def __init__(self, num_input_features, action_dim, std):
        super(Actor, self).__init__()
        print("Actor __init__ arguments:", num_input_features, action_dim, std)
        # Define the actor model
        self.actor = nn.Sequential(
            nn.Linear(num_input_features, 64),
            nn.Tanh(),
            nn.Linear(64, 64),
            nn.Tanh(),
            nn.Linear(64, action_dim),
        )
        self.log_std = nn.Parameter(torch.ones(1, action_dim) * std)
        
    def forward(self, x):
        mean = self.actor(x)
        std = self.log_std.exp().expand_as(mean)
        print(f"Actor Model Output Shape (Mean): {mean.shape}")  # Debug print statement
        return mean, std

class Critic(nn.Module):
    def __init__(self, num_input_features):
        super(Critic, self).__init__()
        # Define the critic model
        self.critic = nn.Sequential(
            nn.Linear(num_input_features, 64),
            nn.Tanh(),
            nn.Linear(64, 64),
            nn.Tanh(),
            nn.Linear(64, 1)
        )
        
    def forward(self, x):
        value = self.critic(x)
        print(f"Critic Model Output Shape (Value): {value.shape}")  # Debug print statement
        return value

class PPOAgent:
    def __init__(self, actor_class, critic_class, gnn_model, action_dim, bounds_low, bounds_high,
                 lr_actor=3e-4, lr_critic=1e-3, gamma=0.99, lamda=0.95, epsilon=0.2, std=0.0):
        print("gnn_model", gnn_model)
        self.actor = actor_class(gnn_model.conv2.out_channels, action_dim, std)  # Initialize actor        
        self.critic = critic_class(gnn_model.conv2.out_channels)  # Initialize critic
        self.gnn_model = gnn_model  # GNN model instance
        self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=lr_actor)
        self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=lr_critic)
        self.gamma = gamma
        self.lamda = lamda
        self.epsilon = epsilon
        #self.bounds_low = torch.tensor(bounds_low).float()
        #self.bounds_high = torch.tensor(bounds_high).float()
        self.bounds_low = bounds_low
        self.bounds_high = bounds_high
        self.std = std
        self.epochs = 10  # Define the number of epochs for policy update

    def select_action(self, state, edge_index, edge_attr):
        # Pass state through the GNN model first to get state’s embedding
        state_embedding = self.gnn_model(state, edge_index, edge_attr)
        
        print(f"State Embedding Shape: {state_embedding.shape}")
        
        # Then, pass the state embedding through the actor network to get action mean and std
        mean, std = self.actor(state_embedding)

        # Rearranging based on the specifications
        rearranged_mean = rearrange_action_output(mean)

        # Scale mean based on action bounds defined
        mean = self.bounds_low + (torch.sigmoid(rearranged_mean) * (self.bounds_high - self.bounds_low))
        dist = Normal(mean, std)
        
        # Sample an action from the distribution and calculate its log probability
        action = dist.sample()

        action = torch.clamp(action, self.bounds_low, self.bounds_high)  # Ensure action is within bounds
        action_log_prob = dist.log_prob(action)
        
        return action.detach(), action_log_prob.detach()

# Initializing the GNN Model (EnhancedGNNModelWithSharedParams)
num_node_features =  24# Define based on your graph data
num_edge_features =  11# Define based on your graph data
num_out_features =  19# Define based on the expected output feature size of your GNN

# Initializing the Actor and Critic with the GNN Model
action_dim =  19# Define according to your action space dimensions
actor = Actor(num_out_features, action_dim, std=0.0)
print("actor", actor)
critic = Critic(num_out_features)

# Create the environment
env = CircuitEnvironment(server_address, username, password, bounds_low, bounds_high, target_metrics, netlist_content)
gnn_model = EnhancedGNNModelWithSharedParams(num_node_features, num_edge_features, num_out_features)
actor_output_features = gnn_model.conv2.out_channels
print(f"Initializing Actor with output feature size: {actor_output_features}")
print("GNN Model Conv2 Out Channels:", gnn_model.conv2.out_channels)
print("Action Dimension:", action_dim)
agent = PPOAgent(actor, critic, gnn_model, action_dim, bounds_low, bounds_high, lr_actor, lr_critic, gamma, lambda_, epsilon, std)

# Train agent
train(env, agent, env.num_episodes, env.max_timesteps, env.batch_size, env.epsilon)