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 "estimation", {}

  #json = estimate_body(image)
  candidate, subset = body_estimation(pil2cv(image))
  print(subset.shape)
  json = "{ \"candidate\": " + candidate_to_json_string(candidate) + ", \"subset\": " + subset_to_json_string(subset) + " }"

  return f"""{image.width}px x {image.height}px, {subset.shape[0]} indivisual(s)""", json

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

with gr.Blocks(css="""button { min-width: 80px; }""") as demo:
  with gr.Row():
    with gr.Column(scale=1):
      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")
        estimationResult = gr.Markdown("""estimation""")
        with gr.Row():
          with gr.Column(min_width=80):
            applySizeBtn = gr.Button(value="Apply size")
          with gr.Column(min_width=80):
            replaceBtn = gr.Button(value="Replace")
          with gr.Column(min_width=80):
            importBtn = gr.Button(value="Import")
      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")
        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, min_width=60):
          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
- "ctrl + z", "shift + ctrl + z" to undo, redo
- "a" add new person
- "q + click" to delete person
- "x + drag" to x-axis pseudo-3D rotation
- "c + drag" to y-axis pseudo-3D rotation

Points to note for pseudo-3D rotation: When performing pseudo-3D rotation on the X and Y axes, the projection is converted to 2D and Z-axis information is lost when the mouse button is released. This means that if you finish dragging while the shape is collapsed, you may not be able to restore it to its original state. In such a case, please use the "undo" function.
""")

  width.change(fn=None, inputs=[width], _js="(w) => { resizeCanvas(w,null); }")
  height.change(fn=None, inputs=[height], _js="(h) => { resizeCanvas(null,h); }")

  source.change(
    fn = image_changed,
    inputs = [source],
    outputs = [estimationResult, json])
  applySizeBtn.click(
    fn = lambda x: (x.width, x.height),
    inputs = [source], 
    outputs = [width, height])
  replaceBtn.click(
    fn = None,
    inputs = [json],
    outputs = [],
    _js="(json) => { initializeEditor(); importPose(json); return []; }")
  importBtn.click(
    fn = None,
    inputs = [json],
    outputs = [],
    _js="(json) => { importPose(json); return []; }")

  saveBtn.click(
    fn = None,
    inputs = [], outputs = [json],
    _js="() => { return [savePose()]; }")
  jsonInputBtn.click(
    fn = lambda x: x,
    inputs = [jsonInput], outputs = [json])
  demo.load(fn=None, inputs=[], outputs=[], _js="() => { initializeEditor(); importPose(); return []; }")
    
gr.mount_gradio_app(app, demo, path="/")