#!/usr/bin/env python3
"""
Inference script for POSNEG sentiment classification model
"""

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import argparse
import json

class POSNEGInference:
    def __init__(self, model_path):
        """
        Initialize the POSNEG sentiment classifier
        
        Args:
            model_path (str): Path to the trained model
        """
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # Load tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
        self.model.to(self.device)
        self.model.eval()
        
        # Label mapping
        self.id2label = {0: "NEGATIVE", 1: "POSITIVE"}
        self.label2id = {"NEGATIVE": 0, "POSITIVE": 1}
    
    def predict(self, text, return_probabilities=False):
        """
        Predict sentiment for given text
        
        Args:
            text (str): Input text to classify
            return_probabilities (bool): Whether to return class probabilities
            
        Returns:
            dict: Prediction results
        """
        # Tokenize input
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding=True,
            max_length=512
        )
        
        # Move to device
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # Get predictions
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            probabilities = torch.nn.functional.softmax(logits, dim=-1)
            predicted_class = torch.argmax(logits, dim=-1).item()
        
        result = {
            'text': text,
            'predicted_label': self.id2label[predicted_class],
            'predicted_class': predicted_class,
            'confidence': probabilities[0][predicted_class].item()
        }
        
        if return_probabilities:
            result['probabilities'] = {
                'NEGATIVE': probabilities[0][0].item(),
                'POSITIVE': probabilities[0][1].item()
            }
        
        return result
    
    def predict_batch(self, texts, return_probabilities=False):
        """
        Predict sentiment for multiple texts
        
        Args:
            texts (list): List of input texts
            return_probabilities (bool): Whether to return class probabilities
            
        Returns:
            list: List of prediction results
        """
        results = []
        for text in texts:
            result = self.predict(text, return_probabilities)
            results.append(result)
        return results

def main():
    parser = argparse.ArgumentParser(description='POSNEG Sentiment Classification Inference')
    parser.add_argument('--model_path', type=str, required=True, help='Path to the trained model')
    parser.add_argument('--text', type=str, help='Single text to classify')
    parser.add_argument('--input_file', type=str, help='JSON file with texts to classify')
    parser.add_argument('--output_file', type=str, help='Output file for results')
    parser.add_argument('--probabilities', action='store_true', help='Include class probabilities')
    
    args = parser.parse_args()
    
    # Initialize classifier
    classifier = POSNEGInference(args.model_path)
    
    if args.text:
        # Single text prediction
        result = classifier.predict(args.text, args.probabilities)
        print(json.dumps(result, indent=2))
    
    elif args.input_file:
        # Batch prediction
        with open(args.input_file, 'r') as f:
            data = json.load(f)
        
        texts = data if isinstance(data, list) else data['texts']
        results = classifier.predict_batch(texts, args.probabilities)
        
        if args.output_file:
            with open(args.output_file, 'w') as f:
                json.dump(results, f, indent=2)
        else:
            print(json.dumps(results, indent=2))
    
    else:
        print("Please provide either --text or --input_file argument")

if __name__ == "__main__":
    main()