import gradio as gr
import os

# PersistDataset -----
import os
import csv
import gradio as gr
from gradio import inputs, outputs
import huggingface_hub
from huggingface_hub import Repository, hf_hub_download, upload_file
from datetime import datetime

# created new dataset as awacke1/MindfulStory.csv
#DATASET_REPO_URL = "https://huggingface.co/datasets/awacke1/MindfulStory.csv"
#DATASET_REPO_ID = "awacke1/MindfulStory.csv"
#DATA_FILENAME = "MindfulStory.csv"
#DATA_FILE = os.path.join("data", DATA_FILENAME)
HF_TOKEN = os.environ.get("HF_TOKEN")

# Download dataset repo using hub download
#try:
#    hf_hub_download(
#        repo_id=DATASET_REPO_ID,
#        filename=DATA_FILENAME,
#        cache_dir=DATA_DIRNAME,
#        force_filename=DATA_FILENAME
#    )
#except:
#    print("file not found")
    
#def AIMemory(title: str, story: str):
#    if title and story:
#        with open(DATA_FILE, "a") as csvfile:
#            writer = csv.DictWriter(csvfile, fieldnames=["title", "story", "time"])
#            writer.writerow({"title": title, "story": story, "time": str(datetime.now())})
        # uncomment line below to begin saving your changes
        #commit_url = repo.push_to_hub()
#    return ""


# Set up cloned dataset from repo for operations
#repo = Repository(
#    local_dir="data", clone_from=DATASET_REPO_URL, use_auth_token=HF_TOKEN
#)

#generator1 = gr.Interface.load("bigscience/bloom", api_key=HF_TOKEN)


generator1 = gr.Interface.load("huggingface/gpt2-large", api_key=HF_TOKEN)
generator2 = gr.Interface.load("huggingface/EleutherAI/gpt-neo-2.7B", api_key=HF_TOKEN)
generator3 = gr.Interface.load("huggingface/EleutherAI/gpt-j-6B", api_key=HF_TOKEN)


def calculator(intro, operator, outro):
    if operator == "add":
        output = generator2(intro) + generator3(outro)
        title = intro + " " + outro
#        saved = AIMemory(title, output)
        return output
    elif operator == "subtract":
        output = generator2(outro) + generator3(intro)
        title = outro + " " + intro
#        saved = AIMemory(title, output)
        output = output.replace(intro, "").replace(outro, "")
        return output
    elif operator == "multiply":
        output = generator1(intro) + generator2(outro) + generator3(intro)
        title = intro + " " + outro + " " + intro
#        saved = AIMemory(title, output)
        return output
    elif operator == "divide":
        output = generator1(outro) + generator2(intro) + generator3(outro)
        title = outro + " " + intro + " " + outro
#        saved = AIMemory(title, output)
        output = output.replace(intro, "").replace(outro, "")
        return output

#with open('Mindfulness.txt', 'r') as file:
#    context = file.read()
#contextBox = gr.Textbox(lines=3, default=context, label="Story starter")

examples = [
    ["Telling my daughter garlic is good for you.  Good immune system and keeps pests away.  Ticks, mosquitos, vampires... men.", "multiply", "Two guys walk into a bar.  The third guy ducks."],
    ["Why was the musician arrested?  He got in treble.", "add", "What did the Spanish musician say when he went fishing?  Castanets!"],
    ["To save the planet with good looks and comedy,", "add", "I was going to organize a space-themed birthday party for my son... ...but I could not planet."],
    ["In America, people rob banks.  In Soviet Russia, Banks rob you!", "multiply", "My roommate just told me he tried to walk to Russia from Alaska.  He pulled up short because he could not get his Bering Strait."],
    ["Why does a chicken coop have two doors?  Because if had four doors it would be a chicken sedan!", "multiply", "Why did the chicken cross the park?  To get to the other slide."],
    ["What type of shoes do ninjas wear?  Sneakers", "add", "Can a ninja bring a ninja star into the airport?  Shuriken."],
    ["To save the planet with good looks and comedy find your", "multiply", "Everybody laughed at me when I told them I was going to be a comedian. I thought well, thats not bad for a start."]
]

demo = gr.Interface(
    calculator,
    [
        "text",
        gr.Radio(["add", "subtract", "multiply", "divide"]),
        "text"
    ],
    "text",
    examples=examples,
    article="Saved story memory dataset: https://huggingface.co/datasets/awacke1/MindfulStory.csv with available models to use from text gen: https://huggingface.co/models?pipeline_tag=text-generation&sort=downloads",
    live=True,
)
demo.launch()