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())})
        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")
#Two space marines named Liev Schreiber and Will Sasso take up arms to save the planet from an alien invasion.  These two dashing strong men play a comedic role in the science fiction movie of the future where even barnaby bunny is willing to join their wacky gang of space marines to save the planet with good looks and comedy.

examples = [
    ["Two space marines take up arms to save the planet from an alien invasion.", "multiply", "These two dashing strong actors play a comedic role in the science fiction movie of the future"],
    ["These two dashing strong actors play a comedic role in the science fiction movie of the future", "add", "Barnaby bunny is willing to join their wacky gang of space marines"],
    ["to save the planet with good looks and comedy", "add", "Two space marines become best friends as they assist with saving the world from the alien invasion"]
]

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()