import sys
import os
import logging
import numpy as np
import torch
import torch.nn as nn
from tqdm import tqdm
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import monai
from monai.data import ImageDataset
from monai.networks.layers.factories import Conv, Dropout, Pool, Norm
from monai.transforms import Compose, AddChannel, ScaleIntensity, EnsureType
from collections import OrderedDict
from typing import Callable, Sequence
from monai.networks.nets import ViT

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


def main():
    # Training data paths
    print("ViT-test")
    data_dir = '/home/marafath/scratch/abide_iq'
    np.random.seed(42)   
    idx = np.random.permutation(850)
    datalist = os.listdir(data_dir)

    fold = 0
    train_image = []
    train_label = []
    validation_image = []
    validation_label = []

    start = fold*170
    end = (fold+1)*170

    for i in range(len(datalist)):
        nifti_name = datalist[idx[i]]
        iq_label = float(nifti_name.split('_')[2])

        if i >= start and i < end:
            validation_image.append(os.path.join(data_dir,nifti_name))
            validation_label.append(iq_label)
        else:
            train_image.append(os.path.join(data_dir,nifti_name))
            train_label.append(iq_label)
            
    # Define transforms
    train_transforms = Compose([ScaleIntensity(), AddChannel(), EnsureType()])
    val_transforms = Compose([ScaleIntensity(), AddChannel(), EnsureType()])
    
    # create a training data loader
    train_ds = ImageDataset(image_files=train_image, labels=train_label, transform=train_transforms)
    train_loader = DataLoader(train_ds, batch_size=1, shuffle=True, num_workers=1, pin_memory=pin_memory)

    # create a validation data loader
    val_ds = ImageDataset(image_files=validation_image, labels=validation_label, transform=val_transforms)
    val_loader = DataLoader(val_ds, batch_size=1, num_workers=1, pin_memory=pin_memory)
    
    # Create ViT
    model = ViT(in_channels=1, img_size=(140,175,130), patch_size=(32,32,32), pos_embed='conv', post_activation=None).to(device)
    loss_function = torch.nn.L1Loss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)
    
    def mae(t1, t2):
        diff = t1 - t2
        return torch.absolute(diff) #/ diff.numel()
    
    val_interval = 1
    best_metric = 1e5
    best_metric_epoch = -1
    epoch_loss_values = []
    metric_values = []
    total_loss = 0
    max_epochs = 50
    
    for epoch in range(max_epochs):
        print("-" * 10)
        print(f"epoch {epoch + 1}/{max_epochs}")
        model.train()
        epoch_loss = 0
        step = 0
        for batch_data in train_loader:
            step += 1
            inputs, labels = batch_data[0].to(device), batch_data[1].to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            print(f"model outputs {outputs}")
            print(f"Actual labels {labels}")
            #outputs = outputs.squeeze(-1)
            loss = loss_function(outputs, labels)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
            epoch_len = len(train_ds) // train_loader.batch_size
            print(f"{step}/{epoch_len}, train_loss: {loss.item():.4f}")

        epoch_loss /= step
        epoch_loss_values.append(epoch_loss)
        print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
    
        if (epoch + 1) % val_interval == 0:
            model.eval()
            metric_count = 0
            total_loss = 0
            for val_data in val_loader:
                val_images, val_labels = val_data[0].to(device), val_data[1].to(device)
                with torch.no_grad():
                    val_outputs = model(val_images)
                    #val_outputs = val_outputs.squeeze(-1)
                    val_loss = mae(val_outputs, val_labels)
                    metric_count += len(val_loss)
                    total_loss += torch.sum(val_loss)

            metric = total_loss / metric_count
            metric_values.append(metric)

            if metric < best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                torch.save(model.state_dict(), '/home/marafath/scratch/abide_saved_models/best_vit_model_f{}.pth'.format(fold))
                print("saved new best metric model")

            print(f"Current epoch: {epoch+1} current accuracy: {metric:.4f} ")
            print(f"Best accuracy: {best_metric:.4f} at epoch {best_metric_epoch}")

    print(f"Training completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}")
    
if __name__ == '__main__':
    main()