# Python std
import argparse
from timeit import default_timer as timer
import math
import random

# project files
import helpers
from R2ATT_UNET import AUNet

# 3rd party
import torch
import numpy as np

import matplotlib.pyplot as plt
import helpers

import polyscope as ps
import polyscope.imgui as psim

def read_file(path, skip_lines=0):
    with open(path, 'r') as file:
        lines = file.readlines()[skip_lines:]
        return lines


def force_cudnn_initialization():
    s = 32
    dev = torch.device('cuda')
    torch.nn.functional.conv2d(torch.zeros(s, s, s, s, device=dev), torch.zeros(s, s, s, s, device=dev))


def normalize(pc, mask):
    pc = pc.to(torch.device('cpu'))
    mask = mask.to(torch.device('cpu'))

    minX = torch.min(pc[0, :, :])
    minY = torch.min(pc[1, :, :])
    minZ = torch.min(pc[2, :, :])

    pc = torch.cat(
        ((pc[0, :, :] - minX).unsqueeze(0), (pc[1, :, :] - minY).unsqueeze(0), (pc[2, :, :] - minZ).unsqueeze(0)), 0)

    mask_x = torch.logical_and(torch.cat((mask[:, 1:], torch.ones((64, 1), dtype=torch.bool)), 1), mask)[:, :-1]
    mask_y = torch.logical_and(torch.cat((mask[1:, :], torch.ones((1, 64), dtype=torch.bool)), 0), mask)[:-1, :]

    diff_x = torch.cat((torch.ones(3, 64, 1), pc), 2) - \
             torch.cat((pc, torch.ones(3, 64, 1)), 2)
    diff_x_loss = diff_x[:, :, 1:-1].norm(p=2, dim=0).mean()
    final_x = mask_x * diff_x[:, :, 1:-1].norm(p=2, dim=0)
    final_x = final_x.sum() / torch.count_nonzero(final_x)

    diff_y = torch.cat((torch.ones(3, 1, 64), pc), 1) - \
             torch.cat((pc, torch.ones(3, 1, 64)), 1)
    diff_y_loss = diff_y[:, 1:-1, :].norm(p=2, dim=0).mean()
    final_y = mask_y * diff_y[:, 1:-1, :].norm(p=2, dim=0)
    final_y = final_y.sum() / torch.count_nonzero(final_y)

    diff = 0.5 * (final_y + final_x)
    pc = pc / (64 * diff)

    print("scale", 64 * diff)
    print("minX", minX)
    print("minY", minY)
    print("minZ", minZ)

    param = {"scale": 64 * diff, "minX": minX, "minY": minY, "minZ": minZ}

    return pc, param


def distanceLoss(a, b):
    return math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2 + (a[2] - b[2]) ** 2)


def computeDistortions(input_norm, mask_param, param):
    target_point = (0.322, 0.457, 0.225)
    cost = float('inf')
    pos = (0, 0)
    values = np.zeros(64 * 64)
    bestMesh = None
    for i in range(12, 50):
        for j in range(12, 50):
            if mask_param[i, j] == 1:
                continue
            mask_param[i, j] = 1
            input_net_in = input_norm * mask_param
            input_net_in[0, i, j] = target_point[0]
            input_net_in[1, i, j] = target_point[1]
            input_net_in[2, i, j] = target_point[2]

            data_in = torch.cat((input_net_in, mask_param.unsqueeze(0)), dim=0).to(torch.device("cuda"))
            data_in = data_in.unsqueeze(0)
            start = timer()
            model(data_in)
            #print(timer() - start)
            loss = distanceLoss(model.im_pred.squeeze(0).view(3, 4096).transpose(0, 1)[i * 64 + j], target_point)
            values[i * 64 + j] = loss
            if cost > loss:
                bestMesh = model.im_pred.squeeze(0).view(3, 4096).transpose(0, 1).detach().clone().cpu().numpy()
                cost = loss
                pos = (i, j)
            #print((i, j), loss)
            mask_param[i, j] = 0
    print("pos : " + str(pos[0]) + ", " + str(pos[1]))
    ref = bestMesh[pos[0] * 64 + pos[1]]
    print("Distance : ", math.sqrt(
        (ref[0] - target_point[0]) ** 2 + (ref[1] - target_point[1]) ** 2 + (ref[2] - target_point[2]) ** 2))
    psbest_cloud = ps.register_point_cloud("BestMatch", bestMesh)
    psbest_cloud.set_enabled(False)
    return values


