class MRnet(nn.Module):
    """MRnet uses pretrained resnet50 as a backbone to extract features
    """
    
    def __init__(self):
        """This function will be used to initialize the 
        MRnet instance."""
        # Initialize nn.Module instance
        super(MRnet,self).__init__()

        # Initialize three backbones for three axis
        # All the three axes will use pretrained AlexNet model
        # The models will be used for extracting features from
        # the input images
        self.axial = models.alexnet(pretrained=True).features
        self.coronal = models.alexnet(pretrained=True).features
        self.saggital = models.alexnet(pretrained=True).features
        
        # Initialize 2D Adaptive Average Pooling layers
        # The pooling layers will reduce the size of
        # feature maps extracted from the previous axes
        self.pool_axial = nn.AdaptiveAvgPool2d(1)
        self.pool_coronal = nn.AdaptiveAvgPool2d(1)
        self.pool_saggital = nn.AdaptiveAvgPool2d(1)
        
        # Initialize a sequential neural network with
        # a single fully connected linear layer
        # The network will output the probability of
        # having a particular disease
        self.fc = nn.Sequential(
            nn.Linear(in_features=3*256,out_features=1)
        )

def forward(self,x):
        """ Input is given in the form of `[image1, image2, image3]` where
        `image1 = [1, slices, 3, 224, 224]`. Note that `1` is due to the 
        dataloader assigning it a single batch. 
        """

        # squeeze the first dimension as there
        # is only one patient in each batch
        images = [torch.squeeze(img, dim=0) for img in x]
        
        # Extract features across each of the three plane
        # using the three pre-trained AlexNet models defined earlier
        image1 = self.axial(images[0])
        image2 = self.coronal(images[1])
        image3 = self.saggital(images[2])

        # Convert the image dimesnsions from [slices, 256, 1, 1] to
        # [slices,256]
        image1 = self.pool_axial(image1).view(image1.size(0), -1)
        image2 = self.pool_coronal(image2).view(image2.size(0), -1)
        image3 = self.pool_saggital(image3).view(image3.size(0), -1)

        # Find maximum value across slices
        # This will reduce the dimensions of image to [1,256]
        # This is done in order to keep only the most prevalent
        # features for each slice
        image1 = torch.max(image1,dim=0,keepdim=True)[0]
        image2 = torch.max(image2,dim=0,keepdim=True)[0]
        image3 = torch.max(image3,dim=0,keepdim=True)[0]

        # Stack the 3 images together to create the output
        # of size [1, 256*3]
        output = torch.cat([image1,image2,image3], dim=1)

        # Feed the output to the sequential network created earlier
        # The network will return a probability of having a specific
        # disease
        output = self.fc(output)
        return output

