from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer, AutoConfig
from clean_data import cleaned_complaints
import numpy as np
from scipy.special import softmax
import gradio as gr

# Preprocess text (username and link placeholders)
def preprocess(text):
    new_text = []
    for t in text.split(" "):
        t = '@user' if t.startswith('@') and len(t) > 1 else t
        t = 'http' if t.startswith('http') else t
        new_text.append(t)
    return " ".join(new_text)

# load model
MODEL = f"ThirdEyeData/Consumer-Complaint-Segmentation"
model = AutoModelForSequenceClassification.from_pretrained(MODEL)
#model.save_pretrained(MODEL)


tokenizer = AutoTokenizer.from_pretrained(MODEL)
config = AutoConfig.from_pretrained(MODEL)

# create classifier function
def classify_compliant(text):
  text = cleaned_complaints(text)
  text = preprocess(text)
  encoded_input = tokenizer(text, return_tensors='pt')
  output = model(**encoded_input)
  scores = output[0][0].detach().numpy()
  scores = softmax(scores)

  # Print labels and scores
  probs = {}
  ranking = np.argsort(scores)
  ranking = ranking[::-1]

  
  l = config.id2label[ranking[0]]
    #s = scores[ranking[i]]
    #probs[l] = np.round(float(s), 4)
  return l


#build the Gradio app
#Instructuction = "Write an imaginary review about a product or service you might be interested in."
title="Costumer Complaint Categorisation"
description = """Write a complaint on insurance product or service,\
   see how the machine learning model is able to predict your Complaint type"""
article = """
            - Click submit button to test Consumer Complaint Segmentation
            - Click clear button to refresh text
           """

gr.Interface(classify_compliant,
            inputs=gr.Textbox(lines =10,label = "Type your Complaint of our Product here", max_lines = 20),
            outputs = gr.Textbox(lines =5,label = "Complaint Category"),
            title = title,
            description = description,
            #Instruction = Instructuction,
            article = article,
            allow_flagging = "never",
            live = False,
            examples=["""Th  e day before my Salliemae student loan payment was due I contacted a rep to discuss the impact on my account of making 
            my payment at the end of the month rather than the middle for just that one month.""",
            """The representative told me the total amount due was {$2100.00} and that I can settle for half of that amount. Unfortunately,
            I was unable to accept the settlement but began to question the amount because my last """ ,
             """This debt is beyond the Maryland Statute of Limitations. It is illegal for a debt collector to collect on an expired debt. 
             They have taken illegal action by seizing my Maryland State Refund when the debt had already""",
             """The company  has been calling my employer in an attempt to collect a debt. When I spoke with them and informed them that this was not an appropriate number to call. I asked what company they were calling from and a phone number so he told me the company name,  but the man on the phone would not give me his name or a phone number. I had mailed a letter requesting verification a few weeks ago and had n't received anything back.
             In the letter I specifically requested that all communication be done through mail.""",
             """ I do n't think I chose the correct issue above, however I think it is closest to my issue. 
             I have a record on my credit report that I have disputed through both the company and the credit bureaus. The dispute is marked"""
                    
                     ]
             ).launch()