from pathlib import Path
import numpy as np
import random
import re
import textwrap
import torch
from shapely.geometry.polygon import Polygon
import aggdraw
from PIL import Image, ImageDraw, ImageOps, ImageFilter, ImageFont, ImageColor

import gradio as gr

from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM

finetuned = AutoModelForCausalLM.from_pretrained('model')
tokenizer = AutoTokenizer.from_pretrained('gpt2')

device = "cuda:0" if torch.cuda.is_available() else "cpu"
print(device)
finetuned = finetuned.to(device)

# Utility functions

housegan_labels = {"living_room": 1, "kitchen": 2, "bedroom": 3, "bathroom": 4, "missing": 5, "closet": 6, 
                         "balcony": 7, "corridor": 8, "dining_room": 9, "laundry_room": 10}

architext_colors = [[0, 0, 0], [249, 222, 182], [195, 209, 217], [250, 120, 128], [126, 202, 234], [190, 0, 198], [255, 255, 255], 
                   [6, 53, 17], [17, 33, 58], [132, 151, 246], [197, 203, 159], [6, 53, 17],]

regex = re.compile(".*?\((.*?)\)")

def draw_polygons(polygons, colors, im_size=(256, 256), b_color="white", fpath=None):

    image = Image.new("RGB", im_size, color="white")
    draw = aggdraw.Draw(image)

    for poly, color, in zip(polygons, colors):
        #get initial polygon coordinates
        xy = poly.exterior.xy
        coords = np.dstack((xy[1], xy[0])).flatten()
        # draw it on canvas, with the appropriate colors
        brush = aggdraw.Brush((0, 0, 0), opacity=255)
        draw.polygon(coords, brush)
        
        
        #get inner polygon coordinates
        small_poly = poly.buffer(-1, resolution=32, cap_style=2, join_style=2, mitre_limit=5.0)
        if small_poly.geom_type == 'MultiPolygon':
            mycoordslist = [list(x.exterior.coords) for x in small_poly]
            for coord in mycoordslist:
                coords = np.dstack((np.array(coord)[:,1], np.array(coord)[:, 0])).flatten()
                brush2 = aggdraw.Brush((0, 0, 0), opacity=255)
                draw.polygon(coords, brush2)
        elif poly.geom_type == 'Polygon':
            #get inner polygon coordinates
            xy2 = small_poly.exterior.xy
            coords2 = np.dstack((xy2[1], xy2[0])).flatten()
            # draw it on canvas, with the appropriate colors
            brush2 = aggdraw.Brush((color[0], color[1], color[2]), opacity=255)
            draw.polygon(coords2, brush2)

    image = Image.frombytes("RGB", (256,256), draw.tobytes()).transpose(Image.FLIP_TOP_BOTTOM)

    if(fpath):
        image.save(fpath, quality=100, subsampling=0)

    return draw, image

def prompt_to_layout(user_prompt, top_p, top_k, fpath=None):
    
    model_prompt = '[User prompt] {} [Layout]'.format(user_prompt)
    #print(model_prompt)
    input_ids = tokenizer(model_prompt, return_tensors='pt').to(device)
    output = finetuned.generate(**input_ids, do_sample=True, top_p=top_p, top_k=top_k, eos_token_id=50256, max_length=400)
    output = tokenizer.batch_decode(output, skip_special_tokens=True)
    #print(output)
    
    layout = output[0].lstrip().split('[User prompt]')[1].split('[Layout]')[1].split(', ')
    spaces = [txt.split(':')[0] for txt in layout]

    coordinates = [txt.split(':')[1] for txt in layout]
    coordinates = [re.findall(regex, coord) for coord in coordinates]
    
    polygons = []
    for coord in coordinates:
        polygons.append([point.split(',') for point in coord])
        
    geom = []
    for poly in polygons:
        geom.append(Polygon(np.array(poly, dtype=int)))
        
    colors = [architext_colors[housegan_labels[space]] for space in spaces]
    
    _, im = draw_polygons(geom, colors, fpath=fpath)
    
    legend = Image.open("legend.png")
    
    im = np.array(im)
    im[:40, :] = np.array(legend)
    im = Image.fromarray(im)
    
    return im, layout, output
    
def mut_txt2layout(mut_output):
    output = mut_output[0].rstrip().split('[User prompt]')[1].split('[Layout]')[1].split(', ')
    spaces = [txt.split(':')[0].strip(' ') for txt in output]
    coordinates = [txt.split(':')[1] for txt in output]
    coordinates = [re.findall(regex, coord) for coord in coordinates]

    polygons = []
    for coord in coordinates:
        polygons.append([point.split(',') for point in coord])

    geom = []
    for poly in polygons:
        geom.append(Polygon(np.array(poly, dtype=int)))

    colors = [architext_colors[housegan_labels[space]] for space in spaces]
    _, im = draw_polygons(geom, colors, fpath=None)

    legend = Image.open("legend.png")

    im = np.array(im)
    im[:40, :] = np.array(legend)
    im = Image.fromarray(im)
    
    return im
    