class MRData():
    """This class used to load MRnet dataset from `./images` dir
    """

    def __init__(self,task = 'acl', train = True, transform = None, weights = None):
        """Initialize the dataset
        Args:
            plane : along which plane to load the data
            task : for which task to load the labels
            train : whether to load the train or val data
            transform : which transforms to apply
            weights (Tensor) : Give wieghted loss to postive class eg. `weights=torch.tensor([2.223])`
        """
        # Define the three planes to use
        self.planes=['axial', 'coronal', 'sagittal']
        # Initialize the records as None
        self.records = None
        # an empty dictionary
        self.image_path={}
        
        # If we are in training loop
        if train:
            # Read data about patient records
            self.records = pd.read_csv('./images/train-{}.csv'.format(task),header=None, names=['id', 'label'])

            for plane in self.planes:
                # For each plane, specify the image path
                self.image_path[plane] = './images/train/{}/'.format(plane)
        else:
            # If we are in testing loop
            # don't use any transformation
            transform = None
            # Read testing/validation data (patients records)
            self.records = pd.read_csv('./images/valid-{}.csv'.format(task),header=None, names=['id', 'label'])
            
            for plane in self.planes:
                # Read path of images for each plane
                self.image_path[plane] = './images/valid/{}/'.format(plane)

        # Initialize the transformation to apply on images
        self.transform = transform 

        # Append 0s to the patient record id
        self.records['id'] = self.records['id'].map(
            lambda i: '0' * (4 - len(str(i))) + str(i))
        # empty dictionary
        self.paths={}
        for plane in self.planes:
            # Get paths of numpy data files for each plane
            self.paths[plane] = [self.image_path[plane] + filename +
                          '.npy' for filename in self.records['id'].tolist()]

        # Convert labels from Pandas Series to a list
        self.labels = self.records['label'].tolist()

        # Total positive cases
        pos = sum(self.labels)
        # Total negative cases
        neg = len(self.labels) - pos

        # Find the wieghts of pos and neg classes
        if weights:
            self.weights = torch.FloatTensor(weights)
        else:
            self.weights = torch.FloatTensor([neg / pos])
        
        print('Number of -ve samples : ', neg)
        print('Number of +ve samples : ', pos)
        print('Weights for loss is : ', self.weights)

    def __len__(self):
        """Return the total number of images in the dataset."""
        return len(self.records)

    def __getitem__(self, index):
        """
        Returns `(images,labels)` pair
        where image is a list [imgsPlane1,imgsPlane2,imgsPlane3]
        and labels is a list [gt,gt,gt]
        """
        img_raw = {}
        
        for plane in self.planes:
            # Load raw image data for each plane
            img_raw[plane] = np.load(self.paths[plane][index])
            # Resize the image loaded in the previous step
            img_raw[plane] = self._resize_image(img_raw[plane])
            
        label = self.labels[index]
        # Convert label to 0 and 1
        if label == 1:
            label = torch.FloatTensor([1])
        elif label == 0:
            label = torch.FloatTensor([0])

        # Return a list of three images for three planes and the label of the record
        return [img_raw[plane] for plane in self.planes], label

    def _resize_image(self, image):
        """Resize the image to `(3,224,224)` and apply 
        transforms if possible.
        """
        # Resize the image
        # Calculate extra padding present in the image
        # which needs to be removed
        pad = int((image.shape[2] - INPUT_DIM)/2)
        # This is equivalent to center cropping the image
        image = image[:,pad:-pad,pad:-pad]
        # Normalize the image by subtracting it by mean and dividing by standard
        # deviation
        image = (image-np.min(image))/(np.max(image)-np.min(image))*MAX_PIXEL_VAL
        image = (image - MEAN) / STDDEV
        
        # If the transformation is not None
        if self.transform:
            # Transform the image based on the specified transformation
            image = self.transform(image)
        else:
            # Else, just stack the image with itself in order to match the required
            # dimensions
            image = np.stack((image,)*3, axis=1)
        # Convert the image to a FloatTensor and return it
        image = torch.FloatTensor(image)
        return image

def load_data(task : str):

    # Define the Augmentation here only
    augments = Compose([
        # Convert the image to Tensor
        transforms.Lambda(lambda x: torch.Tensor(x)),
        # Randomly rotate the image with an angle
        # between -25 degrees to 25 degrees
        RandomRotate(25),
        # Randomly translate the image by 11% of 
        # image height and width
        RandomTranslate([0.11, 0.11]),
        # Randomly flip the image
        RandomFlip(),
        # Change the order of image channels
        transforms.Lambda(lambda x: x.repeat(3, 1, 1, 1).permute(1, 0, 2, 3)),
    ])

    print('Loading Train Dataset of {} task...'.format(task))
    # Load training dataset
    train_data = MRData(task, train=True, transform=augments)
    train_loader = data.DataLoader(
        train_data, batch_size=1, num_workers=11, shuffle=True
    )

    print('Loading Validation Dataset of {} task...'.format(task))
    # Load validation dataset
    val_data = MRData(task, train=False)
    val_loader = data.DataLoader(
        val_data, batch_size=1, num_workers=11, shuffle=False
    )

    return train_loader, val_loader, train_data.weights, val_data.weights

