# from transformers import TFBertForSequenceClassification, BertTokenizer
# import tensorflow as tf
# import numpy as np
# from sklearn.preprocessing import LabelEncoder

# # --- Configuration (ensure these match your training config) ---
# MODEL_NAME = 'bert-base-uncased'
# MAX_LEN = 128

# # Load saved model and tokenizer
# # !!! UPDATE THESE PATHS TO YOUR SAVED MODEL AND TOKENIZER LOCATION !!!
# model_path = r'C:\Users\adity\OneDrive\Desktop\VS Code Files\Project\ChatUI\bert_sentiment_model'
# tokenizer_path = r'C:\Users\adity\OneDrive\Desktop\VS Code Files\Project\ChatUI\bert_sentiment_tokenizer'

# model = TFBertForSequenceClassification.from_pretrained(model_path)
# tokenizer = BertTokenizer.from_pretrained(tokenizer_path)

# print('Model and Tokenizer loaded successfully.')

# if 'le' not in locals():
#     le = LabelEncoder()
#     # You would need to have access to the original 'labels' list or the 'classes_' array
#     # to fit it correctly or reconstruct it. For this example, I'll use the known classes:
#     original_classes = ['Anxiety', 'Bipolar', 'Depression', 'Normal', 'Personality disorder', 'Stress', 'Suicidal']
#     le.fit(original_classes)

# print('LabelEncoder ready. Classes:', le.classes_)


# # --- Inference helper function ---
# # def predict_text(text):
# #     enc = tokenizer(text, padding=True, truncation=True, max_length=MAX_LEN, return_tensors='tf')
# #     logits = model(enc)[0]
# #     pred = tf.argmax(logits, axis=1).numpy()[0]
# #     proba = tf.nn.softmax(logits, axis=1).numpy()[0]
# #     return {
# #         'label': le.inverse_transform([pred])[0],
# #         'label_id': int(pred),
# #         'probs': float(f"{proba[pred]:.2f}") # Probability of the top predicted label, formatted to 2 decimal places
# #     }
# def predict_text(text):
#     enc = tokenizer(text, padding=True, truncation=True, max_length=MAX_LEN, return_tensors='tf')
#     logits = model(enc)[0]

#     print("DEBUG LOGITS SHAPE:", logits.shape)   # <-- ADD THIS

#     pred = tf.argmax(logits, axis=1).numpy()[0]
#     prob = tf.nn.softmax(logits, axis=1).numpy()[0][pred]
# # user_input_text = input("Please enter the text you want to analyze: ")

# # if user_input_text:
# #     prediction = predict_text(user_input_text)
# #     print(f"\nText: '{user_input_text}'")
# #     print(f"Prediction: {prediction['label']}({prediction['probs']})")
# #     #print(f"Probabilities: {prediction['probs']}")
# # else:
# #     print("No text entered for analysis.")
# # --- Example usage ---
# # print(predict_text("I feel so sad and hopeless"))














import numpy as np
from transformers import TFBertForSequenceClassification, BertTokenizer
import tensorflow as tf
from sklearn.preprocessing import LabelEncoder

# Paths
model_path = r'C:\Users\adity\OneDrive\Desktop\VS Code Files\Project\ChatUI\bert_sentiment_model'
tokenizer_path = r'C:\Users\adity\OneDrive\Desktop\VS Code Files\Project\ChatUI\bert_sentiment_tokenizer'

# Load model & tokenizer
model = TFBertForSequenceClassification.from_pretrained(model_path)
tokenizer = BertTokenizer.from_pretrained(tokenizer_path)

# Label order (Very Important)
classes = ['Anxiety', 'Bipolar', 'Depression', 'Normal',
           'Personality disorder', 'Stress', 'Suicidal']

# Label encoder (consistent with model output)
label_encoder = LabelEncoder()
label_encoder.fit(classes)

print("All model checkpoint layers were used when initializing TFBertForSequenceClassification.")
print(f"All the layers of TFBertForSequenceClassification were initialized from: {model_path}")
print("Model and Tokenizer loaded successfully.")
print("LabelEncoder ready. Classes:", classes)


# --------------------------
# PREDICT FUNCTION
# --------------------------
def predict_text(text: str):
    """Predicts label and confidence using the trained BERT model."""

    # Convert text to input tensors
    inputs = tokenizer(text, return_tensors='tf', padding=True, truncation=True)

    # Run model
    outputs = model(inputs)
    logits = outputs.logits

    # Softmax → probabilities
    probs = tf.nn.softmax(logits, axis=1).numpy()[0]

    # Get predicted class
    pred_index = np.argmax(probs)
    pred_label = classes[pred_index]
    confidence = float(probs[pred_index])

    return pred_label, confidence, probs