from pathlib import Path
from num2words import num2words
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

def containsNumber(value):
    for character in value:
        if character.isdigit():
            return True
    return False
    
def creativity(intensity):
    if(intensity == 'Low'):
        top_p = 0.95
        top_k = 10
    elif(intensity == 'Medium'):
        top_p = 0.9
        top_k = 50
    if(intensity == 'High'):
        top_p = 0.85
        top_k = 100
    return top_p, top_k

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=(512, 512), b_color="white", fpath=None):

    image = Image.new("RGBA", 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("RGBA", im_size, 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, intensity, fpath=None):

    if(containsNumber(user_prompt) == True):
        spaced_prompt = user_prompt.split(' ')
        new_prompt = ' '.join([word if word.isdigit() == False else num2words(int(word)).lower() for word in spaced_prompt])
        model_prompt = '[User prompt] {} [Layout]'.format(new_prompt)
        
    top_p, top_k = creativity(intensity)
    model_prompt = '[User prompt] {} [Layout]'.format(user_prompt)
    input_ids = tokenizer(model_prompt, return_tensors='pt')
    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)
    
    layout = output[0].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:
        scaled_poly = scale(Polygon(np.array(poly, dtype=int)), xfact=2, yfact=2, origin=(0,0))
        geom.append(scaled_poly)
        #geom.append(Polygon(np.array(poly, dtype=int)))
        
    colors = [architext_colors2[housegan_labels[space]] for space in spaces]
    
    _, im = draw_polygons(geom, colors, fpath=fpath)
    
    return 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;
}""" 
creative_slider = gr.inputs.Radio(["Low", "Medium", "High"], default="Medium", label='Creativity')
textbox = gr.inputs.Textbox(placeholder='a house with two bedrooms and one bathroom', lines="2", 
                            label="DESCRIBE YOUR DESIGN")
generated = gr.outputs.Image(label='Generated Layout')

iface = gr.Interface(fn=prompt_to_layout, inputs=[textbox, creative_slider], 
                      outputs=[generated],
                      css=custom_css,
                      allow_flagging=False,
                      allow_screenshot=False,
                      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()