def set_grab(input_n, mask_tmp):
    indices = [(32, 32)]
    for i, j in indices:
        mask_tmp[i, j] = 1
        input_n[0, i, j] = model.im_pred[0, 0, i, j]
        input_n[1, i, j] = model.im_pred[0, 1, i, j]
        input_n[2, i, j] = model.im_pred[0, 2, i, j]

    return input_n, mask_tmp

def set_pinch(input_n, mask_tmp):
    indices = [(28, 43), (28, 14)]
    for i, j in indices:
        mask_tmp[i, j] = 1
        input_n[0, i, j] = model.im_pred[0, 0, i, j]
        input_n[1, i, j] = model.im_pred[0, 1, i, j]
        input_n[2, i, j] = model.im_pred[0, 2, i, j]

    return input_n, mask_tmp


def set_twist(input_n, mask_tmp):
    for k in range(-e, e + 1, e):
        mask_tmp[32, 32] = 1
        mask_tmp[32, 32 + k] = 1
        mask_tmp[32 + k, 32] = 1
        input_n[0, 32, 32] = model.im_pred[0, 0, 32, 32]
        input_n[1, 32, 32] = model.im_pred[0, 1, 32, 32]
        input_n[2, 32, 32] = model.im_pred[0, 2, 32, 32]
        input_n[0, 32 + k, 32] = model.im_pred[0, 0, 32 + k, 32]
        input_n[1, 32 + k, 32] = model.im_pred[0, 1, 32 + k, 32]
        input_n[2, 32 + k, 32] = model.im_pred[0, 2, 32 + k, 32]
        input_n[0, 32, 32 + k] = model.im_pred[0, 0, 32, 32 + k]
        input_n[1, 32, 32 + k] = model.im_pred[0, 1, 32, 32 + k]
        input_n[2, 32, 32 + k] = model.im_pred[0, 2, 32, 32 + k]
    return input_n, mask_tmp


def grab(input_n, xyz):
    input_n[0, 32, 32] = x
    input_n[1, 32, 32] = y
    input_n[2, 32, 32] = z
    
    target_cloud.update_point_positions(np.array([[x, y, z]]))
    return input_n

def twist(input_n, xyz, thet):
    thet += 3.141592
    x, y, z = xyz
    for k in range(-e, e+1, e):
        input_n[0, 32, 32] = x
        input_n[1, 32, 32] = y
        input_n[2, 32, 32] = z
        input_n[0, 32 + k, 32] = x + k * 0.0329 * math.cos(thet)
        input_n[1, 32 + k, 32] = y
        input_n[2, 32 + k, 32] = z + k * 0.0329 * math.sin(thet)
        input_n[0, 32, 32 + k] = x + k * 0.0329 * math.sin(thet)
        input_n[1, 32, 32 + k] = y
        input_n[2, 32, 32 + k] = z - k * 0.0329 * math.cos(thet)

    target_cloud.update_point_positions(np.array([[x, y, z],
                                                  [x + (-e) * 0.0329 * math.cos(thet), y, z + (-e) * 0.0329 * math.sin(thet)],
                                                  [x + e * 0.0329 * math.cos(thet), y, z + e * 0.0329 * math.sin(thet)],
                                                  [x + (-e) * 0.0329 * math.sin(thet), y, z - (-e) * 0.0329 * math.cos(thet)],
                                                  [x + e * 0.0329 * math.sin(thet), y, z - e * 0.0329 * math.cos(thet)]]))
    return input_n


