#run_neuro3dunet.py#

import os
import pdb
import h5py # note: importing h5py multiple times can cause an error
import numpy as np
import pandas as pd

import torch as t
import torch.nn as nn
import torch.nn.functional as f
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

# translated from imports of unet3d.model
from unet3d.buildingblocks import DoubleConv, ResNetBlock, ResNetBlockSE, \
    create_decoders, create_encoders
from unet3d.utils import get_class, number_of_features_per_level
from unet3d.model import UNet3D

import logging
logging.basicConfig(filename="unet.log", level=logging.DEBUG, format='%(asctime)s - %(message)s',
                    datefmt='%y-%m-%d %H:%M:%S')

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

def save_checkpoint(checkpoint_dir, epoch, model, optimizer):
    checkpoint_path = os.path.join(checkpoint_dir, f"checkpoint_epoch_{epoch}.pt")
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }
    torch.save(checkpoint, checkpoint_path)


def load_checkpoint(checkpoint_dir, model, optimizer=None):
    checkpoint_list = [f for f in os.listdir(checkpoint_dir) if f.startswith('checkpoint_epoch_')]
    if not checkpoint_list:
        return 0  # Start training from epoch 0 if no checkpoints are found

    latest_checkpoint = max(checkpoint_list)
    checkpoint_path = os.path.join(checkpoint_dir, latest_checkpoint)
    checkpoint = torch.load(checkpoint_path)

    model.load_state_dict(checkpoint['model_state_dict'])
    if optimizer is not None:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

    return checkpoint['epoch']


class HDF5Dataset(Dataset):

    """ A custom Dataset class to iterate over subjects.
        This Dataset assumes that the data take the following form:
            data_dir/
                -- subject0.hdf5 (file with two datasets)
                    -- x_name: 4D array
                      -- y_name: 4D array
                -- subject1.hdf5 (next file with two datasets)
                    -- ...
        Note also that this directory should not contain any other files
        besides h5 files for subjects intended to be included in this dataset.
        -----
        Arguments:
            data_dir
            x_name
            y_name
            ordered_subject_list
        -----       
        Returns:
            Pytorch index-based Dataset where each sample is an x, y pair of tensors
                corresponding to a 3D T1 scan and a 4D set of anatomical labels (one-hot)
        
    """
    
    def __init__(self, 
                 data_dir, 
                 x_name=None,
                 y_name=None,
                 ordered_subject_list=None):
        
        self.data_dir = data_dir

        # parse default args
        x_name = 'raw' if x_name is None else x_name
        y_name = 'label' if y_name is None else y_name
        self.x_name = x_name
        self.y_name = y_name
        
        # parse subject ordering, if specified
        if ordered_subject_list is None:
            ordered_subject_list = sorted(os.listdir(data_dir))
        self.subjects = ordered_subject_list
        

    def __len__(self):
        return len(self.subjects)
    

    def __getitem__(self, index):
        subject = self.subjects[index]  # Select the current datapoint (subject)    
        h5 = h5py.File(f'{self.data_dir}/{subject}', 'r')
        
        x_np = h5.get(self.x_name)
        y_np = h5.get(self.y_name)
        
        x = t.from_numpy(np.array(x_np))
        y = t.from_numpy(np.array(y_np))

        h5.close() # close the h5 file to avoid extra memory usage

        # If necessary, apply any preprocessing or transformations to the data
        # data = ...
        #x = x[:, :64, :64, :64]
        #y = y[:, :64, :64, :64]

        return x, y
    

