from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
import torch

model = AutoModelForSequenceClassification.from_pretrained("DavinciTech/BERT_Categorizer")
tokenizer = AutoTokenizer.from_pretrained("DavinciTech/BERT_Categorizer")

model.to("cuda")

input_texts = [

    ("Title: Scanner not working\n"
     "Description: Good morning Team My scanner is not connecting to the image saving folder I assume it has "
     "something to do with the merge from last week I need to scan to be able to do payroll which needs to all be done "
     "before and scanning is the first step"),

    ("Title: New mouse please\n"
     "Description: My mouse is acting up a little bit could I get a new one please?"),

    ("Title: Internet outage\n"
     "Description: The whole internet is down for everyone in the office"),
    
]

id2label = {k: l for k, l in enumerate(model.config.LABEL_DICTIONARY.keys())}
#label2id = {l: k for k, l in enumerate(model.config.LABEL_DICTIONARY.keys())}

# Encode the text
encoded = tokenizer(input_texts, truncation=True, padding="max_length", max_length=512, return_tensors="pt").to("cuda")

# Call the model to predict under the format of logits of 27 classes
logits = model(**encoded).logits.cpu().detach().numpy()

IMPACT_LABELS = ["I1", "I2", "I3", "I4"]
IMPACT_INDICES = range(0, 4)
URGENCY_LABELS = ["U1", "U2", "U3", "U4"]
URGENCY_INDICES = range(4, 8)
TYPE_LABELS = ["T1", "T2", "T3", "T4", "T5"]
TYPE_INDICES = range(8, 13)
ALL_LABELS = IMPACT_LABELS + URGENCY_LABELS + TYPE_LABELS

def get_preds_from_logits(logits):
    ret = np.zeros(logits.shape)

    # The first 5 columns (IMPACT_INDICES) are for Impact. They should be handled with a multiclass approach
    # i.e. we fill 1 to the class with highest probability, and 0 into the other columns
    best_class = np.argmax(logits[:, IMPACT_INDICES], axis=-1)
    ret[list(range(len(ret))), best_class] = 1

    ret[:, URGENCY_INDICES] = 0  # Initialize all priority indices to 0
    ret[:, TYPE_INDICES] = 0  # Initialize all type indices to 0

    # Find the index with the maximum value in the PRIORITY_INDICES and set it to 1
    max_priority_index = np.argmax(logits[:, URGENCY_INDICES], axis=-1)
    ret[list(range(len(ret))), max_priority_index + URGENCY_INDICES[0]] = 1

    # Find the index with the maximum value in the TYPE_INDICES and set it to 1
    max_type_index = np.argmax(logits[:, TYPE_INDICES], axis=-1)
    ret[list(range(len(ret))), max_type_index + TYPE_INDICES[0]] = 1

    return ret


# Decode the result
preds = get_preds_from_logits(logits)
decoded_preds = [[id2label[i] for i, l in enumerate(row) if l == 1] for row in preds]

print("\n")

for text, pred in zip(input_texts, decoded_preds):
    print(text)
    print("Impact:", [model.config.LABEL_DICTIONARY[l] for l in pred if l.startswith("I")])
    print("Urgency:", [model.config.LABEL_DICTIONARY[l] for l in pred if l.startswith("U")])
    print("Type:", [model.config.LABEL_DICTIONARY[l] for l in pred if l.startswith("T")])
    print("")