import os
import tqdm
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"] = "1"
import tensorflow as tf
import tensorflow_addons as tfa
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

# @tf.function
def fill_tf(img, mask):
    img = img[tf.newaxis,:,:,:]
    coords_int = tf.transpose(tf.stack(tf.experimental.numpy.nonzero(1-mask)), (1,0))
    coords = tf.cast(coords_int, dtype=tf.float32)
    img_2 = tfa.image.interpolate_bilinear(img, coords[tf.newaxis,:,:])
    # coords_2 = list(zip(tf.ones(coords_int.shape[0]), coords_int[:,0], coords_int[:,1]))
    # coords = tf.transpose(tf.stack([coords_int[:,0], coords_int[:,1]]))
    out = tf.tensor_scatter_nd_update(img[0], coords_int, img_2[0])

    return out



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 fill_uv_maps_posed_body(folder, type):
    os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
    mask = np.load('custom_utils/mask_256x512.npy')

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

    for frame in tqdm.tqdm(frames):
        id, _ = frame.split('_')
        # mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
        img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
        # filled = fill(img, mask)
        img = np.zeros((100, 100, 3), dtype=np.float32)
        mask = np.zeros((100, 100), dtype=np.int32)
        img[25:75, 25:75, :] = 1
        mask[25:75, 25:75] = 1

        filled = fill_tf(img, mask)
        np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_posed_body_no_orient_filled = filled)


def fill_uv_maps_curv(folder, type):
    os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
    mask = np.load('custom_utils/mask_256x512.npy')

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

    for frame in tqdm.tqdm(frames):
        id, _ = frame.split('_')
        mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
        img = np.load(os.path.join(folder, type, frame)); img = img['curvature_uv']
        filled = fill_label(img, mask)
        np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_curvature_filled = filled)



def fill_uv_maps_unposed(folder, type):
    os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
    mask = np.load('custom_utils/mask_256x512.npy')

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

    for frame in tqdm.tqdm(frames):
        id, _ = frame.split('_')
        mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
        img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape']
        filled = fill(img, mask)
        np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_unposed_no_shape_filled = filled)



def fill_uv_maps_static(folder, type):
    os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
    mask = np.load('custom_utils/mask_256x512.npy')

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

    for frame in tqdm.tqdm(frames):
        id, _ = frame.split('.')
        mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
        img = np.load(os.path.join(folder, type, frame)); img = img['uv_static']
        filled = fill(img, mask)
        np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_static_filled = filled)



if __name__ == '__main__':
    # fill_uv_maps_static(folder='/home/workspace/local-data/train_100', type='uv_static')
    # fill_uv_maps_unposed(folder='/home/workspace/local-data/train_100', type='uv_unposed_no_shape')
    # fill_uv_maps_curv(folder='/home/workspace/local-data/train_100', type='uv_curvature')
    fill_uv_maps_posed_body(folder='/home/workspace/local-data/train_100', type='uv_posed_body_no_orient')


    print()