import pickle
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.sequence import pad_sequences
import sys

# Constants
MAX_SEQUENCE_LENGTH = 100

def load_resources():
    try:
        # Load the model
        model = load_model('emotion_model.h5')
        
        # Load the tokenizer
        with open('tokenizer.pickle', 'rb') as handle:
            tokenizer = pickle.load(handle)
            
        # Load the label encoder classes
        with open('label_encoder_classes.npy', 'rb') as f:
            classes = np.load(f, allow_pickle=True)
            
        return model, tokenizer, classes
    except Exception as e:
        print(f"Error loading resources: {e}")
        return None, None, None

def predict_emotion(text, model, tokenizer, classes):
    # Tokenize and pad the input text
    sequences = tokenizer.texts_to_sequences([text])
    data = pad_sequences(sequences, maxlen=MAX_SEQUENCE_LENGTH)
    
    # Predict
    prediction = model.predict(data)
    predicted_class_index = np.argmax(prediction)
    predicted_emotion = classes[predicted_class_index]
    confidence = prediction[0][predicted_class_index]
    
    return predicted_emotion, confidence

if __name__ == "__main__":
    model, tokenizer, classes = load_resources()
    
    if model and tokenizer and classes is not None:
        if len(sys.argv) > 1:
            # Command line mode
            text = " ".join(sys.argv[1:])
            emotion, confidence = predict_emotion(text, model, tokenizer, classes)
            print(f"Input: {text}")
            print(f"Predicted Entity: {emotion}")
            print(f"Confidence: {confidence:.2f}")
        else:
            # Interactive mode
            print("Model loaded successfully.")
            print("Type a sentence to analyze its emotion (or 'quit' to exit).")
            
            while True:
                user_input = input("\nEnter text: ")
                if user_input.lower() == 'quit':
                    break
                
                emotion, confidence = predict_emotion(user_input, model, tokenizer, classes)
                print(f"Predicted Entity: {emotion}")
                print(f"Confidence: {confidence:.2f}")