import os
# TODO: Is it possible to put these into requirements.txt?
os.system('pip install paddlepaddle')
os.system('pip install paddleocr')
from paddleocr import PaddleOCR, draw_ocr
from PIL import Image
import gradio as gr
import torch
import matplotlib.pylab as plt
import numpy as np

torch.hub.download_url_to_file('https://i.imgur.com/aqMBT0i.jpg', 'example.jpg')

DISPLAY_WIDTH = 40
CYCLES = [20, 60, 100, 140, 180, 220]
PIXEL = '█' # Better than '#' for printing
EMPTY = ' ' # Better than '.' for printing



def solve(file_obj):
    cycle = 0
    sprite = 1 # middle position, we have pixels at 0 and 2
    total = 0
    X = 1
    crt = []
    with open(file_obj.name) as f:
        data = f.read().split("\n")
    for line in data:
        if cycle in [sprite-1, sprite, sprite+1]:
            crt.append(PIXEL)
        else:
            crt.append(EMPTY)
        cycle += 1
        if cycle in CYCLES:
            total += X * cycle
        if line.startswith("addx"):
            if cycle in [sprite-1, sprite, sprite+1]:
                crt.append(PIXEL)
            else:
                crt.append(EMPTY)
            cycle += 1
            if cycle in CYCLES:
                total += X * cycle
            X += int(line.split()[-1])
            sprite = X
        # Go back to first position after 40 pixels
        cycle = cycle % DISPLAY_WIDTH

    # print("Solution to part 1: ", total)
    res = "\n".join(["".join(crt[i:i+DISPLAY_WIDTH]) for i in range(0, len(crt), DISPLAY_WIDTH)])
    print(res)
    
    # Bonus: make an image and then use OCR
    # to extract the text.
    img = np.array([[1 if c == PIXEL else 0 for c in crt[i:i+DISPLAY_WIDTH]] for i in range(0, len(crt), DISPLAY_WIDTH)])
    plt.imshow(img, cmap="binary")
    plt.axis('off')
    plt.savefig("day_10.png")
    img_path = "day_10.png"
    ocr = PaddleOCR(use_angle_cls=True, use_gpu=False)
    result = ocr.ocr(img_path, cls=True)[0]
    image = Image.open(img_path).convert('RGB')
    boxes = [line[0] for line in result]
    txts = [line[1][0] for line in result]
    scores = [line[1][1] for line in result]
    # TODO: Need to upload a font
    # im_show = draw_ocr(image, boxes, txts, scores)
    # im_show = Image.fromarray(im_show)
    # im_show.save('result.jpg')
    return txts, 'day_10.png'

title = 'Cathode-Ray Tube'
description = 'Day 10 2022 AoC using OCR!!!'
article = "<p style='text-align: center'>Day 10 2022 AoC using OCR!!!</p>"
css = ".output_image {height: 40rem !important; width: 100% !important;}"
gr.Interface(
    solve,
    [gr.File(file_types='text', label='Input')],
    [gr.Textbox(label="OCR result"), gr.outputs.Image(type='file', label='OCR image')],
    title=title,
    description=description,
    article=article,
    css=css,
    enable_queue=True
    ).launch(debug=True)