import gradio as gr
from transformers import pipeline

# Set up the generatove model transformer pipeline
generator = pipeline("text-generation", model="matthh/gpt2-poetry-model")

# A sequence of lines both those typed in and the line so far
# when save is clicked the txt file is downloaded
source_lines = []
generated_lines = []


def twolists(list1, list2):
    newlist = []
    a1 = len(list1)
    a2 = len(list2)

    for i in range(max(a1, a2)):
        if i < a1:
            newlist.append(list1[i])
        if i < a2:
            newlist.append(list2[i])

    return newlist


def build_poem(source_lines, generated_lines):
    # This function structures the space already, set boundarys to the possible.
    # return a list with the two lists interspersed
    return twolists(source_lines, generated_lines)


def add_to_lines(new_line):
    source_lines.append(new_line)
    # NOTE: pass the whole array of lines to the generator
    # There is a compounding of the effect, a spiral of interaction.
    poem = build_poem(source_lines, generated_lines)

    # pass the last two lines of the poem
    prompt = "\n".join(poem[-2:])
    result = generator(prompt, max_length=50, num_return_sequences=3)
    response = result[0]["generated_text"]

    # Get only the new generated text
    response = response.replace(prompt, "")

    generated_lines.append(response)

    lines = []
    for line in build_poem(source_lines, generated_lines):
        if (len(lines) % 2) == 1:
            line = f"<p style='text-align: right;'>{line}</p>"
        else:
            line = f"<p>{line}</p>"
        lines.append(line)

    html = "".join(lines)

    return html


def reset_all(new_line):
    global source_lines
    global generated_lines
    source_lines = []
    generated_lines = []

    return None


# demo = gr.Interface(
#     fn=add_to_lines,
#     inputs=gr.Textbox(
#         lines=1,
#         value="The drought had lasted now for ten million years, and the reign of the terrible lizards had long since ended.",
#     ),
#     outputs="html",
#     allow_flagging="never",
# )

demo = gr.Blocks()

with demo:
    with gr.Row():
        with gr.Column():
            input = gr.Textbox(
                lines=1,
                value="The drought had lasted now for ten million years, and the reign of the terrible lizards had long since ended.",
            )
            with gr.Row():
                add_line_button = gr.Button("Add line", variant="primary")
                clear_all_button = gr.Button("Reset all")
        with gr.Column():
            output = gr.HTML()

    add_line_button.click(fn=add_to_lines, inputs=input, outputs=output)
    clear_all_button.click(fn=reset_all, inputs=input, outputs=output)

if __name__ == "__main__":
    demo.launch()