import json
import numpy as np
from tokenizers import Tokenizer
import onnxruntime as ort
from huggingface_hub import hf_hub_download
import gradio as gr


class ONNXInferencePipeline:
    def __init__(self, repo_id):
        # Download files from Hugging Face Hub
        self.onnx_path = hf_hub_download(repo_id=repo_id, filename="moodlens.onnx")
        self.tokenizer_path = hf_hub_download(repo_id=repo_id, filename="train_bpe_tokenizer.json")
        self.config_path = hf_hub_download(repo_id=repo_id, filename="hyperparameters.json")

        # Load configuration
        with open(self.config_path) as f:
            self.config = json.load(f)

        # Initialize tokenizer
        self.tokenizer = Tokenizer.from_file(self.tokenizer_path)
        self.max_len = self.config["tokenizer"]["max_len"]

        # Initialize ONNX runtime session
        self.session = ort.InferenceSession(self.onnx_path)
        self.providers = ['CPUExecutionProvider']  # Use CUDA if available
        if 'CUDAExecutionProvider' in ort.get_available_providers():
            self.providers = ['CUDAExecutionProvider']
        self.session.set_providers(self.providers)

    def preprocess(self, text):
        encoding = self.tokenizer.encode(text)
        ids = encoding.ids[:self.max_len]
        padding = [0] * (self.max_len - len(ids))
        return np.array(ids + padding, dtype=np.int64).reshape(1, -1)

    def predict(self, text):
        # Preprocess
        input_array = self.preprocess(text)

        # Run inference
        results = self.session.run(
            None,
            {'input': input_array}
        )

        # Post-process
        logits = results[0]
        probabilities = np.exp(logits) / np.sum(np.exp(logits), axis=1, keepdims=True)
        predicted_class = int(np.argmax(probabilities))

        # Map to labels
        label_mapping = {'neg': 'Negative', 'pos': 'Positive'}
        class_labels = ['neg', 'pos']
        return {
            'label': label_mapping[class_labels[predicted_class]],
            'confidence': float(probabilities[0][predicted_class]),
            'probabilities': probabilities[0].tolist()
        }


# Example usage
if __name__ == "__main__":
    # Initialize the pipeline with the Hugging Face repository ID
    pipeline = ONNXInferencePipeline(repo_id="iimran/MoodlensV3")

    # Example texts for testing
    example_texts = [
        "This product is absolutely amazing! I love how efficient and easy it is to use.",
        "I am very disappointed with this service. The experience was terrible and frustrating."
    ]
    
    for text in example_texts:
        result = pipeline.predict(text)
        print(f"Input: {text}")
        print(f"Prediction: {result['label']} ({result['confidence']:.2%})")
        print(f"Confidence Scores: Negative={result['probabilities'][0]:.2%}, Positive={result['probabilities'][1]:.2%}")
        print("-" * 80)

    # Define a function for Gradio to use
    def gradio_predict(text):
        result = pipeline.predict(text)
        return (
            f"Prediction: {result['label']} ({result['confidence']:.2%})\n"
            f"Confidence Scores: Negative={result['probabilities'][0]:.2%}, Positive={result['probabilities'][1]:.2%}"
        )

    # Create a Gradio interface
    iface = gr.Interface(
        fn=gradio_predict,  # Function to call
        inputs=gr.Textbox(lines=7, placeholder="Enter text here..."),
        outputs="text",     # Output type
        title="MoodLens – Service-Focused Sentiment Analysis Agent",
        description="Moodlens is designed to evaluate the quality of service provided by the council. It looks directly at the content of the communication to identify service issues. For instance, when a resident reports a problem — such as a missing bin— Moodlens interprets that as a clear signal that something is wrong. Enter an email/chat to analyze its sentiment.",
        examples=[
            "The new Customer service portal is fantastic! The new images, categories and website data model is outstanding.",
            "I had a horrible experience with the new council website upgrade. It keeps crashing and the customer support is unhelpful.",
            "The public library had great ambiance but the food was mediocre at best.",
            "I'm extremely dissatisfied with the coumcil services, contact center rep wasn't very friendly and wasn't enough knowledgeable. ",
            "This was the worst customer service experience I've ever had.",
            "Thanks for the prompt response! My complaint has been open for three days, and my SLA is already past due. Really impressive turnaround time, you guys are at least faster than a crippled snail. Thumbs up!",
            "I'm extremely frustrated by the constant loud noises from my neighborhood—it's unbearable and needs urgent attention.",
            "Bravo on the swift response! Three days later and my issue still hasn't been touched it’s amazing how you manage to turn every complaint into a prolonged waiting experience."
        ]
    )

    # Launch the Gradio app
    iface.launch()