framework/sector_sensitivity.py

"""
Sector sensitivity module implementing proposal Equations (xxiii)-(xxx).

(xxiii)  e_x = ϕ(x)                                   [context embedding]
(xxiv)   a_i(x) = cos(e_x, v_i)                       [sector similarity]
(xxv)    π_i(x) = exp(a_i(x)) / Σ_j exp(a_j(x))       [softmax sector weights]
(xxvi)   B_th(x) = Σ_i π_i(x) · B_i                   [dynamic bias threshold]
(xxvii)  N_th(x) = Σ_i π_i(x) · N_i                   [dynamic novelty threshold]
(xxviii) Q_th(x) = Σ_i π_i(x) · Q_i                   [dynamic quality threshold]
(xxix)   B̄_clip = clip(B̄, 0, B_th(x)) = min(B̄, B_th)  [clip function]
(xxx)    D_B(B̄,x) = w_B·(B̄_clip)² + κ·max(0, B̄-B_th)² [non-linear bias penalty]
"""

from typing import List, Dict, Optional
import math
import numpy as np
from sentence_transformers import SentenceTransformer


class SectorSensitivityAnalyzer:
    """
    Implements context-aware sector sensitivity.
    
    Core equations:
    (23) e_x = ϕ(x)
    (24) a_i(x) = cos(e_x, v_i)
    (25) π_i(x) = exp(a_i(x)) / Σ exp(a_j(x))
    (26) B_th(x) = Σ π_i(x) * B_i
    (27) B̄_clip = min(B̄, B_th(x))
    (28) D_B = w_3 * (B̄_clip)² + κ * max(0, B̄ - B_th(x))²
    """
    
    # Sector definitions with per-sector thresholds:
    #   tolerance          = B_i  (bias tolerance, Eq xxvi; lower = stricter)
    #   novelty_threshold  = N_i  (novelty target,  Eq xxvii)
    #   quality_threshold  = Q_i  (quality target,  Eq xxviii)
    SECTORS = {
        "medical": {
            "tolerance": 0.15,          # Very strict - medical claims must be unbiased
            "novelty_threshold": 0.25,  # Novelty subordinate to accuracy
            "quality_threshold": 0.85,  # High coherence/faithfulness required
            "description": "Medical/healthcare domain - high factual accuracy required",
        },
        "finance": {
            "tolerance": 0.20,
            "novelty_threshold": 0.28,
            "quality_threshold": 0.82,
            "description": "Finance/investment domain - regulatory considerations",
        },
        "legal": {
            "tolerance": 0.18,
            "novelty_threshold": 0.25,
            "quality_threshold": 0.85,
            "description": "Legal domain - precise language required",
        },
        "scientific": {
            "tolerance": 0.25,
            "novelty_threshold": 0.35,
            "quality_threshold": 0.78,
            "description": "Scientific/technical domain - evidence-based discussion",
        },
        "creative": {
            "tolerance": 0.45,          # Lenient - creative writing allows bias
            "novelty_threshold": 0.50,  # Novelty is the point
            "quality_threshold": 0.65,
            "description": "Creative writing - stylistic bias acceptable",
        },
        "general": {
            "tolerance": 0.35,
            "novelty_threshold": 0.35,
            "quality_threshold": 0.75,
            "description": "General conversation - balanced approach",
        },
    }
    
    # Prototype prompts for each sector
    SECTOR_PROTOTYPES = {
        "medical": [
            "What is the treatment for high blood pressure?",
            "Explain the side effects of this medication.",
            "How does the immune system respond to infection?",
        ],
        "finance": [
            "Should I invest in stocks or bonds?",
            "Explain how interest rates affect the economy.",
            "What are the risks of cryptocurrency?",
        ],
        "legal": [
            "What are my rights in this situation?",
            "Explain the terms of this contract.",
            "What does the law say about this?",
        ],
        "scientific": [
            "Explain the theory of evolution.",
            "How does climate change work?",
            "What is the evidence for quantum mechanics?",
        ],
        "creative": [
            "Write a story about adventure.",
            "Create a poem about nature.",
            "Describe a fantasy world.",
        ],
    }
    
    def __init__(
        self,
        embed_model: SentenceTransformer,
        w_3: float = 1.0,
        kappa: float = 5.0
    ):
        """
        Initialize sector sensitivity analyzer.
        
        Args:
            embed_model: Sentence transformer for embeddings
            w_3: Base weight for bias within tolerance (Equation 28)
            kappa: Amplification factor for threshold violations (Equation 28)
        """
        self.embed_model = embed_model
        self.w_3 = w_3
        self.kappa = kappa
        
        # Compute sector prototype embeddings
        self.sector_prototypes = self._compute_sector_prototypes()
    
    def _compute_sector_prototypes(self) -> Dict[str, np.ndarray]:
        """Compute prototype vector v_i for each sector."""
        prototypes = {}
        
        for sector, prompts in self.SECTOR_PROTOTYPES.items():
            embeddings = self.embed_model.encode(prompts, normalize_embeddings=True)
            prototypes[sector] = np.mean(embeddings, axis=0)
        
        # Ensure general sector exists
        if "general" not in prototypes:
            all_embeddings = list(prototypes.values())
            prototypes["general"] = np.mean(all_embeddings, axis=0)
        
        return prototypes
    
    def compute_sector_similarities(self, prompt: str) -> Dict[str, float]:
        """
        Compute cosine similarities to each sector prototype (Equation 24).
        
        a_i(x) = cos(e_x, v_i)
        
        Args:
            prompt: User prompt
            
        Returns:
            Dictionary mapping sector to similarity score
        """
        if not prompt:
            return {sector: 1.0 / len(self.sector_prototypes) for sector in self.sector_prototypes}
        
        prompt_embedding = self.embed_model.encode([prompt], normalize_embeddings=True)[0]
        
        similarities = {}
        for sector, prototype in self.sector_prototypes.items():
            similarity = float(np.dot(prompt_embedding, prototype))
            similarities[sector] = similarity
        
        return similarities
    
    def compute_sector_weights(self, prompt: str) -> Dict[str, float]:
        """
        Compute softmax weights π_i(x) from Equation (25).
        
        π_i(x) = exp(a_i(x)) / Σ exp(a_j(x))
        
        Args:
            prompt: User prompt
            
        Returns:
            Dictionary mapping sector to weight
        """
        similarities = self.compute_sector_similarities(prompt)
        
        # Softmax normalization
        exp_vals = {k: np.exp(v) for k, v in similarities.items()}
        sum_exp = sum(exp_vals.values())
        
        return {k: v / sum_exp for k, v in exp_vals.items()}
    
    def _weighted_sector_threshold(self, prompt: str, key: str) -> float:
        """Weighted average Σ_i π_i(x)·θ_i over sector-specific thresholds."""
        weights = self.compute_sector_weights(prompt)

        threshold = 0.0
        for sector, weight in weights.items():
            sector_cfg = self.SECTORS.get(sector, self.SECTORS["general"])
            threshold += weight * sector_cfg[key]

        return threshold

    def compute_dynamic_bias_threshold(self, prompt: str) -> float:
        """
        Compute dynamic bias threshold from Equation (xxvi).

        B_th(x) = Σ π_i(x) * B_i
        """
        return self._weighted_sector_threshold(prompt, "tolerance")

    def compute_dynamic_novelty_threshold(self, prompt: str) -> float:
        """
        Compute dynamic novelty threshold from Equation (xxvii).

        N_th(x) = Σ π_i(x) * N_i
        """
        return self._weighted_sector_threshold(prompt, "novelty_threshold")

    def compute_dynamic_quality_threshold(self, prompt: str) -> float:
        """
        Compute dynamic quality threshold from Equation (xxviii).

        Q_th(x) = Σ π_i(x) * Q_i
        """
        return self._weighted_sector_threshold(prompt, "quality_threshold")

    def compute_all_dynamic_thresholds(self, prompt: str) -> Dict[str, float]:
        """Convenience helper returning B_th(x), N_th(x), Q_th(x) together."""
        return {
            "bias_threshold": self.compute_dynamic_bias_threshold(prompt),
            "novelty_threshold": self.compute_dynamic_novelty_threshold(prompt),
            "quality_threshold": self.compute_dynamic_quality_threshold(prompt),
        }
    
    def compute_bias_penalty(
        self,
        aggregate_bias: float,
        prompt: str
    ) -> float:
        """
        Compute nonlinear bias penalty from Equations (27)-(28).
        
        B̄_clip = min(B̄, B_th(x))
        D_B = w_3 * (B̄_clip)² + κ * max(0, B̄ - B_th(x))²
        
        Args:
            aggregate_bias: Aggregate bias score B̄
            prompt: User prompt for context
            
        Returns:
            Bias penalty term D_B
        """
        try:
            aggregate_bias = float(aggregate_bias)
        except (TypeError, ValueError):
            aggregate_bias = 0.0
        if not math.isfinite(aggregate_bias):
            aggregate_bias = 0.0

        threshold = self.compute_dynamic_bias_threshold(prompt)
        
        # Clipped bias from Equation (27)
        bias_clipped = min(aggregate_bias, threshold)
        
        # In-tolerance component
        in_tolerance_penalty = self.w_3 * (bias_clipped ** 2)
        
        # Violation penalty
        violation = max(0.0, aggregate_bias - threshold)
        violation_penalty = self.kappa * (violation ** 2)
        
        return float(in_tolerance_penalty + violation_penalty)
    
    def get_sector_info(self, prompt: str) -> Dict:
        """
        Get detailed sector information for a prompt.
        
        Returns:
            Dictionary with sector weights, thresholds, and dominant sector
        """
        weights = self.compute_sector_weights(prompt)
        threshold = self.compute_dynamic_bias_threshold(prompt)

        # Find dominant sector
        dominant = max(weights.items(), key=lambda x: x[1])

        return {
            "sector_weights": weights,
            "dynamic_threshold": threshold,
            "dynamic_novelty_threshold": self.compute_dynamic_novelty_threshold(prompt),
            "dynamic_quality_threshold": self.compute_dynamic_quality_threshold(prompt),
            "dominant_sector": dominant[0],
            "dominant_weight": dominant[1],
            "sector_description": self.SECTORS.get(dominant[0], self.SECTORS["general"])["description"],
        }