import os
import cv2
from glob import glob

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import PIL

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
import torchvision
from torchvision import transforms
"""
# Preprocess
"""
labels = pd.read_csv ('../input/data/Data_Entry_2017.csv')
labels.sample (3)
pathology_list = ['Cardiomegaly', 'Emphysema', 'Effusion', 'Hernia', 'Nodule', 'Pneumothorax', 'Atelectasis',
                  'Pleural_Thickening', 'Mass','Edema', 'Consolidation', 'Infiltration', 'Fibrosis', 'Pneumonia']

# multi hot encoding
for pathology in pathology_list :
    labels [pathology] = labels ['Finding Labels'].apply (lambda x: 1 if pathology in x else 0)

labels ['No Findings'] = labels ['Finding Labels'].apply (lambda x: 1 if 'No Finding' in x else 0)
labels.sample (3)
labels = labels.drop (list (labels.iloc [:,1:12].columns.values), axis = 1) #12
labels.head (3)
# removing No Finding patient images and its label
indexes = []
for i in range (len (labels)):
    if labels.iloc [i, 15] == 1:
        indexes.append (i)
# remove indexes
labels = labels.drop (indexes, axis = 0)
# remove column
labels = labels.drop(['No Findings'],axis = 1)
# reset indexes
labels = labels.reset_index ()

print ('number of patients:', len (labels))
labels.head (3)
# removing multi label patients
idxs = []
for i in range (len (labels)):
    if len (np.where (labels.iloc [i] == 1)[0]) > 1:
        idxs.append (i)
labels = labels.drop (idxs, axis = 0)
labels = labels.reset_index ()

print ('number of unique patients:', len (labels))
labels.head (3)
# adding path to dataframe
all_image_paths = {os.path.basename(x): x for x in glob ('../input/data/images_*/images/*.png')}
labels ['path'] = labels ['Image Index'].map(all_image_paths.get)
labels.head (3)
# adding class idx to dataframe
idx = []
for i in range (len (labels)):
    idx.append (np.where (labels.iloc [i, 3:] == 1)[0].item ())
    
labels ['class_idx'] = idx
labels.head (3)
# remove additional columns
labels = labels.drop (list (labels.iloc [:,:17].columns.values), axis = 1)
labels.head (3)
# ploting distribution of classes
counts = labels.class_idx.value_counts ()
class_idx = counts.index

plt.bar (class_idx, counts)
# img_dir = '../input/sample/sample/sample/images'
# print ('number of images:', len (os.listdir (img_dir)))
img_channels = 1
# from WGAN paper
z_dim = 100
batch_size = 64
img_size = 64
# features_c = 16
# features_g = 16
# lr = 5e-5 for WGAN
# lr = 1e-4 for WPGAN
# num_classes = 10
# gen_embedding = 100
# could also use two lrs, one for gen and one for critic
# critic_iterations = 5
# weight_clip = 0.01 for WGAN
# lambda_gp = 10 for WPGAN
class lungDataset (Dataset):
    """ create dataset by using image path and their labels """
    def __init__ (self, labels, transform = None):
        self.labels = labels
        self.transform = transform
        
    def __len__ (self):
        return len (self.labels)
    
    def __getitem__ (self, index):
        image_path = self.labels.iloc [index][0]
        image = PIL.Image.open (image_path).convert ('L')
        
        label = self.labels.iloc [index][1]

        if self.transform is not None:
            image = self.transform (image)
            
        return image, label
# class lungDataset (Dataset):
#     """Create Unsupervised Dataset from sample NIH"""
#     def __init__ (self, image_dir, transform = None):
#         self.image_dir = image_dir
#         self.images = os.listdir (image_dir)
#         self.transform = transform
        
#     def __getitem__ (self, index):
#         image_path = self.image_dir + '/' + self.images [index]
#         image = PIL.Image.open (image_path).convert ('L')
        
#         if self.transform is not None:
#             image = self.transform (image)
            
#         return image
    
