import torch
import torchvision

# FOR DATA
from utils.preprocess import prepare_datasets
from utils.show_image import show_image
from torch.utils.data import DataLoader
from torchvision import datasets

from torch import nn
from torchvision.transforms import ToTensor

# import nonechucks as nc     # Used to load data in pytorch even when images are corrupted / unavailable (skips them)

# FOR IMAGE VISUALIZATION
import nibabel as nib

# GENERAL PURPOSE
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import glob



print("--- RUNNING ---")
print("Pytorch Version: " + torch. __version__)


# MAYBE??
'''
import sys
sys.path.append('//data/data_wnx3/data_wnx1/rschuurs/CNN+RNN-2class-1cnn-CLEAN/utils')

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # use id from $ nvidia-smi
'''

# LOADING DATA
# data & training properties:
val_split = 0.2     # % of val and test, rest will be train
seed = 12       # TODO Randomize seed
'''
target_rows = 91
target_cols = 109
depth = 91
axis = 1
num_clinical = 2
CNN_drop_rate = 0.3
RNN_drop_rate = 0.1
CNN_w_regularizer = regularizers.l2(2e-2)
RNN_w_regularizer = regularizers.l2(1e-6)
CNN_batch_size = 10
RNN_batch_size = 5
val_split = 0.2
optimizer = Adam(lr=1e-5)
final_layer_size = 5
'''


# Might have to replace datapaths or separate between training and testing
model_filepath = '//data/data_wnx1/rschuurs/Pytorch_CNN-RNN'
mri_datapath = './ADNI_volumes_customtemplate_float32/'
annotations_datapath = './LP_ADNIMERGE.csv'

# annotations_file = pd.read_csv(annotations_datapath)    # DataFrame

# show_image(17508)

# TODO: Datasets include multiple labels, such as medical info
training_data, val_data, test_data = prepare_datasets(mri_datapath, val_split, seed)
batch_size = 64

# Create data loaders
train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=True)
val_dataloader = DataLoader(val_data, batch_size=batch_size, shuffle=True)

for X, y in train_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break


# Display 10 images and labels.
x = 0
while x < 10:
    train_features, train_labels = next(iter(train_dataloader))
    print(f"Feature batch shape: {train_features.size()}")
    img = train_features[0].squeeze()
    image = img[:, :, 40]
    label = train_labels[0]
    plt.imshow(image, cmap="gray")
    plt.show()
    print(f"Label: {label}")
    x = x+1

print("--- END ---")

# EXTRA

# will I need these params?
'''
params_dict = { 'CNN_w_regularizer': CNN_w_regularizer, 'RNN_w_regularizer': RNN_w_regularizer,
               'CNN_batch_size': CNN_batch_size, 'RNN_batch_size': RNN_batch_size,
               'CNN_drop_rate': CNN_drop_rate, 'epochs': 30,
          'gpu': "/gpu:0", 'model_filepath': model_filepath,
          'image_shape': (target_rows, target_cols, depth, axis),
          'num_clinical': num_clinical,
          'final_layer_size': final_layer_size,
          'optimizer': optimizer, 'RNN_drop_rate': RNN_drop_rate,}

params = Parameters(params_dict)

# WHAT WAS THIS AGAIN?
seeds = [np.random.randint(1, 5000) for _ in range(1)]

# READ THIS TO UNDERSTAND TRAIN VS VALIDATION DATA
def evaluate_net (seed):
    n_classes = 2
    data_loader = DataLoader((target_rows, target_cols, depth, axis), seed = seed)
    train_data, val_data, test_data,rnn_HdataT1,rnn_HdataT2,rnn_HdataT3,rnn_AdataT1,rnn_AdataT2,rnn_AdataT3, test_mri_nonorm = data_loader.get_train_val_test(val_split, mri_datapath)

    print('Length Val Data[0]: ',len(val_data[0]))
'''