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

import argparse
import importlib

parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())

# config_module = importlib.import_module(args['config'])
# config = config_module.config
if args['debug']:
    pass
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
if args['debug']:
    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'


import wandb
import torch
import torch.optim as optim
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.profiler import profile, record_function, ProfilerActivity
# from torchgeometry.losses.dice import DiceLoss
# from torchmetrics import Dice

from data.pose_dataset import Dataset
from utils.utils import Accumulator
from models.pose_encoder import PoseEncoder as PoseEncoder
from models.convnext import ConvNeXt
import plotly.express as ex

from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)


def init_wandb():
    run = wandb.init(project='reconstruction', name='pose_2', notes='have the back of the body also as the main input', dir='/data-local/127/temp', tags=['pose'])
    artifact = wandb.Artifact('code', type='code')
    artifact.add_file('scripts/pose_encoder_train.py')
    artifact.add_file('models/pose_encoder.py')
    artifact.add_file('data/pose_dataset.py')
    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
    for config_file_name in config_file_names:
        file_name = '/'.join(config_file_name.split('.')) + '.py'
        if os.path.exists(file_name):
            artifact.add_file(file_name)
    run.log_artifact(artifact)


class SymmetricLoss(nn.Module):
    def __init__(self):
        super(SymmetricLoss, self).__init__()

    def forward(self, x):
        # Check the difference between the matrix and its transpose
        diff = x - x.transpose(1, 2)
        # Compute the Frobenius norm of the difference
        loss = torch.norm(diff, p='fro', dim=(1, 2))
        # Return the average loss over the batch
        return torch.mean(loss)
    

def train_step(batch):
    model.train()

    # pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
    pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']])

    mse = loss_fn(pred_pose, batch['pose_out'])

    loss = mse #+ mse_mask + symm_rotation
    # print(f'Loss: {loss} MSE: {mse} MSE Mask: {mse_mask} Symm Rotation: {symm_rotation}')

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    with torch.inference_mode():
        mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
        symm_rotation = loss_fn_rotation(pred_rotation)

    return { 'train/loss': loss, 'train/mse': mse, 'train/mse_mask': mse_mask, 'train/symm_rotation': symm_rotation }


def valid_step(batch):
    model.eval()
    with torch.inference_mode():
        pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']])

        mse = loss_fn(pred_pose, batch['pose_out'])
        mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
        symm_rotation = loss_fn_rotation(pred_rotation)
        loss = mse

    return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation }


def plot_everything(batch, type: str):
    model.eval()
    with torch.inference_mode():
        pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']])

    mask = torch.permute(batch['mask'], (0, 2, 3, 1))
    mask_out = torch.permute(batch['mask_out'], (0, 2, 3, 1))
    pose = torch.permute(batch['pose'], (0, 2, 3, 1))
    pose_back = torch.permute(batch['pose_back'], (0, 2, 3, 1))
    pose_out = torch.permute(batch['pose_out'], (0, 2, 3, 1))
    pred_mask = torch.permute(pred_mask, (0, 2, 3, 1))
    pred_pose = torch.permute(pred_pose, (0, 2, 3, 1))

    masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1).cpu().numpy()
    poses = torch.concat((pose, pose_back, pose_out, pred_pose), axis=2).cpu().numpy()
    masks_imgs = [wandb.Image((masks[i] + 1) / 2.0) for i in range(pred_pose.shape[0])]
    poses_imgs = [wandb.Image((poses[i] + 1) / 2.0) for i in range(pred_pose.shape[0])]

    # do a batchwise min-max normalization of pred_rotation
    pred_rotation_img = pred_rotation.cpu().numpy()
    pred_rotation_img = (pred_rotation_img - pred_rotation_img.min(axis=(1, 2,), keepdims=True)) / (pred_rotation_img.max(axis=(1, 2), keepdims=True) - pred_rotation_img.min(axis=(1, 2), keepdims=True))
    pred_rotation_imgs = [wandb.Image(pred_rotation_img[i]) for i in range(pred_pose.shape[0])]

    return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, f'{type}/rotation': pred_rotation_imgs} }


def train():
    iteration = 0
    epoch = 0
    while True:

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

            if iteration % 1000 == 0 or (iteration < 500 and iteration % 100 == 0):
                to_log = dict({ **train_loss,  **{ 'current_step': iteration } })
                if iteration % 5000 == 0 or (iteration < 5000 and iteration % 1000 == 0):
                    train_figs = plot_everything(batch, 'train')
                    to_log = { **to_log, **train_figs }
                # if iteration == 0:
                #     to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }

                train_accumulator.reset()

                # figure out when to log when using validation set as well
                if not args['debug']:
                    wandb.log(to_log, step=iteration)
            iteration += 1
            # print(f'Iteration {iteration} finished')

        epoch += 1
        print(f'Epoch {epoch} finished')

        # if not iteration % 1000 == 0:
        #     continue

        # for batch in valid_dataloader:
        #     valid_loss = valid_step(batch)
        #     valid_accumulator.update(valid_loss)
            
        #     to_log = dict({ **valid_loss, **{ 'current_step': iteration } })
        #     valid_figs = plot_everything(batch, 'train')
        #     to_log = { **to_log, **valid_figs }
        #     valid_accumulator.reset()

        # if not args['debug']:
        #     wandb.log(to_log, step=iteration)


if __name__ == '__main__':
    if not args['debug']:
        wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d')
        init_wandb()

    device = "cuda" if torch.cuda.is_available() else "cpu"

    batch_size = 4
    use_cache = False

    train_dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache)
    valid_dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache)
    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 = PoseEncoder(device=device).to(device)

    loss_fn = nn.MSELoss()
    loss_fn_mask = nn.BCELoss(reduction='mean')
    loss_fn_rotation = SymmetricLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=False)
    # optimizer = optim.SGD(model.parameters(), lr=0.1)

# Use the profiler to track memory usage
    # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
    #     with record_function("model_inference"):
    train()

    # print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))

    pass