#     def __len__ (self):
#         return len (self.images)
tfms = transforms.Compose ([
    transforms.Resize ((64, 64)),
    transforms.ToTensor (),
    transforms.Normalize (
        [0.5 for _ in range (img_channels)], [0.5 for _ in range (img_channels)]
    )
])
# Create dataset and dataloader
dataset = lungDataset (labels, transform = tfms)
dataloader = DataLoader (dataset, batch_size = batch_size, shuffle = True)
"""
# Visualize
"""
def imshow (img, title = None):
    """ a function to show tensor images """
    img = img.numpy ().transpose (1, 2, 0)
    mean = 0.5
    std = 0.5
    img = img * std + mean
    img = np.clip (img, 0, 1)
    
    plt.figure (figsize = (10, 8))
    plt.axis ('off')
    plt.imshow (img)
    if title:
        plt.title (title)
# Visualize random batch of dataset
img, lbl = next (iter (dataloader))
print (img.shape)
print (lbl.shape)
img = torchvision.utils.make_grid (img, nrow = 8)
imshow (img, 'Random Batch of Images')
"""
# Models
"""
class Critic (nn.Module):
    """ a model that judges between real and fake images """
    def __init__ (self, img_channels, features_c, num_classes = 14, img_size = 64):
        super (Critic, self).__init__ ()
        
        # Input: N x channels_img x 64 x 64
        self.critic = nn.Sequential (
            nn.Conv2d (img_channels + 1, features_c, kernel_size = 4, stride = 2, padding = 1), # 32x32
            nn.LeakyReLU (0.2),
            self._block (features_c , features_c * 2, 4, 2, 1),                             # 16x16
            self._block (features_c * 2, features_c * 4, 4, 2, 1),                          # 8x8
            self._block (features_c * 4, features_c * 8, 4, 2, 1),                          # 4x4
            nn.Conv2d (features_c * 8, 1, kernel_size = 4, stride = 2, padding = 0),        # 1x1
        )
        self.embed = nn.Embedding (num_classes, img_size * img_size)
        self.img_size = img_size
      
    def _block (self, in_channels, out_channels, kernel_size, stride, padding):
        return nn.Sequential (
            nn.Conv2d (in_channels, out_channels, kernel_size, stride, padding, bias = False),
            # for WGAN
            # nn.BatchNorm2d (out_channels),
            # for WPGAN
            nn.InstanceNorm2d (out_channels, affine = True),
            nn.LeakyReLU (0.2)
        )
    
    def forward (self, x, labels):
        embedding = self.embed (labels).view (labels.shape [0], 1, self.img_size, self.img_size)
        x = torch.cat ([x, embedding], dim = 1) # N x C x img_size (H) x img_size (W)
        return self.critic (x)
class Generator (nn.Module):
    """ a model generates fake images """
    def __init__ (self, z_dim, img_channels, features_g, embed_size, img_size = 64, num_classes = 14):
        super (Generator, self).__init__ ()
        self.img_size = img_size
        # Input: N x z_dim x 1 x 1
        self.gen = nn.Sequential (
            self._block (z_dim + embed_size, features_g * 16, 4, 2, 0),                               # 4x4
            self._block (features_g * 16, features_g * 8, 4, 2, 1),                      # 8x8
            self._block (features_g * 8, features_g * 4, 4, 2, 1),                       # 16x16
            self._block (features_g * 4, features_g * 2, 4, 2, 1),                       # 32x32
            nn.ConvTranspose2d (
                features_g * 2, img_channels, kernel_size = 4, stride = 2, padding = 1), # 64x64
            nn.Tanh ()
        )
        self.embed = nn.Embedding (num_classes, embed_size)
        
    def _block (self, in_channels, out_channels, kernel_size, stride, padding):
        return nn.Sequential (
            nn.ConvTranspose2d (in_channels, out_channels, kernel_size, stride, padding, bias = False),
            nn.BatchNorm2d (out_channels),
            nn.ReLU ()
        )
    
    def forward (self, x, labels):
        # latent vector z: N x noise_dim x 1 x 1
        embedding = self.embed (labels).unsqueeze (2).unsqueeze (3)
        x = torch.cat ([x, embedding], dim = 1)
        return self.gen (x)
def initialize_weights (model):
    for m in model.modules ():
        if isinstance (m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):
            nn.init.normal_ (m.weight.data, 0.0, 0.02)
