import requests
from cv2 import cv2
import numpy as np
import json
import os

from datetime import datetime


API_URL = 'http://127.0.0.1:8000'

img_dir = '../../usable_imgs/'

cv2.namedWindow('Input Image')
cv2.namedWindow('Predicted Disparity')


def normalize_and_colormap(img):
    ret = (img - img.min()) / (img.max() - img.min()) * 255.0
    ret = ret.astype("uint8")
    ret = cv2.applyColorMap(ret, cv2.COLORMAP_INFERNO)
    return ret


for img in os.scandir(img_dir):
    if 'ir' not in img.path:
        continue
    input_img = cv2.imread(img.path)
    if input_img.shape == (1024, 1280, 3):
        diff = (512 - 480) // 2
        downsampled = cv2.pyrDown(input_img)
        input_img = downsampled[diff:downsampled.shape[0]-diff, 0:downsampled.shape[1]]

    openBin = {'file': ('file', open(img.path, 'rb'), 'image/png')}

    print('sending image')
    start = datetime.now()
    r = requests.put(f'{API_URL}/ir', files=openBin)
    end = datetime.now()
    print('received response')
    print(f'processing took {end - start}')
    r.raise_for_status()

    # FIXME yuck, don't json the json
    pred_disp = np.asarray(json.loads(json.loads(r.text))['disp'], dtype='uint8')
    ref_pat = np.asarray(json.loads(json.loads(r.text))['reference'], dtype='uint8').transpose((2,0,1)).astype('uint8')
    pred_disp = cv2.transpose(pred_disp)

    cv2.imshow('Input Image', input_img)
    # cv2.imshow('Reference Image', ref_pat)
    cv2.imshow('Predicted Disparity', normalize_and_colormap(pred_disp))
    cv2.waitKey()