import os
import sys

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0' #config['gpu']
import torch

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

import importlib
import argparse
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


import wandb
import numpy as np
import torch.optim as optim
import torch.nn as nn
from torch.utils.data import DataLoader
# from torchgeometry.losses.dice import DiceLoss
# from torchmetrics import Dice

from data.view_dataset import Dataset
from utils.utils import Accumulator
from models.view_encoder import ViewEncoder
import plotly.express as ex


def init_wandb():
    run = wandb.init(project='reconstruction', name='view_1', notes='', dir='/tmp', tags=['view'])
    artifact = wandb.Artifact('code', type='code')
    artifact.add_file('scripts/view_encoder_train.py')
    artifact.add_file('models/view_encoder.py')
    artifact.add_file('data/view_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 OrthogonalLoss(nn.Module):
    def __init__(self):
        super(OrthogonalLoss, self).__init__()

    def forward(self, tensor):
        # Reshape the input tensor into a square matrix
        matrix = tensor.view(-1, tensor.size(-1)//2, tensor.size(-1)//2)
        # matrix = matrix / torch.norm(matrix, dim=1, keepdim=True)  # Normalize the rows

        # Compute the Gram matrix
        gram_matrix = torch.matmul(matrix, matrix.t())

        # Compute the Frobenius norm of the difference between the Gram matrix and the identity matrix
        loss = torch.norm(gram_matrix - torch.eye(gram_matrix.size(0)).to(tensor.device))

        return loss
    

def train_step(batch):
    model.train()

    pred, matrix = model(batch['matrix'])

    mse = loss_fn(pred, batch['matrix'])
    ortho_loss = ortho_loss_fn(matrix)
    loss = mse + ortho_loss

    # TODO: calculate max of ortho_loss

    # print(f'Loss: {loss} MSE: {mse}')

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


    return { 'train/loss': loss, 'train/mse': mse, 'train/ortho_loss': ortho_loss }


def valid_step(batch):
    model.eval()
    with torch.inference_mode():
        pred, matrix = model(batch['matrix'])

        mse = loss_fn(pred, batch['matrix'])
        ortho_loss = ortho_loss_fn(matrix)
        loss = mse + ortho_loss


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


def plot_everything(batch, type):
    model.eval()
    with torch.inference_mode():
        pred = model(batch['matrix'])

    matrix = batch['matrix'].cpu().numpy()
    pred = pred.detach().cpu().numpy()

    matrix_str = [np.array2string(matrix[i], precision=2, suppress_small=True) for i in range(4)]
    pred_str = [np.array2string(pred[i], precision=2, suppress_small=True) for i in range(4)]

    logs = { f'{type}/matrix': wandb.Table(data=[[matrix_str, pred_str]], columns=['matrix', 'pred']), }

    return logs


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

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

    batch_size = 128
    use_cache = False

    train_dataset = Dataset(device=device)
    valid_dataset = Dataset(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 = ViewEncoder(bottleneck_size=1024).to(device)

    loss_fn = nn.MSELoss()
    ortho_loss_fn = OrthogonalLoss()
    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)

    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 < 1000 and iteration % 100 == 0):
                to_log = dict({ **train_loss,  **{ 'current_step': iteration } })
                if iteration % 1000 == 0 or (iteration < 1000 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)) } }

                if not args['debug']:
                    wandb.log(to_log, step=iteration)
                train_accumulator.reset()
            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)




    pass