import torch
import torch.nn as nn
from typing import List, Dict, Any

from .config import conf
from .core.tnorms import get_logic_engine
from .core.symbols import Variable
from .core.logic import Rule
from .bridge.reasoner import LogicEvaluator


class SemanticLoss(nn.Module):
    """
    A PyTorch Loss module that calculates how much the data violates the Logic Rules.
    Goal: Minimize (1.0 - TruthValue).
    """
    def __init__(self, concept_map: Dict[str, Any], logic: str = None):
        super().__init__()
        # 1. Setup Logic Engine
        logic_type = logic if logic else conf.logic
        self.tnorm = get_logic_engine(logic_type)
        
        # 2. Setup Reasoner
        self.evaluator = LogicEvaluator(self.tnorm, concept_map)

    def forward(self, rules: List[Rule], data_binding: Dict[Variable, torch.Tensor]) -> torch.Tensor:
        """
        Calculates loss for a list of rules.
        Total Loss = Average(1 - Truth(Rule_i)) for all rules and all batch items.
        """
        total_loss = torch.tensor(0.0, device=next(iter(data_binding.values())).device)
        
        for rule in rules:
            # 1. Evaluate Rule Truth (Vector of size Batch)
            # Result: [0.8, 0.9, 0.2, ...]
            rule_truth = self.evaluator.evaluate(rule, data_binding)
            
            # 2. Calculate Violation (1 - Truth)
            # We want Truth to be 1.0, so Loss is distance from 1.0
            violation = 1.0 - rule_truth
            
            # 3. Aggregate (Mean over batch)
            total_loss += violation.mean()
            
        return total_loss / len(rules)