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')

# epoch 75 ist weird


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


while True:
    for img in os.scandir(img_dir):
        start = datetime.now()
        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')
        r = requests.put(f'{API_URL}/ir', files=openBin)
        print('received response')
        r.raise_for_status()
        data = json.loads(json.loads(r.text))

        # FIXME yuck, don't json the json
        pred_disp = np.asarray(data['disp'], dtype='uint8')
        in_img = np.asarray(data['input'], dtype='uint8').transpose((2,0,1))
        ref_pat = np.asarray(data['reference'], dtype='uint8').transpose((2,0,1)).astype('uint8')
        duration = data['duration']
        pred_disp = cv2.transpose(pred_disp)
        print(f'inference took {duration}s')
        print(f'pipeline and transfer took another {(datetime.now() - start).total_seconds() - float(duration)}s\n')

        cv2.imshow('Input Image', in_img)
        cv2.imshow('Reference Image', ref_pat)
        cv2.imshow('Normalized Predicted Disparity', normalize_and_colormap(pred_disp))
        cv2.imshow('Predicted Disparity', pred_disp)
        key = cv2.waitKey()
        if key == 113:
            quit()
        elif key == 101:
            epoch = input('Enter epoch number or "latest"\n')
            r = requests.post(f'{API_URL}/model/update/{epoch}')
            print(r.text)