import os
import tqdm
import cv2
import tensorflow as tf
import scipy.interpolate as spi
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex

os.environ["CUDA_VISIBLE_DEVICES"] = "0"


def fill(img, mask):
    coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
    img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
    img_data = img[img_idx[:,0], img_idx[:,1], :]

    interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
    img_2 = interpolator(coords)
    res = np.zeros((512,256,3), dtype=np.float32)
    res[coords[:,0], coords[:,1], :] = img_2
    res[img_idx[:,0], img_idx[:,1], :] = img_data

    return res


def fill_label(img, mask):
    coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
    img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
    img_data = img[img_idx[:,0], img_idx[:,1]]

    interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
    img_2 = interpolator(coords)
    res = np.zeros((512,256), dtype=np.float32)
    res[coords[:,0], coords[:,1]] = img_2
    res[img_idx[:,0], img_idx[:,1]] = img_data

    return res


def downscale_uv_maps_posed_body(folder, type, size: int):
    os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)

    frames = os.listdir(os.path.join(folder, type))

    for frame in tqdm.tqdm(frames):
        id, _ = frame.split('_')
        img = np.load(os.path.join(folder, type + '_filled', frame)); img = img['uv_posed_body_no_orient_filled']
        mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
        down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
        down = down * mask[:,:,np.newaxis]
        np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_body_no_orient = down)


def downscale_uv_maps_curv(folder, type, size):
    os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)

    frames = os.listdir(os.path.join(folder, type))

    for frame in tqdm.tqdm(frames):
        id, _ = frame.split('_')
        img = np.load(os.path.join(folder, type + '_filled', frame)); img = img['uv_curvature_filled']
        mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
        down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
        down = np.squeeze(down) * mask[:,:]
        np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_curvature = down)



def downscale_uv_maps_unposed(folder, type, size):
    os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)

    frames = os.listdir(os.path.join(folder, type))

    for frame in tqdm.tqdm(frames):
        id, _ = frame.split('_')
        img = np.load(os.path.join(folder, type + '_filled', frame)); img = img['uv_unposed_no_shape_filled']
        mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
        down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
        down = down * mask[:,:,np.newaxis]
        np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down)


def downscale_uv_maps_static(folder, type, size):
    os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)

    frames = os.listdir(os.path.join(folder, type))

    for frame in tqdm.tqdm(frames):
        id, _ = frame.split('.')
        img = np.load(os.path.join(folder, type + '_filled', frame)); img = img['uv_static_filled']
        mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
        down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
        down = down * mask[:,:,np.newaxis]
        np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_static = down)


def downscale_uv_maps_mask(folder, type, size):
    os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)

    frames = os.listdir(os.path.join(folder, type))

    for frame in tqdm.tqdm(frames):
        id, _ = frame.split('.')
        img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32)
        down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
        down = np.squeeze(down.numpy().astype(bool))
        np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_mask = down)



if __name__ == '__main__':
    size = 32
    # downscale_uv_maps_mask(folder='/home/workspace/local-data/train_100', type='uv_mask', size=size)
    downscale_uv_maps_static(folder='/home/workspace/local-data/train_100', type='uv_static', size=size)
    downscale_uv_maps_unposed(folder='/home/workspace/local-data/train_100', type='uv_unposed_no_shape', size=size)
    downscale_uv_maps_curv(folder='/home/workspace/local-data/train_100', type='uv_curvature', size=size)
    downscale_uv_maps_posed_body(folder='/home/workspace/local-data/train_100', type='uv_posed_body_no_orient', size=size)


    print()