def pinch(input_n, pinch_pos, dist):
    indices = [(28, 43), (28, 14)]
    mid_point = (0.5*(pinch_pos[0][0] + pinch_pos[1][0]), 0.5*(pinch_pos[0][1] + pinch_pos[1][1]), 0.5*(pinch_pos[0][2] + pinch_pos[1][2]))

    input_n[0, indices[0][0], indices[0][1]] = dist * pinch_pos[0][0] + (1 - dist) * mid_point[0]
    input_n[1, indices[0][0], indices[0][1]] = dist * pinch_pos[0][1] + (1 - dist) * mid_point[1]
    input_n[2, indices[0][0], indices[0][1]] = dist * pinch_pos[0][2] + (1 - dist) * mid_point[2]

    input_n[0, indices[1][0], indices[1][1]] = dist * pinch_pos[1][0] + (1 - dist) * mid_point[0]
    input_n[1, indices[1][0], indices[1][1]] = dist * pinch_pos[1][1] + (1 - dist) * mid_point[1]
    input_n[2, indices[1][0], indices[1][1]] = dist * pinch_pos[1][2] + (1 - dist) * mid_point[2]
    aa = (dist * pinch_pos[0][0] + (1 - dist) * mid_point[0]).detach().clone().cpu()
    ab = (dist * pinch_pos[0][1] + (1 - dist) * mid_point[1]).detach().clone().cpu()
    ac = (dist * pinch_pos[0][2] + (1 - dist) * mid_point[2]).detach().clone().cpu()
    ad = (dist * pinch_pos[1][0] + (1 - dist) * mid_point[0]).detach().clone().cpu()
    ae = (dist * pinch_pos[1][1] + (1 - dist) * mid_point[1]).detach().clone().cpu()
    af = (dist * pinch_pos[1][2] + (1 - dist) * mid_point[2]).detach().clone().cpu()
    target_cloud.update_point_positions(np.array([[aa, ab, ac], [ad, ae, af]]))
    return input_n

def gaussian_2d(tensor_size, center, sigma=1.0):
    x, y = torch.meshgrid(torch.arange(tensor_size[0]), torch.arange(tensor_size[1]))
    x_center, y_center = center
    exponent = -((x - x_center) ** 2 + (y - y_center) ** 2) / (2 * sigma ** 2)
    return torch.exp(exponent)

def minimize(target, uv, input_model, mask_model):
    target_tensor = target.view(-1, 1, 1)
    target_tensor = target_tensor.to('cuda')
    mask_cache = mask_model.detach().clone().to('cuda')
    u, v = uv
    visited = {}
    input_model_gpu = input_model.to('cuda')
    step_size = 25.0
    last_loss = 10000000
    last_grad = (0.0, 0.0)
    last_uv = (u, v)

    while True:
        coordinates = torch.tensor([float(u), float(v)], requires_grad=True)
        soft_mask = 0.999 * gaussian_2d((64, 64), center=coordinates, sigma=0.5)
        soft_mask = soft_mask.to('cuda')

        mask_tmp = mask_cache.clone()
        mask_tmp[u, v] = 1
        data_input = input_model_gpu * (1 - soft_mask).expand(3, 64, 64) + soft_mask * target_tensor
        data_input_model = data_input * mask_tmp
        data_input_model = torch.cat((data_input_model, mask_tmp.unsqueeze(0)), dim=0).to(torch.device("cuda"))
        data_input_model = data_input_model.unsqueeze(0)
        model(data_input_model)

        loss = torch.norm(model.im_pred.squeeze(0) * soft_mask.expand(3, 64, 64) - soft_mask * target_tensor)

        loss.backward(retain_graph=True)
        if last_loss < loss:
            (u, v) = last_uv
            step_size /= 2
            continue
        last_loss = loss

        normalized_grad = coordinates.grad / torch.norm(coordinates.grad)
        if (step_size * torch.norm(coordinates.grad) < 1.0):
            grad_u, grad_v = normalized_grad
        last_uv = (u, v)
        u = int(torch.round(u - step_size * grad_u).item())
        v = int(torch.round(v - step_size * grad_v).item())
        print((u, v), loss, coordinates.grad)
        if not (u, v) in visited:
            visited[(u, v)] = loss
        else:
            break

    minLoss = float('inf')
    bestPos = (0, 0)
    path = np.zeros(64 * 64)
    for coord in visited:
        path[coord[0] * 64 + coord[1]] = 1
        if visited[coord] < minLoss:
            minLoss = visited[coord]
            bestPos = coord

    ps_cloud.add_scalar_quantity("Path", path)
    print("Best position : ", bestPos)


