import gradio as gr
import numpy as np
import cv2
from fastapi import FastAPI, Request, Response
from src.body import Body

body_estimation = Body('model/body_pose_model.pth')

def pil2cv(image):
    ''' PIL型 -> OpenCV型 '''
    new_image = np.array(image, dtype=np.uint8)
    if new_image.ndim == 2:  # モノクロ
        pass
    elif new_image.shape[2] == 3:  # カラー
        new_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)
    elif new_image.shape[2] == 4:  # 透過
        new_image = cv2.cvtColor(new_image, cv2.COLOR_RGBA2BGRA)
    return new_image

with open("static/poseEditor.js", "r") as f:
    file_contents = f.read()

app = FastAPI()

@app.middleware("http")
async def some_fastapi_middleware(request: Request, call_next):
    path = request.scope['path']  # get the request route
    response = await call_next(request)
    
    if path == "/":
        response_body = ""
        async for chunk in response.body_iterator:
            response_body += chunk.decode()

        some_javascript = f"""
        <script type="text/javascript" defer>
{file_contents}
        </script>
        """

        response_body = response_body.replace("</body>", some_javascript + "</body>")

        del response.headers["content-length"]

        return Response(
            content=response_body,
            status_code=response.status_code, 
            headers=dict(response.headers),
            media_type=response.media_type
        )

    return response

# make cndidate to json
def candidate_to_json_string(arr):
    a = [f'[{x:.2f}, {y:.2f}]' for x, y, *_ in arr]
    return '[' + ', '.join(a) + ']'

# make subset to json
def subset_to_json_string(arr):
    arr_str = ','.join(['[' + ','.join([f'{num:.2f}' for num in row]) + ']' for row in arr])
    return '[' + arr_str + ']'

def estimate_body(source):
    if source == None:
      return None

    candidate, subset = body_estimation(pil2cv(source))
    return "{ \"candidate\": " + candidate_to_json_string(candidate) + ", \"subset\": " + subset_to_json_string(subset) + " }"
    
def image_changed(image):
  if (image == None):
    return {}, 512, 512
  json = estimate_body(image)
  return json, image.width, image.height

html_text = f"""
    <canvas id="canvas" width="512" height="512"></canvas>
    <script type="text/javascript" defer>{file_contents}</script>
    """

with gr.Blocks() as demo:
  gr.Markdown("""### Usage

Choose one of the following methods to edit the pose:

| Style            | Description                                                                               |
| -----------------| ----------------------------------------------------------------------------------------- |
| Pose recognition | Upload an image and click "Start edit".                                               |
| Input json       | Input json to "Json source" and click "Input Json", edit the width/height, then click "Start edit".    |
| Free style       | Edit the width/height, then click "Start edit".                                        |

To save the pose image, click "Save".  
To export the pose data, click "Save" and "Copy to clipboard" of "Json" section.
""")
  with gr.Row():
    with gr.Column(scale=1):
      source = gr.Image(type="pil")
      width = gr.Slider(label="Width", mininmum=512, maximum=1920, step=64, value=512, key="Width", interactive=True)
      height = gr.Slider(label="Height", mininmum=512, maximum=1080, step=64, value=512, key="Height", interactive=True)
      startBtn = gr.Button(value="Start edit")
      json = gr.JSON(label="Json", lines=10)
      jsonInput = gr.Textbox(label="Json source", lines=10)
      jsonInputBtn = gr.Button(value="Input Json")
    with gr.Column(scale=2):
      html = gr.HTML(html_text)
      saveBtn = gr.Button(value="Save")
      gr.HTML("<ul><li>ctrl + drag to scale</li><li>alt + drag to translate</li><li>shift + drag to rotate(move right first, then up or down)</li></ul>")

  source.change(
    fn = image_changed,
    inputs = [source],
    outputs = [json, width, height])
  startBtn.click(
    fn = None,
    inputs = [json, width, height], 
    outputs = [],
    _js="(json, w, h) => { initializePose(json,w,h); return []; }")
  saveBtn.click(
    fn = None,
    inputs = [], outputs = [json],
    _js="() => { return [savePose()]; }")
  jsonInputBtn.click(
    fn = lambda x: x,
    inputs = [jsonInput], outputs = [json])
    
gr.mount_gradio_app(app, demo, path="/")