import numpy as np
import pandas as pd
import pydicom
import cv2
from sklearn.model_selection import train_test_split
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from keras_preprocessing.image.dataframe_iterator import DataFrameIterator

import os
"""
### Quick Notebook on making a dataframe iterator for training in tensorflow
"""
#Built from https://gist.github.com/ReemRashwan/8c92086d3104d01978a16e05ca93a165
class DCMFlowFromDataFrame(DataFrameIterator):
    def __init__(self, *arg, **kwargs):
        self.white_list_formats = ('dcm')
        super(DCMFlowFromDataFrame, self).__init__(*arg, **kwargs)
        self.dataframe = kwargs['dataframe']
        self.x = self.dataframe[kwargs['x_col']]
        self.y = self.dataframe[kwargs['y_col']]
        self.target_size = kwargs['target_size']

    def _get_batches_of_transformed_samples(self, indices_array):
        batch_x = np.array([self.read_dcm_as_array(dcm_path, self.target_size)
                            for dcm_path in self.x.iloc[indices_array]])

        batch_y = np.array(self.y.iloc[indices_array].astype(np.uint8))

        if self.image_data_generator is not None:
            for i, (x, y) in enumerate(zip(batch_x, batch_y)):
                transform_params = self.image_data_generator.get_random_transform(x.shape)
                batch_x[i] = self.image_data_generator.apply_transform(x, transform_params)

        return batch_x, batch_y

    @staticmethod
    def read_dcm_as_array(dcm_path, target_size=(256, 256)):
        image_array = pydicom.dcmread(dcm_path).pixel_array
        image_array = cv2.resize(image_array, target_size, interpolation=cv2.INTER_NEAREST)
        image_array = np.expand_dims(image_array, -1)
        
        return image_array
"""
### Creating a full df of paths to valid images (i.e. images that are not empty), with the labels of each study
"""
train = pd.read_csv("../input/rsnamiccai-brain-tumor-radiogenomic-valid-images/valid_trains.csv", dtype=str)
train.head()
train_labels = pd.read_csv('../input/rsna-miccai-brain-tumor-radiogenomic-classification/train_labels.csv',dtype=str)
train = train.merge(train_labels,left_on='sub_dir',right_on='BraTS21ID').drop('BraTS21ID',axis=1)
train['path'] = np.vectorize(os.path.join)(train['main_dir'], train['sub_dir'], train['mpMRI_scan'], train['file'])
train = train.drop(['main_dir','sub_dir','mpMRI_scan','file'],axis=1)
"""
### Sample train for some rows to test on
"""
sample = train.sample(n=320, random_state=1)
sample.head()
"""
### Model without much thought just to test on
"""
def make_model():
    model = tf.keras.Sequential()
    model.add(tf.keras.layers.Conv2D(32,(3,3),activation='relu',input_shape=(256,256,1)))
    model.add(tf.keras.layers.MaxPooling2D((2,2)))
    model.add(tf.keras.layers.Conv2D(32,(3,3),activation='relu'))
    model.add(tf.keras.layers.MaxPooling2D((2,2)))
    model.add(tf.keras.layers.Conv2D(32,(3,3),activation='relu'))
    model.add(tf.keras.layers.MaxPooling2D((2,2)))
    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(128,activation='relu'))
    model.add(tf.keras.layers.Dense(32,activation='relu'))
    model.add(tf.keras.layers.Dense(1,activation='sigmoid'))
    
    model.compile(loss='binary_crossentropy',optimizer=tf.keras.optimizers.Adam(),metrics=[tf.keras.metrics.AUC()])
    return model
model = make_model()
train_datagen = ImageDataGenerator(rescale=1./255)
train_generator = DCMFlowFromDataFrame(dataframe=sample,
                                       x_col='path',
                                       y_col='MGMT_value',
                                       image_data_generator=train_datagen,
                                       target_size=(256,256),
                                       batch_size=32,
                                       class_mode='binary')
history = model.fit(
    train_generator,
    epochs=1
)