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:
  with gr.Row():
    with gr.Column(scale=1):
      startBtn = gr.Button(value="Start edit")
      width = gr.Slider(label="Width", minimum=512, maximum=1024, step=64, value=512, interactive=True)
      height = gr.Slider(label="Height", minimum=512, maximum=1024, step=64, value=512, interactive=True)
      with gr.Accordion(label="Pose estimation", open=False):
        source = gr.Image(type="pil")
        gr.Markdown("""push "Start edit" after estimation done.""")
      with gr.Accordion(label="Json", open=False):
        gr.Markdown("""
| inout            | how to                                                                               |
| -----------------| ----------------------------------------------------------------------------------------- |
| Import | Paste json to "Json source" and click "Import Json", edit the width/height, then click "Start edit". |
| Export | click "Save" and "Copy to clipboard" of "Json" section.                                             |
""")
        json = gr.JSON(label="Json", lines=10)
        jsonInput = gr.Textbox(label="Json source", lines=10)
        jsonInputBtn = gr.Button(value="Import Json")
    with gr.Column(scale=2):
      html = gr.HTML(html_text)
      with gr.Row():
        with gr.Column(scale=1):
          saveBtn = gr.Button(value="Save")
        with gr.Column(scale=7):
          gr.Markdown("""
- "ctrl + drag" to scale
- "alt + drag" to translate
- "shift + drag" to rotate(move right first, release shift, then up or down)
- "space + drag" to move within range
- "[", "]" or mouse wheel to shrink or expand range
""")

  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="/")