#!/usr/bin/env python3
"""
MLP Gate (gϕ) for EvoRM Two-Stage Inference Controller
========================================================
Paper: Section III-E, "MLP gating"
Implements a lightweight MLP gate that decides whether ambiguous
(Survival) pairs require full LLM inference.

Architecture:
  Input: concat[entity_name_sim, rule_type_dist, conf_stats, topo_features] = 64 dims
  Hidden: 256 → 128 → 64
  Output: gating probability ∈ [0,1]

Training: Self-supervised
  - Positive: pairs where Stage 2 LLM returned match
  - Negative: pairs where Stage 2 LLM returned non-match
  - N_warmup = 500 trajectories before training
"""

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from typing import Dict, List, Tuple, Optional
import os
import json


class MLPGate(nn.Module):
    """Lightweight MLP gate for Stage 1→Stage 2 routing."""
    
    def __init__(self, 
                 input_dim: int = 64,
                 hidden_dims: List[int] = None,
                 theta_gate: float = 0.5,
                 n_warmup: int = 500,
                 device: str = 'cuda'):
        super().__init__()
        
        if hidden_dims is None:
            hidden_dims = [256, 128, 64]
        
        self.theta_gate = theta_gate
        self.n_warmup = n_warmup
        self.device = device
        
        # Build MLP layers
        layers = []
        prev_dim = input_dim
        for h_dim in hidden_dims:
            layers.append(nn.Linear(prev_dim, h_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(0.2))
            prev_dim = h_dim
        layers.append(nn.Linear(prev_dim, 1))
        layers.append(nn.Sigmoid())
        
        self.net = nn.Sequential(*layers)
        self.to(device)
        
        # Training state
        self.is_trained = False
        self.samples_features = []  # list of np arrays
        self.samples_labels = []    # list of floats
        self.trajectory_count = 0
        
        # Optimizer (created when training starts)
        self.optimizer = None
        self.criterion = nn.BCELoss()
        
        # Statistics
        self.total_predictions = 0
        self.llm_skipped = 0
    
    def extract_features(self, 
                         es_context: Dict,
                         et_context: Dict,
                         triggered_rules: List) -> np.ndarray:
        """Extract 64-dim feature vector from entity pair context and triggered rules."""
        features = []
        
        # (i) Entity name similarity features (16 dims)
        name1 = str(es_context.get('entity_name', '')).lower()
        name2 = str(et_context.get('entity_name', '')).lower()
        
        tokens1 = set(name1.split())
        tokens2 = set(name2.split())
        jaccard = len(tokens1 & tokens2) / max(1, len(tokens1 | tokens2))
        
        chars1 = set(name1)
        chars2 = set(name2)
        char_jaccard = len(chars1 & chars2) / max(1, len(chars1 | chars2))
        
        len_ratio = min(len(name1), len(name2)) / max(1, max(len(name1), len(name2)))
        len_diff = abs(len(name1) - len(name2)) / max(1, max(len(name1), len(name2)))
        
        name_feats = [jaccard, char_jaccard, len_ratio, len_diff,
                      float(len(name1) > 0), float(len(name2) > 0),
                      float(len(name1) > 10), float(len(name2) > 10)]
        while len(name_feats) < 16:
            name_feats.append(name_feats[len(name_feats) % 8] * 0.5)
        features.extend(name_feats[:16])
        
        # (ii) Rule-type distribution profile (16 dims)
        atom_type_counts = {'SameValue': 0, 'DifferValue': 0,
                           'ShareNeighbor': 0, 'DifferNeighbor': 0,
                           'SemanticEquiv': 0, 'SemanticConflict': 0}
        total_atoms = 0
        for rule in triggered_rules:
            for atom in rule.atoms:
                atype = getattr(atom, 'atom_type', 'Unknown')
                if atype in atom_type_counts:
                    atom_type_counts[atype] += 1
                total_atoms += 1
        
        rule_feats = []
        for atype in sorted(atom_type_counts.keys()):
            rule_feats.append(atom_type_counts[atype] / max(1, total_atoms))
        n_match = sum(1 for r in triggered_rules if r.conclusion == 1)
        n_nonmatch = len(triggered_rules) - n_match
        rule_feats.extend([
            n_match / max(1, len(triggered_rules)),
            n_nonmatch / max(1, len(triggered_rules)),
            float(len(triggered_rules)),
            min(1.0, len(triggered_rules) / 10.0),
        ])
        while len(rule_feats) < 16:
            rule_feats.append(0.0)
        features.extend(rule_feats[:16])
        
        # (iii) Max historical confidence + trigger stats (16 dims)
        conf_feats = []
        if triggered_rules:
            confs = [r.conf for r in triggered_rules]
            triggers = [r.trigger_count for r in triggered_rules]
            conf_feats = [
                max(confs), min(confs),
                sum(confs) / len(confs),
                float(np.std(confs)) if len(confs) > 1 else 0.0,
                max(triggers) / max(1, max(triggers)),
                min(triggers) / max(1, max(triggers)),
                sum(triggers) / max(1, sum(triggers) + len(triggers)),
                float(len(triggered_rules)),
            ]
        while len(conf_feats) < 16:
            conf_feats.append(0.0)
        features.extend(conf_feats[:16])
        
        # (iv) Topological features (16 dims)
        topo_feats = []
        for ctx in [es_context, et_context]:
            n_neighbor_keys = sum(1 for k in ctx if k.startswith('neighbors_'))
            total_n = sum(len(v) if isinstance(v, set) else 1
                         for k, v in ctx.items() if k.startswith('neighbors_'))
            topo_feats.append(float(n_neighbor_keys))
            topo_feats.append(float(total_n) / max(1, total_n))
        
        es_rels = set(k.replace('neighbors_', '') for k in es_context if k.startswith('neighbors_'))
        et_rels = set(k.replace('neighbors_', '') for k in et_context if k.startswith('neighbors_'))
        shared_rels = es_rels & et_rels
        topo_feats.extend([
            float(len(shared_rels)),
            len(shared_rels) / max(1, len(es_rels | et_rels)),
        ])
        while len(topo_feats) < 16:
            topo_feats.append(0.0)
        features.extend(topo_feats[:16])
        
        return np.array(features, dtype=np.float32)
    
    def predict_proba(self, features: np.ndarray) -> float:
        """Predict gating probability."""
        self.eval()
        with torch.no_grad():
            x = torch.from_numpy(features).float().unsqueeze(0).to(self.device)
            prob = self.net(x).item()
        self.total_predictions += 1
        return prob
    
    def should_invoke_llm(self, features: np.ndarray) -> bool:
        """Decide whether to invoke LLM based on gating probability."""
        if not self.is_trained:
            return True  # During warmup, always invoke LLM
        
        prob = self.predict_proba(features)
        if prob < self.theta_gate:
            self.llm_skipped += 1
            return False
        return True
    
    def collect_sample(self, features: np.ndarray, label: float):
        """Collect a training sample during warmup phase."""
        if len(self.samples_features) < self.n_warmup * 2:
            self.samples_features.append(features)
            self.samples_labels.append(label)
        self.trajectory_count += 1
        
        if self.trajectory_count >= self.n_warmup and not self.is_trained:
            self.fit_model()
    
    def fit_model(self, epochs: int = 50, batch_size: int = 32, verbose: bool = True):
        """Self-supervised training on collected samples."""
        if len(self.samples_features) < 10:
            if verbose:
                print(f"[MLPGate] Not enough samples ({len(self.samples_features)})")
            return
        
        # Set module to training mode
        super().train()
        
        X = torch.from_numpy(np.stack(self.samples_features)).float().to(self.device)
        y = torch.tensor(self.samples_labels).float().to(self.device)
        
        n_train = int(0.8 * len(self.samples_features))
        indices = torch.randperm(len(self.samples_features))
        X_train, y_train = X[indices[:n_train]], y[indices[:n_train]]
        X_val, y_val = X[indices[n_train:]], y[indices[n_train:]]
        
        if self.optimizer is None:
            self.optimizer = optim.Adam(self.parameters(), lr=1e-3, weight_decay=1e-5)
        
        for epoch in range(epochs):
            super().train()  # training mode
            total_loss = 0.0
            for i in range(0, len(X_train), batch_size):
                batch_X = X_train[i:i+batch_size]
                batch_y = y_train[i:i+batch_size]
                
                self.optimizer.zero_grad()
                pred = self.net(batch_X).squeeze()
                loss = self.criterion(pred, batch_y)
                loss.backward()
                self.optimizer.step()
                total_loss += loss.item()
            
            self.eval()
            with torch.no_grad():
                val_pred = self.net(X_val).squeeze()
                val_loss = self.criterion(val_pred, y_val).item()
                val_acc = ((val_pred > 0.5) == y_val).float().mean().item()
            
            if verbose and epoch % 10 == 0:
                print(f"[MLPGate] Epoch {epoch}: loss={total_loss/max(1,len(X_train)):.4f}, "
                      f"val_loss={val_loss:.4f}, val_acc={val_acc:.4f}")
        
        self.is_trained = True
        self.eval()
        
        if verbose:
            with torch.no_grad():
                final_pred = self.net(X).squeeze()
                final_acc = ((final_pred > 0.5) == y).float().mean().item()
                n_pos = int((y == 1).sum().item())
                n_neg = int((y == 0).sum().item())
            print(f"[MLPGate] Trained: {len(self.samples_features)} samples, "
                  f"acc={final_acc:.4f}, pos={n_pos}, neg={n_neg}")
    
    def save(self, path: str):
        """Save MLP gate state."""
        state = {
            'model_state': self.state_dict(),
            'optimizer_state': self.optimizer.state_dict() if self.optimizer else None,
            'is_trained': self.is_trained,
            'trajectory_count': self.trajectory_count,
            'n_samples': len(self.samples_features),
            'theta_gate': self.theta_gate,
            'n_warmup': self.n_warmup,
            'total_predictions': self.total_predictions,
            'llm_skipped': self.llm_skipped,
        }
        os.makedirs(os.path.dirname(path) if os.path.dirname(path) else '.', exist_ok=True)
        torch.save(state, path)
        print(f"[MLPGate] Saved to {path}")
    
    def load(self, path: str) -> bool:
        """Load MLP gate state."""
        if not os.path.exists(path):
            return False
        state = torch.load(path, map_location=self.device)
        self.load_state_dict(state['model_state'])
        if state.get('optimizer_state'):
            if self.optimizer is None:
                self.optimizer = optim.Adam(self.parameters(), lr=1e-3)
            self.optimizer.load_state_dict(state['optimizer_state'])
        self.is_trained = state.get('is_trained', False)
        self.trajectory_count = state.get('trajectory_count', 0)
        self.theta_gate = state.get('theta_gate', 0.5)
        self.n_warmup = state.get('n_warmup', 500)
        self.total_predictions = state.get('total_predictions', 0)
        self.llm_skipped = state.get('llm_skipped', 0)
        print(f"[MLPGate] Loaded: trained={self.is_trained}, samples={state.get('n_samples', 0)}")
        return True
    
    def get_stats(self) -> Dict:
        return {
            'is_trained': self.is_trained,
            'trajectory_count': self.trajectory_count,
            'n_samples': len(self.samples_features),
            'total_predictions': self.total_predictions,
            'llm_skipped': self.llm_skipped,
            'skip_rate': self.llm_skipped / max(1, self.total_predictions),
        }


# ==============================================================================
# Test
# ==============================================================================
if __name__ == "__main__":
    print("MLP Gate - Self Test")
    print("=" * 60)
    
    gate = MLPGate(input_dim=64, n_warmup=20, device='cpu')
    print(f"Architecture:\n{gate.net}")
    print(f"Params: {sum(p.numel() for p in gate.parameters())}")
    
    es_ctx = {'entity_name': 'Test Entity A', 'neighbors_rel1': {'B', 'C'}}
    et_ctx = {'entity_name': 'Test Entity B', 'neighbors_rel1': {'C', 'D'}}
    
    class MockAtom:
        def __init__(self, atom_type, attr):
            self.atom_type = atom_type
            self.attr = attr
    class MockRule:
        def __init__(self, conclusion, conf, trigger_count):
            self.conclusion = conclusion
            self.conf = conf
            self.trigger_count = trigger_count
            self.atoms = [MockAtom('SameValue', 'name')]
    
    triggered = [MockRule(1, 0.8, 5), MockRule(0, 0.3, 2)]
    
    features = gate.extract_features(es_ctx, et_ctx, triggered)
    print(f"Features: shape={features.shape}, range=[{features.min():.3f}, {features.max():.3f}]")
    
    prob = gate.predict_proba(features)
    print(f"Prob (before training): {prob:.4f}")
    print(f"Should invoke: {gate.should_invoke_llm(features)}")
    
    print("\nCollecting samples...")
    for i in range(30):
        label = 1.0 if i < 15 else 0.0
        gate.collect_sample(features + np.random.normal(0, 0.1, 64), label)
    
    stats = gate.get_stats()
    print(f"Stats: {stats}")
    
    prob2 = gate.predict_proba(features)
    print(f"Prob (after training): {prob2:.4f}")
    print(f"Should invoke: {gate.should_invoke_llm(features)}")
    
    print("\n✅ MLP Gate test complete!")