import gradio as gr
import cameratransform as ct
import json
import itertools
import numpy as np

keys = [
  ('focallength_mm', gr.Slider(0, 10, value=7, step=0.1, label='focallength_mm')),
  ('sensor_width_mm', gr.Slider(0, 10, value=6.7, step=0.1)),
  ('sensor_height_mm', gr.Slider(0, 10, value=5.6, step=0.1)),
  ('image_width_px', gr.Slider(0, 3000, value=640, step=1)),
  ('image_height_px', gr.Slider(0, 3000, value=480, step=1)),
  ('elevation_m', gr.Slider(0, 20, value=1.8, step=0.1)),
  ('pos_x_m', gr.Slider(-20, 20, value=0, step=0.1)),
  ('pos_y_m', gr.Slider(-20, 20, value=0, step=0.1)),
  ('tilt_deg', gr.Slider(-180, 180, value=0, step=0.1)),
  ('heading_deg', gr.Slider(-180, 180, value=0, step=0.1)),
  ('roll_deg', gr.Slider(-180, 180, value=0, step=0.1)),
  ('distortion_k1', gr.Slider(-30, 30, value=0, step=0.1)),
  ('distortion_k2', gr.Slider(-30, 30, value=0, step=0.1)),
  ('distortion_k3', gr.Slider(-30, 30, value=0, step=0.1)),
  ('xmin', gr.Slider(-30, 30, value=-10, step=1)),
  ('xmax', gr.Slider(-30, 30, value=10, step=1)),
  ('xtickcount', gr.Slider(0, 50, value=31, step=1)),
  ('ymin', gr.Slider(-30, 30, value=0, step=1)),
  ('ymax', gr.Slider(-30, 30, value=20, step=1)),
  ('ytickcount', gr.Slider(0, 50, value=31, step=1)),
  ('detections', gr.Textbox())
]


def main(form):
  projparams = {}
  spatparams = {}
  distparams = {}
  ptsparams = {}

  result = {}
  result['success']=True
  result['errormsgs']=[]

  def check_key(key, destdict, caster=float):
      if key not in form:
          result['success'] = False
          result['errormsgs'].append('{} is missing'.format(key))
      else:
          destdict[key] = caster(form[key])

  check_key("focallength_mm", projparams)
  check_key("sensor_width_mm", projparams)
  check_key("sensor_height_mm", projparams)
  check_key("image_width_px", projparams, int)
  check_key("image_height_px", projparams, int)

  check_key("elevation_m", spatparams)
  check_key("pos_x_m", spatparams)
  check_key("pos_y_m", spatparams)
  check_key("tilt_deg", spatparams)
  check_key("heading_deg", spatparams)
  check_key("roll_deg", spatparams)

  check_key("distortion_k1", distparams)
  check_key("distortion_k2", distparams)
  check_key("distortion_k3", distparams)

  check_key("xmin", ptsparams)
  check_key("xmax", ptsparams)
  check_key("xtickcount", ptsparams, int)
  check_key("ymin", ptsparams)
  check_key("ymax", ptsparams)
  check_key("ytickcount", ptsparams, int)

  result['projection_params']=projparams
  result['spatial_params']=spatparams
  result['distortion_params']=distparams
  result['points_params']=ptsparams
  try:
    result['detections']=json.loads(form['detections'])
  except:
    result['detections']=[]
  result['image_points']=[]

  if result['success']:
      if distparams['distortion_k1'] and \
          (distparams['distortion_k1'] != 0 or distparams['distortion_k2'] != 0 or distparams['distortion_k3'] != 0):
          dis = ct.BrownLensDistortion(distparams['distortion_k1'], distparams['distortion_k2'], distparams['distortion_k3']) 
      else:
          dis = None

      cam = ct.Camera(ct.RectilinearProjection(**projparams),
                      ct.SpatialOrientation(**spatparams),
                      dis)
      fakepts=list(itertools.product(np.linspace(ptsparams['xmin'],ptsparams['xmax'],ptsparams['xtickcount']),
                                     np.linspace(ptsparams['ymin'],ptsparams['ymax'],ptsparams['ytickcount'])))
      for worldpt in fakepts:
          campt=cam.imageFromSpace((worldpt[0], worldpt[1], 0))
          if np.isnan(campt).any(): continue
          campt_uint=campt.astype(np.int32)
          #print("worldpt={} campt={} _uint={}".format(worldpt,campt,campt_uint))
          x, y = (int(campt_uint[0]), int(campt_uint[1]))
          if 0 <= x and x < projparams['image_width_px'] and\
             0 <= y and y < projparams['image_height_px']:
             result['image_points'].append([x,y])
      for det in result['detections']:
          bottomCentre = np.array([det['x'] + det['w']/2, det['y'] + det['h']])
          td = cam.spaceFromImage(bottomCentre)[:2]
          det['td_x']=td[0]
          det['td_y']=td[1]
  return result

def make_fun(parameters):
    exec("def f_make_fun({}): return main(locals())".format(', '.join(parameters)))
    return locals()['f_make_fun']

run = make_fun(list(map(lambda x: x[0], keys)))

inputComponents = list(map(lambda x: x[1], keys))

iface = gr.Interface(fn=run, inputs=inputComponents, outputs='json')
iface.launch()