def gradient_penalty (critic, images, labels, fake_images, device):
    batch_size, C, W, H = images.shape
    epsilon = torch.rand ((batch_size, 1, 1, 1)).repeat (1, C, H, W).to(device)
    interpolated_images = images * epsilon + fake_images * (1 - epsilon)
    
    # calculate critic scores
    mixed_scores = critic (interpolated_images, labels)
    
    gradient = torch.autograd.grad (
        inputs = interpolated_images,
        outputs = mixed_scores,
        grad_outputs = torch.ones_like (mixed_scores),
        create_graph = True,
        retain_graph = True,
    )[0]
    
    gradient = gradient.view (gradient.shape [0], -1)
    gradient_norm = gradient.norm (2, dim = 1)
    gradient_penalty = torch.mean ((gradient_norm - 1) ** 2)
    
    return gradient_penalty
def test ():
    """ test function to test models """
    N, in_channels, H, W = 8, 1, 64, 64
    x = torch.randn ((N, in_channels, H, W))
    y = torch.randint (0, 9, (8,))
    z_dim = 100
    C = Critic(in_channels, 8)
    initialize_weights (C)
    assert C (x, y).shape == (N, 1, 1, 1)
    G = Generator (z_dim, in_channels, 8, 100)
    initialize_weights (G)
    z = torch.randn ((N, z_dim, 1, 1))
    assert G(z, y).shape == (N, in_channels, H, W)
    print ('Succes')
    
test ()
"""
# Training
"""
# Hyper Function and Initializiation
device = torch.device ('cuda' if torch.cuda.is_available () else 'cpu')

critic = Critic(img_channels = 1, features_c = 64).to(device)
gen = Generator (z_dim = 100, img_channels = 1, features_g = 64, embed_size = 100).to(device)

initialize_weights (critic)
initialize_weights (gen)

# WPGAN
critic_optimizer = torch.optim.Adam (critic.parameters (), lr = 1e-4, betas = (0.0, 0.9))
gen_optimizer = torch.optim.Adam (gen.parameters (), lr = 1e-4, betas = (0.0, 0.9))

# WGAN
# critic_optimizer = torch.optim.RMSprop (critic.parameters (), lr = 5e-5)
# gen_optimizer = torch.optim.RMSprop (gen.parameters (), lr = 5e-5)
def denorm (x):
    out = x * 0.5 + 0.5
    return out.clamp (0, 1)
def train_critic (critic, images, labels, critic_iteration, device):
    critic.train ()
    
    batch_size = images.shape [0]
    images = images.to(device)
    labels = labels.to(device)
    
    for _  in range (critic_iteration):
        z = torch.randn (batch_size, z_dim, 1, 1).to(device)
        fake_images = gen (z, labels).to(device)
        
        # scores
        real_score = critic (images, labels).reshape (-1)
        fake_score = critic (fake_images, labels).reshape (-1)
        
        # gradiant penalty
        gp = gradient_penalty (critic, images, labels, fake_images, device)
        
        # calculate loss
        critic_loss = -(torch.mean (real_score) - torch.mean (fake_score)) + 10 * gp

        # Reset gradients
        critic.zero_grad ()

        # Compute gradients
        critic_loss.backward (retain_graph = True)

        # Adjust the parameters using backprop
        critic_optimizer.step ()
        
        # for WGAN
#         for p in critic.parameters ():
#             p.data.clamp_ (-0.01, 0.01)
    
    return fake_images, labels, critic_loss, real_score, fake_score
def train_generator (gen, critic, fake_images, labels):
    gen.train ()

    # calculate generator score and loss
    output = critic (fake_images, labels).reshape (-1)
    gen_loss = -torch.mean (output)
    
    # Backprop and optimize
    gen.zero_grad ()
    gen_loss.backward ()
    gen_optimizer.step ()
    
    return gen_loss
from IPython.display import Image
from torchvision.utils import save_image

def save_fake_images (index, sample_dir = 'samples'):
    if not os.path.exists (sample_dir):
        os.makedirs (sample_dir)
        
    fixed_noise = torch.randn (25, z_dim, 1, 1).to(device)
    labels = torch.randint (0, 13, (25,)).to(device)
    fake = gen (fixed_noise, labels)
    
    fake_fname = 'fake_images-{0:0=4d}.png'.format (index)
    print ('Svaing', fake_fname)
    save_image (denorm (fake), os.path.join (sample_dir, fake_fname), nrow = 5)
