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

import sys

import torch

import torch.optim as optim
import torch.nn as nn

from torch.utils.data import DataLoader

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')

from data.dataset import Dataset
from utils.utils import Accumulator
from models.depth_estim import DepthEstimation
from models.convnext import ConvNeXt




def plot3d(depth, mask):
    import kornia
    import plotly.graph_objects as go
    from data.render import intrinsic
    K = torch.tensor(intrinsic()[None, :, :], device='cuda')
    pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True)
    pc = pc[0,:,mask[0]].detach().cpu().numpy()

    go.Figure(data = [go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),
                        ]).show()


def chamfer_loss(depth_image, vertices, depth_mask, mesh_mask):
    import kornia
    import pytorch3d
    from data.render import intrinsic
    from pytorch3d.structures import Pointclouds
    from utils.pytorch3d_chamfer import chamfer_distance
    from pytorch3d.ops.points_normals import estimate_pointcloud_normals
    from utils.data_utils import get_normals_from_depth
    K = torch.tensor(intrinsic()[None, :, :], device='cuda')
    pc = kornia.geometry.depth.depth_to_3d(depth_image, K, normalize_points=True)
    pc = [torch.permute(pc[i,:,depth_mask[i]], [1,0]) for i in range(depth_image.shape[0])]
    pc_norm = get_normals_from_depth(depth_image, depth_mask, K)
    pc_norm = [pc_norm[i,depth_mask[i],:] for i in range(pc_norm.shape[0])]
    mesh_gt = [vertices[i, mesh_mask[i], :] for i in range(vertices.shape[0])]
    pc_pc = Pointclouds(pc, normals=pc_norm)
    mesh_gt_pc = Pointclouds(mesh_gt)
    mesh_norm = list(torch.unbind(estimate_pointcloud_normals(mesh_gt_pc, neighborhood_size=50)))
    mesh_gt_pc = Pointclouds(mesh_gt, normals=mesh_norm)

    chamfer = chamfer_distance(pc_pc, mesh_gt_pc)

    return chamfer


def normal_loss(depth_image, depth_mask, normal_image_gt):
    import kornia
    import pytorch3d
    from data.render import intrinsic
    K = torch.tensor(intrinsic()[None, :, :], device='cuda')
    pc = kornia.geometry.depth.depth_to_3d(depth_image, K, normalize_points=True)
    pc = torch.stack([pc[i,:,depth_mask[i]] for i in range(depth_image.shape[0])])
    normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K)
    pass

def train_step(batch):
    ### Training
    model.train() # train mode is on by default after construction

    # 1. Forward pass
    y_pred = model([batch['rgb_image'], batch['view']])

    # 2. Calculate loss
    mse = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
    chamfer, normal = chamfer_loss(y_pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
    loss = mse + 1e-1 * (chamfer + normal)

    # 3. Zero grad optimizer
    optimizer.zero_grad()

    # 4. Loss backward
    loss.backward()

    # 5. Step the optimizer
    optimizer.step()

    return { 'train/loss': loss.detach().cpu().numpy() }


def valid_step(batch):
    ### Testing
    model.eval() # put the model in evaluation mode for testing (inference) 
    # 1. Forward pass
    with torch.inference_mode():
        pred = model([batch['rgb_image'], batch['view']])

        # 2. Calculate the loss
        mse = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
        chamfer, normal = chamfer_loss(pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
        loss = mse + 1e-1 * (chamfer + normal)

    return { 'valid/loss': loss.detach().cpu().numpy() }


def batch_to_device(batch, device):
    filter = {'mesh_posed'}
    for key in batch:
        if key in filter:
            batch[key] = batch[key].to(device)


if __name__ == '__main__':
    device = "cuda" if torch.cuda.is_available() else "cpu"

    random_view = False
    use_skip = False
    batch_size = 2

    train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt', random_view=random_view, device=device)
    valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt', random_view=random_view, device=device)
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
    valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
    train_accumulator = Accumulator()
    valid_accumulator = Accumulator()

    model = DepthEstimation(device=device, use_skip=use_skip)

    loss_fn = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    iteration = 0
    while True:

        for batch in train_dataloader:
            train_loss = train_step(batch)
            train_accumulator.update(train_loss)

            if iteration % 1000 == 0:
                print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}")
                train_accumulator.reset()

            iteration += 1

        for batch in valid_dataloader:
            valid_loss = valid_step(batch)
            valid_accumulator.update(valid_loss)

        print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}")
        valid_accumulator.reset()



    pass