import json

import gradio as gr
import os

import spacy
spacy.cli.download('en_core_web_sm')
nlp = spacy.load('en_core_web_sm')

import nltk
nltk.download('stopwords')
nltk.download('punkt')

from rake_nltk import Rake
r = Rake()

import time

import wikipediaapi
wiki_wiki = wikipediaapi.Wikipedia('Organika (cmatthewbrown@gmail.com)', 'en')

## ctransformers disabled for now
# from ctransformers import AutoModelForCausalLM
# model = AutoModelForCausalLM.from_pretrained(
#     "Colby/StarCoder-3B-WoW-JSON", 
#     model_file="StarCoder-3B-WoW-JSON-ggml.bin",
#     model_type="gpt_bigcode"
#    )

# Use a pipeline as a high-level helper
from transformers import pipeline
topic_model = pipeline("zero-shot-classification", model="valhalla/distilbart-mnli-12-9", device=0)
#model = pipeline("text-generation", model="Organika/StarCoder-7B-WoW-JSON_1", device=0)

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)

model_name = "umm-maybe/StarCoder-7B-WoW-JSON_1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config, device_map="auto")

def generate_text(prompt):
    inputs = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(
        inputs,
        do_sample=True,
        max_new_tokens=200,
        temperature=0.6,
        top_p=0.9,
        top_k=40,
        repetition_penalty=1.1
        )
    results = tokenizer.decode(outputs[0], clean_up_tokenization_spaces=False)
    return results

def merlin_chat(message, history):
    chat_text = ""
    chat_list = []
    for turn in history[-3:]:
        chat_text += f"{turn[0]}\n\n{turn[1]}\n\n"
        chat_list.append({"role": "user", "content": turn[0]})
        chat_list.append({"role": "assistant", "content": turn[1]})
    chat_text += f"{message}\n"
    doc = nlp(chat_text)
    ents_found = []
    if doc.ents:
        for ent in doc.ents:
            if len(ents_found) == 3:
                break
            if ent.text.isnumeric() or ent.label in ["DATE","TIME","PERCENT","MONEY","QUANTITY","ORDINAL","CARDINAL"]:
                continue
            if ent.text in ents_found:
                continue
            ents_found.append(ent.text.title().lower())
    r.extract_keywords_from_text(chat_text)
    for phrase in r.get_ranked_phrases()[:3]:
        phrase = phrase.lower()
        if phrase not in ents_found:
            ents_found.append(phrase)
    context = ""
    scores = topic_model(chat_text, ents_found, multi_label=True)['scores']
    if ents_found:
        max_score = 0
        for k in range(len(ents_found)):
            if scores[k] < 0.5:
                continue
            entity = ents_found[k]
            if scores[k] > max_score:
                max_score = scores[k]
                max_topic = entity
            print(f'# Looking up {entity} on Wikipedia... ', end='')
            wiki_page = wiki_wiki.page(entity)
            if wiki_page.exists():
                print("page found... ")
                entsum = wiki_page.summary
                if "may refer to" in entsum or "may also refer to" in entsum:
                    print(" ambiguous, skipping.")
                    continue
                else:
                    context += entsum + '\n\n'
            else:
                print("not found.")
    system_msg = {
        'role': 'system', 'content': context
    }
    chat_list.insert(0,system_msg)
    user_msg = {'role': 'user', 'content': message}
    chat_list.append(user_msg)
    prompt = json.dumps(chat_list)[:-1] + ",{\"role\": \"assistant\", \"content\": \""
    print(f"PROMPT: {prompt}")
    for attempt in range(3):
        #result = generate_text(prompt, model_path, parameters, headers)
        #result = model(prompt,return_full_text=False, max_new_tokens=256, temperature=0.8, repetition_penalty=1.1)
        #response = result[0]['generated_text']
        result = generate_text(prompt)
        response = result.replace(prompt,"")
        print(f"COMPLETION: {response}") # so we can see it in logs
        start = 0
        end = 0
        cleanStr = response.lstrip()
        #start = cleanStr.find('{')
        #if start<=0:
        #    continue
        end = cleanStr.find('}') + 1
        if end<=0:
            continue
        cleanStr = cleanStr[:end]
        messageStr = prompt + cleanStr + ']'
        messages = json.loads(messageStr)
        message = messages[-1]
        if message['role'] != 'assistant':
            continue
        msg_text = message['content']
        if chat_text.find(msg_text) >= 0:
            continue
        return message['content']
    return "🤔"

gr.ChatInterface(merlin_chat).launch()