# Befor training
sample_dir = 'samples'
save_fake_images (0, sample_dir)
Image (os.path.join (sample_dir, 'fake_images-0000.png'))
"""
# Train
"""
def fit (critic, gen, dataloader, num_epochs, critic_iteration, device, save_fn = None):
    
    writer_real = SummaryWriter (f'logs/real')
    writer_fake = SummaryWriter (f'logs/fake')
    step = 0
    
    d_losses, g_losses, real_scores, fake_scores = [], [], [], []
    total_step = len (dataloader)
    
    for epoch in range (num_epochs):
        for i, (images, labels) in enumerate (dataloader):

            # Train the descriminator and generator
            fake_images, labels, critic_loss, real_score, fake_score = train_critic (critic, images, labels, critic_iteration, device)
            gen_loss = train_generator (gen, critic, fake_images, labels)

            # Inspect the losses
            if (i + 1) % 100 == 0:
                d_losses.append (critic_loss.item ())
                g_losses.append (gen_loss.item ())
                real_scores.append (real_score.mean ().item ())
                fake_scores.append (fake_score.mean ().item ())
                print ('Epoch [{}/{}], Step [{}/{}], critic_loss: {:.4f}, gen_loss: {:.4f}, D(x): {:.2f}, D (G(z)): {:.2f}'
                .format (epoch + 1, num_epochs, i + 1, total_step, critic_loss.item (), gen_loss.item (), real_score.mean ().item (), fake_score.mean ().item ()))
                
                with torch.no_grad ():
                    
                    img_grid_real = torchvision.utils.make_grid (images [:16], normalize = True)
                    img_grid_fake = torchvision.utils.make_grid (fake_images [:16], normalize = True)
                    
                    writer_real.add_image ('Real', img_grid_real, global_step = step)
                    writer_real.add_image ('Fake', img_grid_fake, global_step = step)
                    
                step += 1
                    

        # Sample and save images
        if save_fn is not None:
            save_fn (epoch + 1)
        
    return dict (critic_loss = d_losses, gen_loss = g_losses, real_score = real_scores, fake_score = fake_scores)
history = fit (critic, gen, dataloader, 29, 5, device, save_fake_images)
# Save the model checkpoints
torch.save (gen.state_dict (), 'NIHGen.pth')
torch.save (critic.state_dict (), 'NIHCrit.pth')
"""
## Create training video
"""
from IPython.display import FileLink

vid_fname = 'lung_wpgan_training.avi'

files = [os.path.join (sample_dir, f) for f in os.listdir (sample_dir) if 'fake_images' in f]
files.sort ()
files
# create a video that show training results
out = cv2.VideoWriter (vid_fname, cv2.VideoWriter_fourcc (*'FMP4'), 8, (332, 332))
[out.write (cv2.imread (fname)) for fname in files]
out.release ()
FileLink ('lung_wpgan_training.avi')
"""
## ploting results
"""
plt.plot (history ['critic_loss'], '-')
plt.plot (history ['gen_loss'], '-')
plt.xlabel ('epoch')
plt.ylabel ('loss')
plt.legend (['Discriminator', 'Generator'])
plt.title ('Losses');
plt.plot (history ['real_score'], '-')
plt.plot (history ['fake_score'], '-')
plt.xlabel ('epoch')
plt.ylabel ('score')
plt.legend (['Real Score', 'Fake Score'])
plt.title ('Scores');
"""
## show gernerated images
"""
fixed_noise = torch.randn (16, 100, 1, 1).to(device)
labels = torch.randint (0, 13, (16,)).to(device)
fake = gen (fixed_noise, labels)
imgs = denorm (fake).cpu ().detach ()

plt.figure (figsize = (20, 20))
for i in range (16):
    plt.subplot (4, 4, i + 1)
    plt.imshow (imgs [i][0], cmap = 'gray')
    plt.title (pathology_list [labels [i]])
    
plt.show ()