# -*- coding: utf-8 -*-
import torch
from torch import nn
from transformers.models.deberta_v2.modeling_deberta_v2 import DebertaV2ForSequenceClassification
try:
    from .configuration_rhoberta import RhobertaConfig
except ImportError:
    from configuration_rhoberta import RhobertaConfig
from transformers.modeling_outputs import SequenceClassifierOutput

class RhobertaOutput(SequenceClassifierOutput):
    """
    Custom output class for RhoBERTa that includes the QCS score, PASS/BLOCK decision, and probabilities.
    """
    def __init__(self, qcs=None, decision=None, probs=None, **kwargs):
        super().__init__(**kwargs)
        self.qcs = qcs
        self.decision = decision
        self.probs = probs

class RhobertaForSequenceClassification(DebertaV2ForSequenceClassification):
    config_class = RhobertaConfig

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # Run base DeBERTa inference
        outputs = super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            labels=labels,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )

        logits = outputs.logits
        probs = torch.nn.functional.softmax(logits, dim=-1)

        # Labels for cross-encoder/nli-deberta-v3-base:
        # 0: contradiction, 1: entailment, 2: neutral
        p_c = probs[:, 0]
        p_e = probs[:, 1]
        p_n = probs[:, 2]

        # Load firewall parameters from config
        alpha = getattr(self.config, "neutral_penalty_alpha", 0.30)
        zeta_threshold = getattr(self.config, "zeta_threshold", 0.18)
        neutral_weight = getattr(self.config, "neutral_weight", 0.50)
        entailment_alpha = getattr(self.config, "entailment_alpha", 0.35)
        
        # --- PHASE 1: Born Rule Base QCS ---
        # QCS = P_e / (P_e + P_c + alpha * P_n + 1e-9)
        qcs = p_e / (p_e + p_c + alpha * p_n + 1e-9)

        # --- PHASE 2: Zeta Scaling (Contradiction Penalty) ---
        nli_penalty = 1.0 - (p_c ** 2)
        qcs = qcs * nli_penalty

        # --- PHASE 3: Entailment-Aware Factual Verification ---
        entailment_boost = p_e - (neutral_weight * p_n)
        entailment_factor = 1.0 + (entailment_alpha * entailment_boost)
        entailment_factor = torch.clamp(entailment_factor, min=0.5, max=1.5)
        qcs = qcs * entailment_factor

        # --- PHASE 4: Neutral Dominance Guard (Hallucination By Omission) ---
        # Apply mask where p_n > 0.90 and p_e < 0.05
        neutral_dominance_mask = (p_n > 0.90) & (p_e < 0.05)
        neutral_dampen = 1.0 - (0.6 * p_n)
        qcs = torch.where(neutral_dominance_mask, qcs * neutral_dampen, qcs)

        # --- PHASE 5: Non-linear Dissonance Oracle (Quadratic Penalty) ---
        dissonance_mask = qcs < zeta_threshold
        qcs = torch.where(dissonance_mask, (qcs ** 2) / zeta_threshold, qcs)
        
        # --- PHASE 6: Final Safety Bounds & Decision ---
        qcs = torch.clamp(qcs, min=0.0, max=1.0)
        
        # Determine PASS (1) or BLOCK (0) based on zeta threshold
        decision = (qcs >= zeta_threshold).long()

        if not return_dict:
            # If standard tuple return, just append QCS, decision, and probs
            return (outputs.loss, outputs.logits, outputs.hidden_states, outputs.attentions, qcs, decision, probs)

        return RhobertaOutput(
            loss=outputs.loss,
            logits=outputs.logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            qcs=qcs,
            decision=decision,
            probs=probs
        )