def read_dicom(path):
    '''This function reads dicoms and sorts them by Image position patient attribute ''' 
    slices = [pydicom.dcmread(path + '/' + s, force=True) for s in os.listdir(path)]
    slices.sort(key = lambda x: float(x.ImagePositionPatient[2]))
    return slices

def dicom_to_array(slices):
    '''This function extracts pixel data from the pydicom objects '''
    image = np.stack([s.pixel_array for s in slices])
    image = image.astype(np.int16)
    return image

def convert_to_hu(dicom_image_serie, dicom_image_array):
    '''This function converts pixel images to HU units based on Dicom metadata  RescaleIntercept and RescaleSlope '''
    intercept = dicom_image_serie[0].RescaleIntercept
    slope = dicom_image_serie[0].RescaleSlope
    hu_image = dicom_image_array * slope + intercept
    
    return hu_image

""" Full assembly of the parts to form the complete network """

from .unet_parts import *


class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=False):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = (DoubleConv(n_channels, 64))
        self.down1 = (Down(64, 128))
        self.down2 = (Down(128, 256))
        self.down3 = (Down(256, 512))
        factor = 2 if bilinear else 1
        self.down4 = (Down(512, 1024 // factor))
        self.up1 = (Up(1024, 512 // factor, bilinear))
        self.up2 = (Up(512, 256 // factor, bilinear))
        self.up3 = (Up(256, 128 // factor, bilinear))
        self.up4 = (Up(128, 64, bilinear))
        self.outc = (OutConv(64, n_classes))

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

""" Parts of the U-Net model """

import torch
import torch.nn as nn
import torch.nn.functional as F


class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels, mid_channels=None):
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)


class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])
        # if you have padding issues, see
        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

from totalsegmentator.python_api import totalsegmentator
output_path ="inference/organ_masks_demo2/'"
totalsegmentator(path, output_path, task ="body")
totalsegmentator(path, output_path, task ="lung_vessels")

def inference_array(organ,wed_array, mA, path, slice_location, path_output):
    '''This function generates inference data sturcture from organ mask and dicom data and saves it the path_output'''
    mask = np.rot90(nib.load(path + organ + '.nii.gz').get_fdata())
    mask = np.transpose(mask, (2, 0, 1))
    mask[mask != 1] = np.nan

    #retrieving the organ's HU values
    masked_organ = hu_image*mask

    organ_training_array = np.empty((len(dicom_image_serie), 5), dtype=np.float32)
    idx = []
    for j in range(len(dicom_image_serie)):
        if np.all(np.isnan(masked_organ[j])):
            idx.append(j)
        else:
            organ_training_array[j][0] = np.nanmean(masked_organ[j])
            organ_training_array[j][1] = np.nanstd(masked_organ[j])
            organ_training_array[j][2] = float(mA[j])
            organ_training_array[j][3] = wed_array[j]
            organ_training_array[j][4] = float(slice_location[j])

    clean_organ_training_array = np.delete(organ_training_array, idx, axis=0)

    #save the inference array
    os.makedirs(path_output + organ, exist_ok=True)
    file_name = f'{path_output}/{organ}/{organ}_inference_array.npy'
    np.save(file_name, clean_organ_training_array)

#scaler for the organ
data = load('training_arrays/chest_merged_training_arrays/'+ organ + '_training_array.npy')
X = data[:,:-1]

# fit scaler 
scaler = StandardScaler()
X = scaler.fit_transform(X)

import torch.nn as nn

class RegressionModel(nn.Module):
    def __init__(self, input_size, hidden_sizes, output_size, num_layers, dropout):
        super(RegressionModel, self).__init__()

        layers = []
        layers.append(nn.Linear(input_size, hidden_sizes[0]))
        layers.append(nn.ReLU())

        for i in range(1, num_layers - 1):
            layers.append(nn.Linear(hidden_sizes[i - 1], hidden_sizes[i]))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout))
        
        layers.append(nn.Linear(hidden_sizes[num_layers - 2], output_size))
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)

#Model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
hidden_sizes = [512, 256, 128, 64]
model = RegressionModel(input_size=5, hidden_sizes=hidden_sizes, output_size=1, num_layers=5, dropout=0.1)


print('Predicted dose for the organs of interest' )
print('-----------------------------------------')

def evaluate(organ):

    #model loading for the organ
    model.load_state_dict(torch.load('chest_models/'+ organ + '_model.pth'))
    model.to(device)

    #scaler for the organ
    data = load('training_arrays/chest_merged_training_arrays/'+ organ + '_training_array.npy')
    X = data[:,:-1]
    
    # Find rows containing NaN values in X
    nan_rows_X = np.isnan(X).any(axis=1)
    
    # Drop rows with NaN values from X 
    X = X[~nan_rows_X]
    
    # fit scaler 
    scaler = StandardScaler()
    X = scaler.fit_transform(X)

    X_test  = load(f'inference/inference_arrays/{organ}/{organ}_inference_array.npy')
    
    # Find rows containing NaN values in X
    nan_rows_X_test = np.isnan(X_test).any(axis=1)

    # Drop rows with NaN values from X and y
    X_test = X_test[~nan_rows_X_test]
    
    #scaling
    X_test = scaler.transform(X_test)
    
    #data to tensor
    X_test = torch.tensor(X_test, dtype=torch.float32).to(device)

    #evaluate
    with torch.no_grad():
        model.eval()
        predictions = model(X_test).cpu().numpy()

    print(f'{organ} : {np.mean(predictions):.1f} mGy')
    return np.mean(predictions)


organs_list = ['esophagus', 'aorta', 'lung_lower_lobe_left', 'lung_lower_lobe_right',
                'lung_upper_lobe_left', 'lung_upper_lobe_right', 'lung_middle_lobe_right', 
                 'skin', 'lung_vessels', 'lung_trachea_bronchia']
                 

organ_dose_est_list = []
for organ in organs_list:
    organ_dose_estimate  = evaluate(organ)
    organ_dose_est_list.append(organ_dose_estimate)