"""
Custom modeling file for calibrated sentiment prediction.
Auto-generated - do not edit manually.
"""

import torch
import torch.nn as nn
from transformers import AutoModel, PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
import json
import os
import numpy as np


class CalibratedRegressionModel(PreTrainedModel):
    """
    Sentiment model with built-in calibration.
    
    Usage:
        from transformers import AutoTokenizer
        from modeling_calibrated import CalibratedRegressionModel
        
        model = CalibratedRegressionModel.from_pretrained("your-username/model-name")
        tokenizer = AutoTokenizer.from_pretrained("your-username/model-name")
        
        # Single prediction
        result = model.predict_sentiment("This is great!", tokenizer)
        print(result)  # {'score': 0.85, 'category': 'Very Positive'}
    """
    
    def __init__(self, config):
        super().__init__(config)
        
        # Load base transformer
        self.base_model = AutoModel.from_config(config)
        
        # Regression head
        self.dropout = nn.Dropout(0.1)
        self.regressor = nn.Linear(config.hidden_size, 1)
        
        # Load calibration config
        self.calibrator = None
        self._load_calibrator()
    
    def _load_calibrator(self):
        """Load calibration configuration."""
        calibrator_path = os.path.join(
            os.path.dirname(__file__), 
            "calibrator_config.json"
        )
        
        if not os.path.exists(calibrator_path):
            print("Warning: No calibrator found - using raw predictions")
            return
        
        try:
            with open(calibrator_path, 'r') as f:
                config = json.load(f)
            
            self.calibrator = config
            print(f"Loaded {config['method']} calibrator")
            
        except Exception as e:
            print(f"Warning: Could not load calibrator: {e}")
            self.calibrator = None
    
    def _calibrate_score(self, score):
        """Apply calibration to a score."""
        if self.calibrator is None:
            return score
        
        method = self.calibrator['method']
        
        if method in ['isotonic', 'quantile_mapping']:
            # Linear interpolation from mapping
            mapping = self.calibrator['mapping']
            x = np.array(mapping['input_scores'])
            y = np.array(mapping['output_scores'])
            
            # Simple linear interpolation
            calibrated = np.interp(score, x, y)
            
        elif method == 'piecewise':
            # Apply correction from anchors
            anchors = self.calibrator['anchors']
            anchor_points = sorted([float(k) for k in anchors.keys()])
            anchor_corrections = [anchors[str(p)] for p in anchor_points]
            
            correction = np.interp(score, anchor_points, anchor_corrections)
            calibrated = score + correction
        
        else:
            calibrated = score
        
        return float(np.clip(calibrated, 0.0, 1.0))
    
    def forward(self, input_ids, attention_mask=None, token_type_ids=None, labels=None):
        """Forward pass with automatic calibration."""
        
        # Get base model outputs
        outputs = self.base_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids
        )
        
        pooled_output = outputs.pooler_output
        pooled_output = self.dropout(pooled_output)
        logits = self.regressor(pooled_output).squeeze(-1)
        
        # Clip to valid range
        logits = torch.clamp(logits, 0.0, 1.0)
        
        # Apply calibration during inference (not training)
        if not self.training and self.calibrator is not None:
            # Calibrate each score in the batch
            scores = logits.detach().cpu().numpy()
            calibrated_scores = np.array([self._calibrate_score(s) for s in scores])
            logits = torch.tensor(calibrated_scores, device=logits.device, dtype=logits.dtype)
        
        # Calculate loss if labels provided
        loss = None
        if labels is not None:
            loss_fn = nn.MSELoss()
            loss = loss_fn(logits, labels)
        
        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states if hasattr(outputs, 'hidden_states') else None,
            attentions=outputs.attentions if hasattr(outputs, 'attentions') else None,
        )
    
    @staticmethod
    def score_to_category(score):
        """Convert continuous score to category label."""
        if score <= 0.20:
            return "Very Negative"
        elif score <= 0.40:
            return "Negative"
        elif score <= 0.60:
            return "Neutral"
        elif score <= 0.80:
            return "Positive"
        else:
            return "Very Positive"
    
    def predict_sentiment(self, text, tokenizer, device=None):
        """
        Predict sentiment for a single text (convenience method).
        
        Args:
            text: Input text string
            tokenizer: Loaded tokenizer
            device: Device to use (auto-detected if None)
        
        Returns:
            dict: {'score': float, 'category': str}
        """
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        
        self.eval()
        self.to(device)
        
        # Tokenize
        inputs = tokenizer(
            text,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512
        )
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        # Predict
        with torch.no_grad():
            outputs = self(**inputs)
            score = outputs.logits.item()
        
        return {
            'score': score,
            'category': self.score_to_category(score)
        }