import gradio as gr
import torch
import random
import numpy as np
from typing import Dict
import transformers
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from transformers import BertForSequenceClassification,AutoTokenizer, AutoConfig, AutoModel


def seed(language):
    if language == 'python':
        seed = 4
    else:
        seed = 22
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)



languages = ["pharo","python","java"]  # Add more languages as needed
language_choice = gr.Dropdown(languages, label="Select Language")


examples = [
    ["Add offset values to classPool."],
    [ "For example, let's consider a structure that models a fraction, i.e.,"],
    ["The {@link java.text.Collator} class."],
    ["@see  java.lang.Object#toString()"],
    ["Fit the model according to the given training data."],
    ["This is an implementation that uses the result of the previous model to speed up computations along the set of solutions."]
]


"""
examples = [
    ["Add offset values to classPool."],
    ["For example, let's consider a structure that models a fraction, i.e.,"]

]
'2':{"The {@link java.text.Collator} class",
"@see  java.lang.Object#toString()"},
'3':{"Fit the model according to the given training data.",
"This is an implementation that uses the result of the previous model to speed up computations along the set of solutions."}

"""

device = 'cpu'

model_name = 'ZarahShibli/pharo-code-comment-classification3'
MAX_LEN = 200

target_cols_pharo = ['Classreferences', 'Collaborators', 'Example', 'Intent','Keyimplementationpoints', 'Keymessages', 'Responsibilities']
target_cols_python = ['DevelopmentNotes', 'Expand','Parameters', 'Summary', 'Usage']
target_cols_java =[ 'Expand', 'Ownership','Pointer', 'deprecation', 'rational', 'summary', 'usage']




models = {}

class BERTClass(transformers.PreTrainedModel): #torch.nn.Module
    def __init__(self, languages, class_n):
        super(BERTClass, self).__init__(config= AutoConfig.from_pretrained(f'ZarahShibli/{languages}-code-comment-classification3'))
        self.bert = AutoModel.from_pretrained(f'ZarahShibli/{languages}-code-comment-classification3')
        self.fc = torch.nn.Linear(768,class_n) 
        
        
    def forward(self, ids, mask, token_type_ids):
        _, features = self.bert(ids, attention_mask = mask, token_type_ids = token_type_ids, 
                                   return_dict=False)
        output = self.fc(features)  
        return output
    
        





def Multi_Label_Classification_of_Pubmed_Articles(model_input: str, language: str) -> Dict[str, float]: #This wrapper function will pass the article into the model
    
    

    
    class_n = len(eval(f"target_cols_{language}"))
    seed(language)
    model = BERTClass(language, class_n)#.to(device)
    #example_choice = gr.Radio(examples[language], label="Select Example")

    tokenizer = AutoTokenizer.from_pretrained(model_name)


    dict_custom={}
    pred_label = []

    
    print('--------------')
    print(language)
    #target_cols = target_cols_pharo
    if language == 'pharo':
        target_cols = target_cols_pharo
    elif language == 'python':
        target_cols = target_cols_python
    else :
        target_cols = target_cols_java

    # Tokenize the comment sentence using the BERT tokenizer and encode it with special tokens
    inputs = tokenizer.encode_plus(
        model_input,
        truncation=True,
        add_special_tokens=True,
        max_length=MAX_LEN,
        padding='max_length',
        return_token_type_ids=True,
        return_tensors='pt'
    )#.to(device)
    
    # Retrieve the input ids, attention mask, and token type ids from the encoded inputs
    ids = inputs['input_ids']#.to(device, dtype=torch.long)
    mask = inputs['attention_mask']#.to(device, dtype=torch.long)
    token_type_ids = inputs['token_type_ids']#.to(device, dtype=torch.long)

    # Forward pass through model
    print(language)
    print("class_n:", class_n)
    with torch.no_grad():
        outs = model(ids , mask, token_type_ids)
        
    b_logit_pred = outs[0]
    #pred_label = torch.sigmoid(b_logit_pred)
    
    
    pred_label.extend(torch.sigmoid(b_logit_pred).cpu().detach().numpy().tolist())

    print(pred_label)
    # Convert the outputs to boolean values based on the threshold
    outputs_boolean = np.array(pred_label) >= 0.25
    print(outputs_boolean)
    print(np.where(outputs_boolean))
    # Get the indices where outputs are true
    true_indices = np.where(outputs_boolean)[0]

    print("target_cols: ", target_cols)
    # Map the indices to their corresponding categories
    predicted_categories = [target_cols[idx] for idx in true_indices]

    return  {category: pred_label[i] for i, category in zip(true_indices, predicted_categories)}
    

model_input = gr.Textbox("Add offset values to classPool.", show_label=False)

model_output = gr.Label("Predicted Categories", label="Predicted Categories")

title = "Multi Label Code Comment Classification using BERT"
description = "This model uses a fine-tuned BERT model to classify code comments into predefined categories for three programming languages: Python, Pharo, and Java."


#example_choice = gr.Radio(examples['pharo'], label="Select Example")


app = gr.Interface(
    Multi_Label_Classification_of_Pubmed_Articles, 
	inputs=[model_input, language_choice],
    outputs=model_output, 
    #examples=examples,
    title=title,
	description=description,
    allow_flagging='never',
    analytics_enabled=True,
)

app.launch()