from transformers import AutoTokenizer, AutoModelForCausalLM
import argparse
import random
import re
import time
import torch

newline2 = "\n\n"

latin = []
corporate = []
headings = []

# Load data from files
with open('data/latin.txt') as latin_file:
    for line in latin_file:
        latin.append(line)
with open('data/corporate.txt') as corporate_file:
    for line in corporate_file:
        corporate.append(line)
with open('data/headings.txt') as headings_file:
    for line in headings_file:
        headings.append(line)


def append_to_file(filename, text):
    with open(filename, "a") as f:
        f.write(text)


def capitalize(text):
    return re.sub("(?<=^)\w|(?<=\.\s)\w", lambda x: x.group().upper(), text)


def generate_latin(prompt):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
    gen_tokens = model.generate(
        input_ids,
        max_length=40,
        min_length=10,
        do_sample=True,
        early_stopping=True,
        top_k=50,
        top_p=0.95,
        temperature=0.7,
        repetition_penalty=1.0,
        no_repeat_ngram_size=2,
        length_penalty=0,
        pad_token_id=0,
    )

    gen_text = tokenizer.batch_decode(gen_tokens)

    return gen_text[0]


def generate_corporate(prompt):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
    gen_tokens = model.generate(
        input_ids,
        max_length=50,
        min_length=20,
        do_sample=True,
        early_stopping=True,
        top_k=50,
        top_p=0.90,
        temperature=0.65,
        repetition_penalty=1.0,
        no_repeat_ngram_size=2,
        pad_token_id=0,
    )

    gen_text = tokenizer.batch_decode(gen_tokens)
    return gen_text[0]


def generate_section(heading):
    latin_prompts = random.sample(latin, 5)
    latin_results = map(generate_latin, latin_prompts)

    corporate_prompts = random.sample(corporate, 5)
    corporate_results = map(generate_corporate, corporate_prompts)

    # Combine the results
    joined = " ".join([x for t in zip(latin_results, corporate_results) for x in t])
    joined = capitalize(joined)

    return f"{heading}{newline2}{joined}{newline2}"


def generate_document(model_name):
    print("Generating document...")
    start = time.time()

    # Pick headings (in order)
    heading_count = random.randint(3, 5)
    heading_indices = random.sample(range(len(headings)), heading_count)
    heading_prompts = [headings[i] for i in sorted(heading_indices)]

    # Generate each section
    sections = map(generate_section, map(lambda h: "# " + h, heading_prompts))

    contents = f"{newline2.join(sections)}{newline2}"

    # Write to file with datetime name
    filename = time.strftime("%Y-%m-%d-%H%M%S")
    append_to_file(f"output/{filename}.md", contents)

    end = time.time()
    duration = round(end - start, 2)
    print("Time taken: ", duration, "seconds")

    append_to_file(
        f"output/{filename}.md",
        f"----{newline2}Generated by [lorem-insight](https://github.com/baumandm/lorem-insight) using [{model_name}](https://huggingface.co/{model_name}) (_{duration} seconds_)",
    )


##### Start of main program #####
parser = argparse.ArgumentParser()
parser.add_argument("--model", dest="model", default="EleutherAI/gpt-neo-125M", help="Model to use")
parser.add_argument("--count", dest="count", default=1, type=int, help="Number of documents to generate")

args = parser.parse_args()

# Determine if we're using the GPU
print("Using GPU:", torch.cuda.is_available())

tokenizer = AutoTokenizer.from_pretrained(args.model)
model = AutoModelForCausalLM.from_pretrained(args.model)

# Generate multiple documents
for i in range(args.count):
    generate_document(args.model)

print("Done!")
