import torch
import torch.nn as nn
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer

class RouterNetwork(nn.Module):
    def __init__(self, embedding_dim = 384, hidden_dim = 128, num_classes = 3):
        super(RouterNetwork, self).__init__()
        
        # shared representation layer
        self.shared_layer = nn.Sequential(
            nn.Linear(embedding_dim, hidden_dim),
            nn.ReLU(),
            nn.LayerNorm(hidden_dim)
        )
        
        # Actor Head: outputs logits for the 3 modules
        # 0: unstructured, 1: website, 2: structured
        self.actor_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, num_classes)
        )
        
        # Critic Head: outputs a single value for the state value estimation
        self.critic_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, 1)
        )
        
    def forward(self, x):
        shared_features = self.shared_layer(x)
        action_logits = self.actor_head(shared_features)
        state_value = self.critic_head(shared_features)
        return action_logits, state_value
    
    def get_action(self, x):
        action_logits, value = self.forward(x)
        action_probs = F.softmax(action_logits, dim=-1)
        dist = torch.distributions.Categorical(action_probs)
        action = dist.sample()
        return action, dist.log_prob(action), value