import torch
import torch.nn as nn
import os

# ==========================================
# 1. CONFIGURATION (V5.4 Standard)
# ==========================================
AEG_CONFIG = {
    "num_phonemes": 64,
    "num_errors": 9,
    "num_major_types": 3,
    "num_positions": 3,
    "num_structures": 8,
    "max_seq_len": 10,
    "hidden_size": 128
}

class AEGPolicyNetwork_V5_4(nn.Module):
    def __init__(self):
        super(AEGPolicyNetwork_V5_4, self).__init__()
        
        # Embeddings for categorical inputs
        self.emb_error = nn.Embedding(AEG_CONFIG['num_errors'], 8)
        self.emb_major = nn.Embedding(AEG_CONFIG['num_major_types'], 4)
        self.emb_pos = nn.Embedding(AEG_CONFIG['num_positions'], 4)
        
        # Input Dimension Calculation:
        # Cluster(64) + Trap(64) + Feat(8) + Emb_Err(8) + Emb_Maj(4) + Emb_Pos(4) + Flags(6) + Stats(4) = 162
        input_dim = 64 + 64 + 8 + 8 + 4 + 4 + 6 + 4 
        
        # Shared Encoder (Patient Embedding)
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(256, AEG_CONFIG['hidden_size']), nn.ReLU()
        )
        
        # Head A: Volume (Scalar)
        self.head_volume = nn.Linear(AEG_CONFIG['hidden_size'], 1)
        
        # Head B: Difficulty Progression (LSTM Sequence)
        self.lstm_diff = nn.LSTM(input_size=AEG_CONFIG['hidden_size'], hidden_size=64, batch_first=True)
        self.fc_diff = nn.Linear(64, 5) # 5 Difficulty levels
        self.emb_diff_cond = nn.Embedding(5, 8) # Embedding for the predicted difficulty to feed into other heads
        
        # Head D: Format Selector
        self.head_format = nn.Sequential(
            nn.Linear(AEG_CONFIG['hidden_size'] + 8 + 8 + 4, 64),
            nn.ReLU(),
            nn.Linear(64, 8) # 8 Activity Formats
        )

        # Head E: Targets (Residual Anchor)
        self.head_targets = nn.Sequential(
            nn.Linear(64 + 8 + 64, 128),
            nn.ReLU(),
            nn.Linear(128, 64)
        )
        
        # Head F: Structure Scaffolding
        self.head_struct = nn.Sequential(
            nn.Linear(64 + 8 + 8, 64),
            nn.ReLU(),
            nn.Linear(64, AEG_CONFIG['num_structures'])
        )
        
        # Head C: Safety/Forbidden (Residual Bypass)
        self.head_forbidden = nn.Linear(AEG_CONFIG['hidden_size'], 64)
        
        # Head C2: Dynamic Parameters
        self.head_dynamic = nn.Linear(AEG_CONFIG['hidden_size'], 5)

    def forward(self, error_idx, major_idx, pos_idx, cluster, trap, feat, flags, stats):
        # 1. Embed Categoricals
        e_err = self.emb_error(error_idx)
        e_maj = self.emb_major(major_idx)
        e_pos = self.emb_pos(pos_idx)
        
        # 2. Shared Encoding
        x_cat = torch.cat([cluster, trap, feat, e_err, e_maj, e_pos, flags, stats], dim=1)
        h_shared = self.encoder(x_cat)
        
        # 3. Head A: Volume
        pred_vol = self.head_volume(h_shared)
        
        # 4. Head B: Difficulty (LSTM)
        # Repeat shared vector for sequence length
        lstm_input = h_shared.unsqueeze(1).repeat(1, AEG_CONFIG['max_seq_len'], 1)
        lstm_out, _ = self.lstm_diff(lstm_input)
        pred_diff_logits = self.fc_diff(lstm_out)
        
        # Get predicted difficulty indices to condition other heads
        cond_diff = torch.argmax(pred_diff_logits, dim=2)
        cond_vec = self.emb_diff_cond(cond_diff)
        
        # Prepare repeated embeddings for sequence generation
        h_rep = h_shared.unsqueeze(1).repeat(1, AEG_CONFIG['max_seq_len'], 1)
        e_err_rep = e_err.unsqueeze(1).repeat(1, AEG_CONFIG['max_seq_len'], 1)
        e_maj_rep = e_maj.unsqueeze(1).repeat(1, AEG_CONFIG['max_seq_len'], 1)
        cluster_rep = cluster.unsqueeze(1).repeat(1, AEG_CONFIG['max_seq_len'], 1)
        
        # 5. Head D: Format
        fmt_input = torch.cat([h_rep, cond_vec, e_err_rep, e_maj_rep], dim=2)
        pred_fmt = self.head_format(fmt_input)
        
        # 6. Head E: Targets (With Residual Anchor)
        # Output = Net(x) + (Input_Cluster * 5.0)
        target_input = torch.cat([lstm_out, cond_vec, cluster_rep], dim=2)
        pred_targ = self.head_targets(target_input) + (cluster_rep * 5.0)
        
        # 7. Head F: Structure
        struct_input = torch.cat([lstm_out, cond_vec, e_err_rep], dim=2)
        pred_struct = self.head_struct(struct_input)
        
        # 8. Head C: Forbidden (With Residual Bypass)
        # Output = Net(x) + (Input_Trap * 10.0)
        pred_forbid = self.head_forbidden(h_shared) + (trap * 10.0)
        
        # 9. Head C2: Dynamic
        pred_dyn = self.head_dynamic(h_shared)
        
        return pred_vol, pred_diff_logits, pred_fmt, pred_targ, pred_struct, pred_forbid, pred_dyn

# ==========================================
# 2. LOADER UTILITY
# ==========================================
def load_aeg_model(model_path, device):
    """
    Initializes the model and loads weights safely.
    """
    print(f"--- Loading AEG V5.4 from {model_path} ---")
    model = AEGPolicyNetwork_V5_4().to(device)
    
    if os.path.exists(model_path):
        try:
            state_dict = torch.load(model_path, map_location=device)
            model.load_state_dict(state_dict)
            model.eval() # Set to inference mode
            print("✅ AEG Weights Loaded Successfully")
        except Exception as e:
            print(f"❌ CRITICAL: Failed to load AEG weights: {e}")
            raise e
    else:
        print(f"⚠️ WARNING: Model file not found at {model_path}. Running with random weights (Debugging Mode).")
    
    return model