import gradio as gr
from llama_cpp import Llama
import os
import json
from gradio_modal import Modal

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
## Imports
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
# Initialize the Llama model

## Download the GGUF model
model_name = "johnpaulbin/qwen2.5-3b-tokipona-v4-merged-fp16-chat-Q8_0-GGUF"
model_file = "qwen2.5-3b-tokipona-v4-merged-fp16-chat-q8_0.gguf"  # This is the specific model file we'll use in this example.
model_path = hf_hub_download(model_name, filename=model_file)

## Instantiate model from downloaded file
llm = Llama(
    model_path=model_path,
    n_ctx=1024,           # Context length to use
    n_threads=2,          # Number of CPU threads to use
    n_gpu_layers=-1,      # Number of model layers to offload to GPU
    chat_format="chatml",
    flash_attn=True
)

# Initialize the conversation with a system prompt
system_prompt = "sina ilo Elelen. sina jan pona pi wile pali. sina pona tawa jan ante li wile pali pona tawa ona."
initial_messages = [{"role": "system", "content": system_prompt}]

# Initialize the dataset
dataset = {
    "prompt": [],
    "completion": [],
    "label": []
}

def generate_response(user_message, chat_history, messages):
    # Append the user's message to the conversation history
    messages = messages + [{"role": "user", "content": user_message}]
    print(messages)
    # Call the model to get assistant's response
    try:
        response = llm.create_chat_completion(
            messages=messages,
            temperature=0.0,
            top_p=1,
            top_k=1,
            frequency_penalty=1
        )
        assistant_response = response['choices'][0]['message']['content']
        # Append the assistant's response to the conversation history
        messages = messages + [{"role": "assistant", "content": assistant_response}]

        # Update the chat history
        chat_history = chat_history + [(user_message, assistant_response)]
        return "", chat_history, messages
    except ValueError as e:
        print(f"Error: You've reached the max context length for the model. Clear the conversation to continue typing.")

        assistant_response = "Error: You've reached the max context length for the model. Clear the conversation to continue typing."
        # Append the assistant's response to the conversation history
        messages = messages + [{"role": "assistant", "content": assistant_response}]

        # Update the chat history
        chat_history = chat_history + [(user_message, assistant_response)]
        return "", chat_history, messages

def upvote_click(chat_history, messages, dataset):
    notification = "Your feedback has been submitted! Feel free to keep chatting or clear the history."
    if len(messages) >= 2:
        # Construct the full conversation history
        full_conversation = ""
        for msg in messages:
            if msg['role'] == 'system':
                full_conversation += f"
system\n{msg['content']}
\n"
            elif msg['role'] == 'user':
                full_conversation += f"
user\n{msg['content']}
\n"
            elif msg['role'] == 'assistant':
                full_conversation += f"