def train(dl_train, 
          dl_val, 
          model, 
          optimizer,
          criterion,
          device,
          checkpoint_dir,
          epochs=10,
          lr_scheduler=None
          ):
    """ Function to wrap the main training loop.
    """
    
    # parse default parameters
    if lr_scheduler is None:
        default_lr_scheduler_patience = 3
        default_lr_scheduler_factor = 0.1
        lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 
                                                            patience=default_lr_scheduler_patience, 
                                                            factor=default_lr_scheduler_factor
                                                           )
  
    # Training loop
    best_val_loss = float('inf')

    # Load the last saved epoch and optimizer state (if available)
    start_epoch = load_checkpoint(checkpoint_dir, model, optimizer)
    
    for epoch in range(epochs):
        # Training
        logging.debug("\nBeginning epoch \n")
        model.train()
        train_loss = 0.0
        sample_set = 0
        for inputs, labels in dl_train:
            print("-------------------------\n")
            logging.debug(f"New octant {sample_set + 1} / 3840 in epoch {epoch} / {epochs}\n\n")
            print(f"New octant {sample_set + 1} / 3840 in epoch {epoch} / {epochs}\n\n")
            print("Creating optimizer")
            optimizer.zero_grad()
            print("Sending data to device")
            inputs = inputs.to(torch.bfloat16).to(device, dtype=float)
            labels = labels.to(torch.bfloat16).to(device, dtype=float)
            # Forward pass

            ## comments showing cuda activity
            #print(torch.cuda.memory_summary(abbreviated=False))
            #os.system("nvidia-smi")
            print("Calculating outputs")
            outputs = model(inputs)
            print("Calculating loss")
            loss = criterion(outputs, labels)

            # Backward pass and optimization
            print("Backward pass")
            loss.backward()
            optimizer.step()

            
            train_loss += loss.item()
            inputs.detach()
            labels.detach()
            torch.cuda.empty_cache()
            sample_set += 1
            save_checkpoint(checkpoint_dir, epoch, model, optimizer)
            logging.debug("Saving checkpoint")
        # Validation
        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for inputs, labels in dl_val:
                inputs = inputs.to(torch.bfloat16).to(device, dtype=float)
                labels = labels.to(torch.bfloat16).to(device, dtype=float)
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                val_loss += loss.item()
                inputs.detach()
                labels.detach()
                torch.cuda.empty_cache()
        # Compute average loss
        
        train_loss /= len(dl_train)
        val_loss /= len(dl_val)

        # Update learning rate scheduler
        lr_scheduler.step(val_loss)

        # Print progress
        print(f"Epoch {epoch+1}/{epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")
        logging.info(f"Epoch {epoch+1}/{epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

        # Check if current validation loss is the best so far
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            # Save the model checkpoint if desired

        # Check early stopping condition if desired
        # TODO

    # Training complete
    return model
    
def main():
    train_dir = '/home/weiner/bparker/NotBackedUp/train_chunks'
    val_dir = '/home/weiner/bparker/NotBackedUp/test_chunks'
    
    #train_dir = 'data/h5/train'
    #val_dir = 'data/h5/val'
    # option to use DistributedSampler to distribute data over multiple GPUs

    batch_size = 1

    ds_train = HDF5Dataset(data_dir=train_dir)
    ds_val = HDF5Dataset(data_dir=val_dir)

    dl_train = DataLoader(ds_train, batch_size=batch_size, shuffle=True)
    dl_val = DataLoader(ds_val, batch_size=batch_size, shuffle=False)
    print("Loaded Datasets\n")


    
    ## Define model
    in_channels = 1
    out_channels = 102
    print(f"Creating model with {in_channels} in channels, and {out_channels} out channels.")
    model = UNet3D(in_channels=in_channels, out_channels=out_channels)

    # specifying "float" may or may not be necessary on the GPU
    # but it is required on CPU

    print("Checking for CUDA availability. . . \n")
    if torch.cuda.is_available():
        # GPU is available
        print("CUDA is available! \nAssigning model to CUDA")
        device = torch.device("cuda")

        ### DistributedDataParallel chunk
        #os.environ['MASTER_ADDR'] = 'localhost'
        #os.environ['MASTER_PORT'] = '12345'
        #torch.distributed.init_process_group(backend='nccl',
        #                                     init_method = "env://",
        #                                     world_size = 2,
        #                                     rank = 1)
        #model = nn.parallel.DistributedDataParallel(model, device_ids=[0,1])
        #model = nn.DataParallel(model, device_ids=[0,1])
        model.to(device, dtype=float)
    else:
        # GPU is not available, fall back to CPU
        device = torch.device("cpu")
        model.to(device, dtype=float) 


        
    ### DistributedDataParallel chunk
    ## When this is run, it hangs at this stage
    ## Attempting to distribute the data with DistributedSampler()
    
    #dl_train = DataLoader(ds_train, batch_size=batch_size, shuffle=True)
    #dl_val = DataLoader(ds_val, batch_size=batch_size, shuffle=False)
    #print("Loaded Datasets\n")


    #ds_train = HDF5Dataset(data_dir=train_dir)
    #ds_train = DistributedSampler(ds_train)
    
    #ds_val = HDF5Dataset(data_dir=val_dir)
    #ds_val = DistributedSampler(ds_val)
    #####

    ## Training
    
    checkpoint_dir = './checkpoints'  # change this based on your OS and preferences
    os.makedirs(checkpoint_dir, exist_ok=True)
    print("Defining optimizer for model")
    optimizer = optim.Adam(model.parameters())
    criterion = nn.CrossEntropyLoss()
    
    # Other training parameters
    epochs = 10
    lr_scheduler_patience = 3
    lr_scheduler_factor = 0.1
    logging.info("\nTraining model\n\nepochs = {epochs}\nlr_scheduler_patience = {lr_scheduler_patience}\nlr_scheduler_factor = {lr_scheduler_factor}")
    print("\nTraining model\n\nepochs = {epochs}\nlr_scheduler_patience = {lr_scheduler_patience}\nlr_scheduler_factor = {lr_scheduler_factor}")
    #pdb.set_trace()
    train(dl_train, dl_val, model, optimizer, criterion, device, checkpoint_dir)

if __name__ == "__main__":
    main()