def prompt_with_mutation(user_prompt, top_p, top_k, mut_rate, fpath=None):
    
    #Create initial layout based on prompt
    im, layout, output = prompt_to_layout(user_prompt, top_p=top_p, top_k=top_k)
        
    #Create mutated layout based on initial
    mut_len = int((1-mut_rate)*len(layout))
    index1 = random.randrange(0,len(layout)-mut_len)
    rooms = layout[index1:index1+mut_len]
    rooms[-1] = rooms[-1].split(':')[0] + ':'
    rooms = ', '.join(rooms)# + ', '
    new_prompt = '[User prompt] {} [Layout] {}'.format(user_prompt, rooms)
    input_ids = tokenizer(new_prompt, return_tensors='pt')
    mut_output = finetuned.generate(**input_ids, do_sample=True, top_p=top_p, top_k=top_k, eos_token_id=50256, max_length=400)
    mut_output = tokenizer.batch_decode(mut_output, skip_special_tokens=True)
    mut_im = mut_txt2layout(mut_output)
    
    return im, mut_im
    
# Gradio App

custom_css="""
@import url("https://use.typekit.net/nid3pfr.css");
.gradio_page {
  display: flex;
  width: 100vw;
  min-height: 50vh;
  flex-direction: column;
  justify-content: center;
  align-items: center;
  margin: 0px;
  max-width: 100vw;
  background: #FFFFFF;
}
.gradio_interface {
  width: 100vw;
  max-width: 1500px;
}
.gradio_interface[theme=default] .panel_buttons {
  justify-content: flex-end;
}
.gradio_interface[theme=default] .panel_button {
  flex: 0 0 0;
  min-width: 150px;
}
.gradio_interface[theme=default] .panel_button.submit { 
  background: #11213A;
  border-radius: 5px;
  color: #FFFFFF;
  text-transform: uppercase;
  min-width: 150px;
  height: 4em;
  letter-spacing: 0.15em;
  flex: 0 0 0;
}
.gradio_interface[theme=default] .panel_button.submit:hover { 
  background: #000000;
}

.input_text {
  font: 200 50px garamond-premier-pro-display, serif;
  line-height: 115%;
  color: #11213A;
  border-radius: 0px;
  border: 3px solid #11213A;
}
.input_text:focus {
  border-color: #FA7880;
}
.gradio_interface[theme=default] .input_text input, 
.gradio_interface[theme=default] .input_text textarea {
  padding: 30px;
}
.input_text textarea:focus-visible { 
  outline: none;
}
.panel:nth-child(1) {
  margin-left: 50px;
  margin-right: 50px;
  margin-top: 80px;
  margin-bottom: 80px;
  max-width: 750px;
}
.panel:nth-child(2) {
  background: #D3ECF5;
}
.gradio_interface[theme=default] .output_image .image_preview_holder {
  background: #D3ECF5;
}
.gradio_interface[theme=default] .component_set  {
  background: transparent;
  opacity: 1 !important;
}"""

def gen_and_mutate(user_prompt, mutate=False, top_p=0.94, top_k=100, mut_rate=0.2):    
    if(mutate):
        im, mut_im = None, None
        while (mut_im is None):
            try:
                im, mut_im = prompt_with_mutation(user_prompt, top_p, top_k, mut_rate)
            except:
                pass
    else:
        mut_im=Image.open("empty.png")
        im, _, _ = prompt_to_layout(user_prompt, top_p, top_k)
        
    return im, mut_im
    
checkbox =  gr.inputs.Checkbox(label='Mutate')
topp_slider = gr.inputs.Slider(0.1, 1.0, 0.01, default=0.94, label='top_p')
topk_slider = gr.inputs.Slider(0, 100, 25, default=0, label='top_k')
mut_slider = gr.inputs.Slider(0.2, 0.8, 0.1, default=0.3, label='Mutation rate')
textbox = gr.inputs.Textbox(placeholder='house with two bedrooms and one bathroom', lines="2", 
                            label="DESCRIBE YOUR DESIGN")
generated = gr.outputs.Image(label='Generated Layout')
mutated = gr.outputs.Image(label='Mutated Layout')

iface = gr.Interface(fn=gen_and_mutate, inputs=[textbox, checkbox, topp_slider, topk_slider, mut_slider], outputs=[generated, mutated],
                      css=custom_css,
                      thumbnail="thumbnail_gradio.PNG",
                      description='Demo of Semantic Generation of Residential Layouts \n',
                      article='''<div>
    <p> This app allows users the use of natural language prompts for appartment layout generation, using a variety of semantic information:</p>
     <ul>
      <li> <strong>typology</strong>: "a house with two bedrooms and two bathrooms"</li>
      <li> <strong>enumeration</strong>: "a house with five rooms"</li>
      <li> <strong>adjacency</strong>: "the kitchen is adjacent to a bedroom", "the living room is not adjacent to the bathroom"</li>
      <li> <strong>location</strong>: "a house with a bedroom in the north east side"</li>
    </ul>
    <p>You can also create a mutation of the generated layout by enabling the 'Mutate' option.</p>
    <p> Made by: <a href='https://www.linkedin.com/in/theodorosgalanos/'>Theodoros </a> <a href='https://twitter.com/TheodoreGalanos'> Galanos</a> and <a href='https://twitter.com/tylerlastovich'>Tyler Lastovich</a>, using a finetuned <a href='https://huggingface.co/EleutherAI/gpt-neo-125M'> GPT-Neo</a> model. </p>
</div>''')

iface.launch()