assistant\n{msg['content']}
"

        # Split the conversation into prompt and completion
        prompt_parts = full_conversation.rsplit("
assistant", 1)
        prompt = prompt_parts[0]
        completion = f"
assistant{prompt_parts[1]}" if len(prompt_parts) > 1 else ""

        # Read existing data from feedback.json
        try:
            with open('feedback.json', 'r') as f:
                existing_data = json.load(f)
        except FileNotFoundError:
            existing_data = {'prompt': [], 'completion': [], 'label': []}

        # Check for duplicates
        duplicate = False
        for p, c in zip(existing_data['prompt'], existing_data['completion']):
            if p == prompt and c == completion:
                duplicate = True
                break

        if not duplicate:
            existing_data['prompt'].append(prompt)
            existing_data['completion'].append(completion)
            existing_data['label'].append(True)
            # Save to feedback.json
            with open('feedback.json', 'w') as f:
                json.dump(existing_data, f)
            dataset = existing_data
            print("Upvoted the response! YAAY!!!")
        else:
            notification = "You have already submitted feedback for this conversation."
    else:
        notification = "No conversation to submit feedback for."
    return dataset, notification

def downvote_click(chat_history, messages, dataset):
    notification = "Your feedback has been submitted! Feel free to keep chatting or clear the history."
    if len(messages) >= 2:
        # Construct the full conversation history
        full_conversation = ""
        for msg in messages:
            if msg['role'] == 'system':
                full_conversation += f"
system\n{msg['content']}
\n"
            elif msg['role'] == 'user':
                full_conversation += f"
user\n{msg['content']}
\n"
            elif msg['role'] == 'assistant':
                full_conversation += f"
assistant\n{msg['content']}
"

        # Split the conversation into prompt and completion
        prompt_parts = full_conversation.rsplit("
assistant", 1)
        prompt = prompt_parts[0]
        completion = f"
assistant{prompt_parts[1]}" if len(prompt_parts) > 1 else ""

        # Read existing data from feedback.json
        try:
            with open('feedback.json', 'r') as f:
                existing_data = json.load(f)
        except FileNotFoundError:
            existing_data = {'prompt': [], 'completion': [], 'label': []}

        # Check for duplicates
        duplicate = False
        for p, c in zip(existing_data['prompt'], existing_data['completion']):
            if p == prompt and c == completion:
                duplicate = True
                break

        if not duplicate:
            existing_data['prompt'].append(prompt)
            existing_data['completion'].append(completion)
            existing_data['label'].append(False)
            # Save to feedback.json
            with open('feedback.json', 'w') as f:
                json.dump(existing_data, f)
            dataset = existing_data
            print("Downvoted the response!")
        else:
            notification = "You have already submitted feedback for this conversation."
    else:
        notification = "No conversation to submit feedback for."
    return dataset, notification


def calculate_overall_score(ratings):
    weights = {
        'accuracy': 0.3,
        'fluency': 0.2,
        'helpfulness': 0.3,
        'appropriateness': 0.2
    }
    overall_score = sum(ratings[category] * weight for category, weight in weights.items())
    return overall_score

def submit_feedback(chat_history, messages, dataset, upvote, accuracy, fluency, helpfulness, appropriateness, comment):
    notification = "Your feedback has been submitted! Feel free to keep chatting or clear the history."
    if len(messages) >= 2:
        full_conversation = ""
        for msg in messages:
            if msg['role'] == 'system':
                full_conversation += f"
system\n{msg['content']}
\n"
            elif msg['role'] == 'user':
                full_conversation += f"
user\n{msg['content']}
\n"
            elif msg['role'] == 'assistant':
                full_conversation += f"
assistant\n{msg['content']}
"

        prompt_parts = full_conversation.rsplit("
assistant", 1)
        prompt = prompt_parts[0]
        completion = f"
assistant{prompt_parts[1]}" if len(prompt_parts) > 1 else ""

        try:
            with open('feedback.json', 'r') as f:
                existing_data = json.load(f)
        except FileNotFoundError:
            existing_data = {'prompt': [], 'completion': [], 'label': [], 'ratings': [], 'comments': []}

        duplicate = False
        for p, c in zip(existing_data['prompt'], existing_data['completion']):
            if p == prompt and c == completion:
                duplicate = True
                break

        if not duplicate:
            ratings = {
                'accuracy': accuracy,
                'fluency': fluency,
                'helpfulness': helpfulness,
                'appropriateness': appropriateness
            }
            overall_score = calculate_overall_score(ratings)
            label = overall_score >= 3.5 if upvote else overall_score >= 4.0

            existing_data['prompt'].append(prompt)
            existing_data['completion'].append(completion)
            existing_data['label'].append(label)
            existing_data['ratings'].append(ratings)
            existing_data['comments'].append(comment)

            with open('feedback.json', 'w') as f:
                json.dump(existing_data, f)
            dataset = existing_data
            print(f"Feedback submitted! Overall score: {overall_score}, Label: {label}")
        else:
            notification = "You have already submitted feedback for this conversation."
    else:
        notification = "No conversation to submit feedback for."
    return dataset, notification

def clear_history():
    return [], initial_messages

def retry_response(chat_history, messages):
    notification = ""
    # Check if there is an assistant response to retry
    if len(messages) >= 2 and messages[-1]['role'] == 'assistant':
        # Remove the last assistant message from messages
        messages = messages[:-1]
        # Remove the last exchange from chat_history
        if len(chat_history) >= 1:
            chat_history = chat_history[:-1]
        # Regenerate the assistant's response
    try:
        response = llm.create_chat_completion(
            messages=messages,
            temperature=0.2,
            top_p=0.95,
            top_k=5,
            frequency_penalty=1
        )
        assistant_response = response['choices'][0]['message']['content']
        # Append the assistant's response to the conversation history
        messages = messages + [{"role": "assistant", "content": assistant_response}]
        chat_history = chat_history + [(messages[-2]['content'], assistant_response)]
        return chat_history, messages, notification
    except ValueError as e:
        assistant_response = "Error: You've reached the max context length for the model. Clear the conversation to continue typing."
        # Append the assistant's response to the conversation history
        messages = messages + [{"role": "assistant", "content": assistant_response}]

        # Update the chat history
        chat_history = chat_history + [(user_message, assistant_response)]
        return chat_history, messages, notification
    else:
        notification = "Nothing to retry."
    return chat_history, messages, notification


with gr.Blocks() as demo:
    gr.Markdown("# Toki Pona LLM")
    gr.Markdown("Help improve the model by providing detailed feedback after testing it!")
    messages_state = gr.State(initial_messages)
    dataset_state = gr.State(dataset)
    chat_history = gr.Chatbot()
    with gr.Row():
        user_message = gr.Textbox(placeholder="Enter your message")
        send_button = gr.Button("Send")
    with gr.Row():
        upvote_button = gr.Button("Upvote 👍")
        downvote_button = gr.Button("Downvote 👎")
        retry_button = gr.Button("Retry ♻️")
        clear_button = gr.Button("Clear History 🗑️")
    notification = gr.Markdown("")

    with Modal(visible=False) as feedback_modal:
        upvote = gr.Checkbox(label="Upvote")
        accuracy = gr.Slider(minimum=1, maximum=5, step=1, label="Accuracy")
        fluency = gr.Slider(minimum=1, maximum=5, step=1, label="Fluency")
        helpfulness = gr.Slider(minimum=1, maximum=5, step=1, label="Helpfulness")
        appropriateness = gr.Slider(minimum=1, maximum=5, step=1, label="Appropriateness")
        comment = gr.Textbox(label="Additional Comments", lines=3)
        submit_btn = gr.Button("Submit Feedback")

    send_button.click(
        generate_response,
        inputs=[user_message, chat_history, messages_state],
        outputs=[user_message, chat_history, messages_state]
    )
    user_message.submit(
        generate_response,
        inputs=[user_message, chat_history, messages_state],
        outputs=[user_message, chat_history, messages_state]
    )

    def show_upvote_modal():
        upvote.value = True
        return feedback_modal(visible=True)

    def show_downvote_modal():
        upvote.value = False
        return feedback_modal(visible=True)

    upvote_button.click(show_upvote_modal)
    downvote_button.click(show_downvote_modal)

    submit_btn.click(
        submit_feedback,
        inputs=[chat_history, messages_state, dataset_state, upvote, accuracy, fluency, helpfulness, appropriateness, comment],
        outputs=[dataset_state, notification]
    ).then(lambda: feedback_modal(visible=False))

    retry_button.click(
        retry_response,
        inputs=[chat_history, messages_state],
        outputs=[chat_history, messages_state, notification]
    )

    clear_button.click(
        clear_history,
        outputs=[chat_history, messages_state]
    )

demo.launch(debug=True)


ChatGPT, edit the code above to NOT have any modal functions, instead, replace it with a more gradio native solution.