def get_i_j(a, N):
    return a // N, a % N


def get_border_indices(mesh_x):
    N, _ = mesh_x.shape
    group = set([0, 63, 4032, 4095])
    to_visit = set([1, N])
    helpers.add_neighbours_idx(0, 63, N, to_visit, group)
    helpers.add_neighbours_idx(63, 0, N, to_visit, group)
    helpers.add_neighbours_idx(63, 63, N, to_visit, group)
    while to_visit:
        index = to_visit.pop()
        i, j = get_i_j(index, N)
        if mesh_x[i, j]:
            continue
        helpers.add_neighbours_idx(i, j, N, to_visit, group)
        group.add(index)
    in_v = []
    for i in range(N*N):
        if i not in group:
            in_v.append(i)
    return in_v

""" Parse args """
parser = argparse.ArgumentParser()
parser.add_argument('weights', type=str, help='Filepath to model weights')
parser.add_argument('datapath', type=str, help='Filepath to input data')
parser.add_argument('editingType', type=str, nargs='?', default='grab', help='Pretty Huge Depression')
args = parser.parse_args()

""" Initialization / Load model / Normalization """
ps.init()
force_cudnn_initialization()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
N = 64

model = AUNet().to(device)

model.load_state_dict(torch.load(args.weights)['weights'])

# Validation.
model.eval()

# Not database
file = read_file(args.datapath)
point_set = np.loadtxt(file).astype(np.float32)
input_damaged = torch.from_numpy(point_set[:, :3])
input_damaged = input_damaged.transpose(0, 1)
input_damaged = input_damaged.view(3, 64, 64)

mask = (input_damaged[0] != 0) | (input_damaged[1] != 0) | (input_damaged[2] != 0)
in_indices = get_border_indices(mask)

input_normalized, param = normalize(input_damaged, mask)


""" Network Pass """
input_net = input_normalized * mask

arr = input_net.view(3, 4096).transpose(0, 1).detach().clone().cpu().numpy()
nonzero_columns_mask = ~np.all(arr == 0, axis=1)
ps_cloud = ps.register_point_cloud("Init", arr[nonzero_columns_mask, :])
#ps_cloud = ps.register_point_cloud("Init", tmp.view(3, 4096).transpose(0, 1).detach().clone().cpu().numpy())
helpers.write_ply_file(input_net.view(3, 4096).transpose(0, 1).detach().clone().cpu().numpy(), "input.ply")
data = torch.cat((input_net, mask.unsqueeze(0)), dim=0).to(torch.device("cuda"))
data = data.unsqueeze(0)
start = timer()
model(data)
end = timer()

e = 8
if args.editingType == 'grab':
    input_normalized, mask = set_grab(input_normalized, mask)
elif args.editingType == 'twist':
    input_normalized, mask = set_twist(input_normalized, mask)
elif args.editingType == 'pinch':
    input_normalized, mask = set_pinch(input_normalized, mask)

input_net = input_normalized * mask


data = torch.cat((input_net, mask.unsqueeze(0)), dim=0).to(torch.device("cuda"))
data = data.unsqueeze(0)
model(data)
print("timer : ", end - start)

