import torch
import torch.nn as nn
from typing import Callable, Optional

class VeritasSteering:
    """
    The engine that applies Differentiable Logic corrections during inference.
    """
    def __init__(self, 
                 constraint_fn: Callable[[torch.Tensor], torch.Tensor], 
                 steps: int = 10, 
                 lr: float = 0.05,
                 epsilon: float = 1e-4,
                 verbose: bool = False):
        self.constraint_fn = constraint_fn
        self.steps = steps
        self.lr = lr
        self.epsilon = epsilon
        self.verbose = verbose

    def steer(self, module, inputs, output):
        # Enable grad to allow optimization during inference time
        with torch.enable_grad():
            latents = output.clone().detach()
            latents.requires_grad_(True)
            
            # SGD is often more stable for single-point TTO than Adam
            optimizer = torch.optim.SGD([latents], lr=self.lr)
            
            for i in range(self.steps):
                optimizer.zero_grad()
                loss = self.constraint_fn(latents)
                
                if loss.item() < self.epsilon:
                    if self.verbose and i > 0:
                        print(f"[Veritas] Corrected in {i} steps.")
                    break
                
                loss.backward()
                optimizer.step()
            
            if self.verbose:
                final_loss = self.constraint_fn(latents).item()
                if final_loss > self.epsilon:
                    print(f"[Veritas] Warning: Constraint not fully met. Loss: {final_loss:.5f}")

        return latents.detach()

def attach_veritas(model: nn.Module, layer_name: str, constraint_fn: Callable, **kwargs):
    """
    Attaches the logical steering hook to a specific layer in the model.
    """
    target_layer = None
    for name, module in model.named_modules():
        if name == layer_name:
            target_layer = module
            break
            
    if target_layer is None:
        raise ValueError(f"Layer '{layer_name}' not found in model.")
    
    engine = VeritasSteering(constraint_fn, **kwargs)
    target_layer.register_forward_hook(engine.steer)
    print(f"✅ Veritas Logic Layer attached to '{layer_name}'")