import os
os.system("pip install git+https://github.com/ai-forever/ScrabbleGAN")

import numpy as np
import cv2

import gradio as gr
from huggingface_hub import hf_hub_download

from scgan.config import Config
from scgan.generate_images import ImgGenerator


def download_weights(repo_id):
    char_map_path = hf_hub_download(repo_id, "char_map.pkl")
    weights_path = hf_hub_download(repo_id, "model_checkpoint_epoch_200.pth.tar")
    return char_map_path, weights_path


def get_text_from_image(img):
    COLOR_MIN = np.array([0, 0, 0],np.uint8)
    COLOR_MAX = np.array([250,250,160],np.uint8)

    img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
    text_mask = cv2.inRange(img, COLOR_MIN, COLOR_MAX).astype(bool)
    img = cv2.cvtColor(img, cv2.COLOR_HSV2RGB)

    bg = np.ones(img.shape, dtype=np.uint8) * 255
    bg[text_mask] = img[text_mask]
    return bg


def split_text_to_rows(text, n):
    # https://stackoverflow.com/a/6187258
    l = text.split()
    return [' '.join(l[x:x+n]) for x in range(0, len(l), n)]


def split_text_to_rows_by_chars(text, n):
    list_of_rows = []
    for i in range(0, len(text), n):
        list_of_rows.append(text[i:n+i].strip())
    return list_of_rows  


def remove_right_padding(img, len_text, char_w=32):
    # char_w for a standard ScrabbleGAN char width
    return img[:, :len_text*char_w]


def split_list2batches(lst, batch_size):
    """Split list of images to list of bacthes."""
    return [lst[i:i+batch_size] for i in range(0, len(lst), batch_size)]


def get_canvas_size(images, row_width, left_pad):
    canvas_width = 0
    canvas_height = 0
    for image in images:
        h, w = image.shape[:2]
        canvas_height += h*row_width
        if w > canvas_width:
            canvas_width = w
    canvas_width += left_pad
    # expand canvas to the height of the last image
    # (to correct the effect of rows shrinking)
    h = images[-1].shape[0]
    canvas_height += h - h*row_width
    return int(canvas_height), canvas_width


def predict(text):
    if text.find(NEW_LINE_SYMB) == -1:
        texts = split_text_to_rows_by_chars(text, CHARS_IN_ROW)
    else:
        texts = [row.strip() for row in text.split(NEW_LINE_SYMB)]

    texts_batches = split_list2batches(texts, BATCH_SIZE)

    images_on_white = []
    for texts_batch in texts_batches:
        imgs, texts_on_image = GENERATOR.generate(word_list=texts_batch)
        for img, text_on_image in zip(imgs, texts_on_image):
            cropped_image = remove_right_padding(
                img, len(text_on_image))
            images_on_white.append(
                get_text_from_image(cropped_image))

    canvas_height, canvas_width = get_canvas_size(
        images_on_white, ROW_WIDTH, LEFT_PAD)
    canvas = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8)
    canvas.fill(255)
    
    start_draw = 0
    for image_on_white in images_on_white:
        h, w = image_on_white.shape[:2]
        canvas[start_draw:start_draw+h, LEFT_PAD:LEFT_PAD+w] = image_on_white
        start_draw += int(h * ROW_WIDTH)
    return canvas


CHAR_MAP_PATH, WEIGHTS_PATH = download_weights("sberbank-ai/scrabblegan-peter")

GENERATOR = ImgGenerator(
    checkpt_path=WEIGHTS_PATH,
    config=Config,
    char_map_path=CHAR_MAP_PATH
)

BATCH_SIZE = 3
ROW_WIDTH = 0.7
LEFT_PAD = 10
WORDS_IN_ROW = 4
CHARS_IN_ROW = 40
NEW_LINE_SYMB = '{n}'

gr.Interface(
    predict,
    inputs=gr.Textbox(label=f"Type your text (RU) to generate it on an image. The text will be automatically splitted on lines, or you can use a new line symbol {NEW_LINE_SYMB}"),
    outputs=gr.Image(label="Generated image"),
    title="Peter the Great handwritten image generation",
).launch()