pts = model.im_pred.squeeze(0)
pts = pts.view(3, 4096).transpose(0, 1).detach().clone().cpu().numpy()
faces = np.zeros((7938, 3))
for i in range(N - 1):
    for j in range(N - 1):
        faces[2 * ((N - 1) * i + j)] = [N * i + j, N * (i + 1) + j, N * i + j + 1]
        faces[2 * ((N - 1) * i + j) + 1] = [N * (i + 1) + j, N * (i + 1) + j + 1, N * i + j + 1]

ps_cloud = ps.register_point_cloud("Prediction pc", pts)
ps_mesh = ps.register_surface_mesh("Prediction mesh", pts, faces)
start = timer()
distortions = computeDistortions(input_normalized, mask, param)
end = timer()
print("Timing brute force : ", end - start)
ps_mesh.add_scalar_quantity("Distortion", distortions)

selected_u = 40
selected_v = 40
x = model.im_pred[0, 0, 32, 32]
y = model.im_pred[0, 1, 32, 32]
z = model.im_pred[0, 2, 32, 32]
theta = 0.0
d = 1.0
init_pinch = ((model.im_pred[0, 0, 28, 14], model.im_pred[0, 1, 28, 14], model.im_pred[0, 2, 28, 14]),
              (model.im_pred[0, 0, 28, 43], model.im_pred[0, 1, 28, 43], model.im_pred[0, 2, 28, 43]))

if args.editingType == 'grab':
    target_cloud = ps.register_point_cloud("Target", np.array([[x.item(), y.item(), z.item()]]))
elif args.editingType == 'twist':
    target_cloud = ps.register_point_cloud("Target", np.array([[0.322, 0.457, 0.225], [0.322, 0.457, 0.225], [0.322, 0.457, 0.225], [0.322, 0.457, 0.225], [0.322, 0.457, 0.225]]))
elif args.editingType == 'pinch':
    target_cloud = ps.register_point_cloud("Target", np.array([[0.322, 0.457, 0.225], [0.322, 0.457, 0.225]]))

def callback():
    global selected_u, selected_v, input_normalized, input_net, model, x, y, z, theta, d
    changed_x, x = psim.SliderFloat("change x", x, v_min=-1, v_max=2)
    changed_y, y = psim.SliderFloat("change y", y, v_min=-1, v_max=1)
    changed_z, z = psim.SliderFloat("change z", z, v_min=-1, v_max=2)
    changed_theta, theta = psim.SliderFloat("change theta", theta, v_min=-1.4, v_max=1.4)
    changed_d, d = psim.SliderFloat("change dist", d, v_min=-1.0, v_max=1.0)


    if changed_x or changed_y or changed_z or changed_theta or changed_d:
        if args.editingType == 'grab':
            input_normalized = grab(input_normalized, (x, y, z))
        elif args.editingType == 'twist':
            input_normalized = twist(input_normalized, (x, y, z), theta)
        elif args.editingType == 'pinch':
            input_normalized = pinch(input_normalized, init_pinch, d)

        input_net = input_normalized * mask
        data = torch.cat((input_net, mask.unsqueeze(0)), dim=0).to(torch.device("cuda"))
        data = data.unsqueeze(0)
        model(data)
        ps_cloud.update_point_positions(
            model.im_pred.squeeze(0).view(3, 4096).transpose(0, 1).detach().clone().cpu().numpy())
        ps_mesh.update_vertex_positions(
            model.im_pred.squeeze(0).view(3, 4096).transpose(0, 1).detach().clone().cpu().numpy())


    if psim.Button("Minimize Distortion"):

        start = timer()
        minimize(torch.tensor([0.322, 0.457, 0.225]), (selected_u, selected_v), input_normalized, mask)
        end = timer()
        print("Timing GD : ", end - start)


ps.set_user_callback(callback)
ps.show()
ps.clear_user_callback()