def _train_model(model, train_loader, epoch, num_epochs, optimizer, criterion, writer, current_lr, log_every=100):
    
    # Set to train mode
    model.train()

    # Initialize the predicted probabilities
    y_probs = []
    # Initialize the groundtruth labels
    y_gt = []
    # Initialize the loss between the groundtruth label
    # and the predicted probability
    losses = []

    # Iterate over the training dataset
    for i, (images, label) in enumerate(train_loader):
        # Reset the gradient by zeroing it
        optimizer.zero_grad()
        
        # If GPU is available, transfer the images and label
        # to the GPU
        if torch.cuda.is_available():
            images = [image.cuda() for image in images]
            label = label.cuda()

        # Obtain the prediction using the model
        output = model(images)

        # Evaluate the loss by comparing the prediction
        # and groundtruth label
        loss = criterion(output, label)
        # Perform a backward propagation
        loss.backward()
        # Modify the weights based on the error gradient
        optimizer.step()

        # Add current loss to the list of losses
        loss_value = loss.item()
        losses.append(loss_value)

        # Find probabilities from output using sigmoid function
        probas = torch.sigmoid(output)

        # Add current groundtruth label to the list of groundtruths
        y_gt.append(int(label.item()))
        # Add current probabilities to the list of probabilities
        y_probs.append(probas.item())

        try:
            # Try finding the area under ROC curve
            auc = metrics.roc_auc_score(y_gt, y_probs)
        except:
            # Use default value of area under ROC curve as 0.5
            auc = 0.5
        
        # Add information to the writer about training loss and Area under ROC curve
        writer.add_scalar('Train/Loss', loss_value,
                          epoch * len(train_loader) + i)
        writer.add_scalar('Train/AUC', auc, epoch * len(train_loader) + i)

        if (i % log_every == 0) & (i > 0):
            # Display the information about average training loss and area under ROC curve
            print('''[Epoch: {0} / {1} | Batch : {2} / {3} ]| Avg Train Loss {4} | Train AUC : {5} | lr : {6}'''.
                  format(
                      epoch + 1,
                      num_epochs,
                      i,
                      len(train_loader),
                      np.round(np.mean(losses), 4),
                      np.round(auc, 4),
                      current_lr
                  )
                  )
    # Add information to the writer about total epochs and Area under ROC curve
    writer.add_scalar('Train/AUC_epoch', auc, epoch + i)

    # Find mean area under ROC curve and training loss
    train_loss_epoch = np.round(np.mean(losses), 4)
    train_auc_epoch = np.round(auc, 4)

    return train_loss_epoch, train_auc_epoch

def _evaluate_model(model, val_loader, criterion, epoch, num_epochs, writer, current_lr, log_every=20):
    """Runs model over val dataset and returns auc and avg val loss"""

    # Set to eval mode
    model.eval()
    # List of probabilities obtained from the model
    y_probs = []
    # List of groundtruth labels
    y_gt = []
    # List of losses obtained
    losses = []

    # Iterate over the validation dataset
    for i, (images, label) in enumerate(val_loader):
        # If GPU is available, load the images and label
        # on GPU
        if torch.cuda.is_available():
            images = [image.cuda() for image in images]
            label = label.cuda()

        # Obtain the model output by passing the images as input
        output = model(images)
        # Evaluate the loss by comparing the output and groundtruth label
        loss = criterion(output, label)
        # Add loss to the list of losses
        loss_value = loss.item()
        losses.append(loss_value)
        # Find probability for each class by applying
        # sigmoid function on model output
        probas = torch.sigmoid(output)
        # Add the groundtruth to the list of groundtruths
        y_gt.append(int(label.item()))
        # Add predicted probability to the list
        y_probs.append(probas.item())

        try:
            # Evaluate area under ROC curve based on the groundtruth label
            # and predicted probability
            auc = metrics.roc_auc_score(y_gt, y_probs)
        except:
            # Default area under ROC curve
            auc = 0.5
        # Add information to the writer about validation loss and Area under ROC curve
        writer.add_scalar('Val/Loss', loss_value, epoch * len(val_loader) + i)
        writer.add_scalar('Val/AUC', auc, epoch * len(val_loader) + i)

        if (i % log_every == 0) & (i > 0):
            # Display the information about average validation loss and area under ROC curve
            print('''[Epoch: {0} / {1} | Batch : {2} / {3} ]| Avg Val Loss {4} | Val AUC : {5} | lr : {6}'''.
                  format(
                      epoch + 1,
                      num_epochs,
                      i,
                      len(val_loader),
                      np.round(np.mean(losses), 4),
                      np.round(auc, 4),
                      current_lr
                  )
                  )
    # Add information to the writer about total epochs and Area under ROC curve
    writer.add_scalar('Val/AUC_epoch', auc, epoch + i)
    # Find mean area under ROC curve and validation loss
    val_loss_epoch = np.round(np.mean(losses), 4)
    val_auc_epoch = np.round(auc, 4)

    return val_loss_epoch, val_auc_epoch