import torch as t
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import Adam

import numpy as np
import cv2
import os
from itertools import product

batch_size = 16
img_size = 32
latent_dim = 64
channels = 4
lr = 0.00008 #0.0001
b1, b2 = 0.7, 0.999
n_epochs = 50
n_ch = 128  # Number of channels in initial convolution layers
sample_interval = 30  # Save a generated image every sample_interval number of batches

write_dir = '../images/dcgan_quad_gen_images/'
data_dir = '../data/bub_single_24new/'
#data_dir = '../data/bub_single_24_shifted/'
model_path = '../models/dcgan_quad_model.pt'
os.makedirs(write_dir, exist_ok=True)


def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find("BatchNorm2d") != -1:
        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
        torch.nn.init.constant_(m.bias.data, 0.0)


class Data(Dataset):
    def __init__(self):
        
        def sort_key(im_name):
            return int(im_name[-7])
        
        # Create a list containing sublists of the image names for each event
        self.ims_by_event = []
        im_list = os.listdir(data_dir)
        run_ev_set = set([im[:-8] for im in im_list])
        for run_ev in run_ev_set:
            event_ims = []
            for im in im_list:
                if not '-0.png' in im:
                    continue
                if im[:-8] == run_ev:
                    event_ims.append(im)
            if len(event_ims) > 0 and len(event_ims) < 5:
                event_ims = sorted(event_ims, key=sort_key)
                self.ims_by_event.append(event_ims)
        
        # For each sublist in self.ims_by_event, append 'blank' if the camera did not image the bubble in that event
        for j, event_ims in enumerate(self.ims_by_event):
            if len(event_ims) == 4:
                continue
            cams_with_ims = [int(im_name[-7]) for im_name in event_ims]
            for k in range(4):
                if not k in cams_with_ims:
                    self.ims_by_event[j].insert(k, 'blank')
        
        # Load images. Load a blank image where 'blank' is found
        self.data = []
        blank = np.zeros((img_size, img_size))
        for event_ims in self.ims_by_event:
            loaded_ev_ims = []
            for im_name in event_ims:
                if im_name == 'blank':
                    loaded_ev_ims.append(blank)
                else:
                    im = cv2.imread(data_dir + im_name, cv2.IMREAD_GRAYSCALE)//255
                    if im.shape != (img_size, img_size):
                        loaded_ev_ims = None
                        break
                    loaded_ev_ims.append(im)
            if loaded_ev_ims is not None:
                self.data.append(loaded_ev_ims)
        
        self.data = t.Tensor(self.data)
        
        self.len = self.data.shape[0]

    def __len__(self):
        return self.len

    def __getitem__(self, index):
        return self.data[index]

dataset = Data()

dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

class Generator(nn.Module):
    # Takes as input a random tensor of size (batch_size, latent_dim) generated by noise()
    def __init__(self):
        super(Generator, self).__init__()
        
        self.init_size = img_size//4
        
        self.l1 = nn.Sequential(nn.Linear(latent_dim, n_ch * self.init_size**2))
        
        self.conv_blocks = nn.Sequential(
            nn.BatchNorm2d(n_ch),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(n_ch, n_ch, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(n_ch, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(n_ch, n_ch//2, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(n_ch//2, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(n_ch//2, channels, kernel_size=3, stride=1, padding=1),
            nn.Sigmoid()
            #nn.Tanh()
            )
            
    def forward(self, z):
        out = self.l1(z)
        out = out.view(out.shape[0], n_ch, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img
            

class Discriminator(nn.Module):
    # Takes as input a data image or image generated by Generator
    def __init__(self):
        super(Discriminator, self).__init__()

        def discriminator_block(in_filters, out_filters, bn=True):
            block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]
            if bn:
                block.append(nn.BatchNorm2d(out_filters, 0.8))
            return block

        self.model = nn.Sequential(
            *discriminator_block(channels, n_ch//8, bn=False),
            *discriminator_block(n_ch//8, n_ch//4),
            *discriminator_block(n_ch//4, n_ch//2),
            *discriminator_block(n_ch//2, n_ch),
        )

        # The height and width of downsampled image
        ds_size = img_size // 2 ** 4
        self.adv_layer = nn.Sequential(nn.Linear(n_ch * ds_size ** 2, 1), nn.Sigmoid())

    def forward(self, img):
        out = self.model(img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)
        return validity
        

def generate_noise(name):
    if name == 'Gaussian':
        return t.normal(mean=0, std=1, size=(batch_size, latent_dim))
    return t.rand((batch_size, latent_dim))
    
def save_model():
    print('Saving Saving generator model to ' + model_path)
    t.save(G, model_path)

# Loss function
adversarial_loss = nn.BCELoss()

G = Generator()
D = Discriminator()

# Initialize weights
#G.apply(weights_init_normal)
#D.apply(weights_init_normal)

# Optimizers
optimizer_G = Adam(G.parameters(), lr=lr, betas=(b1, b2))
optimizer_D = Adam(D.parameters(), lr=lr, betas=(b1, b2))
try:
    for epoch in range(n_epochs):
        for i, data_imgs in enumerate(dataloader):
            
            # Configure generator input
            data_shape = data_imgs.shape
            data_imgs = data_imgs.view(data_imgs.shape[0], channels, data_imgs.shape[-2], data_imgs.shape[-1])
            
            # Adversarial ground truths
            valid = t.ones((data_shape[0], 1), requires_grad=False)
            fake = t.zeros((data_shape[0], 1), requires_grad=False)
            
            # ----------------
            # Train Generator
            # ----------------
            
            optimizer_G.zero_grad()
            
            # Sample noise as generator input
            noise = generate_noise('Gaussian')
            
            # Generate a batch of images
            gen_imgs = G(noise)
            gen_imgs = gen_imgs[:data_shape[0]]

            # Loss measures generator's ability to fool the discriminator
            G_loss = adversarial_loss(D(gen_imgs), valid)

            G_loss.backward()
            optimizer_G.step()
            
            # --------------------
            # Train Discriminator
            # --------------------
            
            optimizer_D.zero_grad()
            
            # Measure discriminator's ability to classify real from generated samples
            real_loss = adversarial_loss(D(data_imgs), valid)
            fake_loss = adversarial_loss(D(gen_imgs.detach()), fake)
            D_loss = (real_loss + fake_loss) / 2
            
            D_loss.backward()
            optimizer_D.step()
            
            # ----------------
            # Log Progress
            # ----------------
            
            print(
                "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
                % (epoch, n_epochs, i, len(dataloader), D_loss.item(), G_loss.item())
            )

            batches_done = epoch*len(dataloader) + i
            if batches_done % sample_interval == 0:
                os.makedirs(write_dir + str(batches_done), exist_ok=True)
                for j, gen_im in enumerate(gen_imgs):
                    for k, im in enumerate(gen_im):
                        cv2.imwrite('%s%d/%d_%d.png' % (write_dir, batches_done, j, k), im.detach().numpy()*255)
    
    save_model()
            
except KeyboardInterrupt:
    save_model()
