#!/usr/bin/env python
# coding:utf-8

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
import gzip
import pickle

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import os


class Generator(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Generator, self).__init__()
        self.map1 = nn.Linear(input_size, hidden_size)
        self.map2 = nn.Linear(hidden_size, hidden_size)
        self.map3 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = F.elu(self.map1(x))
        x = F.sigmoid(self.map2(x))
        return self.map3(x)


class Discriminator(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Discriminator, self).__init__()
        self.map1 = nn.Linear(input_size, hidden_size)
        self.map2 = nn.Linear(hidden_size, hidden_size)
        self.map3 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = F.elu(self.map1(x))
        x = F.elu(self.map2(x))
        return F.sigmoid(self.map3(x))


class DataGen(object):
    """docstring for DataGen"""

    def __init__(self, idx=-1):
        super(DataGen, self).__init__()
        f = gzip.open("../../../datasets/mnist.pkl.gz", "rb")
        (x_train, y_train), (x_test, y_test) = pickle.load(f, encoding="bytes")
        f.close()
        if idx != -1:
            x_train = x_train[y_train == idx]
            y_train = y_train[y_train == idx]
        self.x_train, self.y_train = x_train, y_train
        self.num = x_train.shape[0]
        self.idx = 0
        self.x_dim = 28 * 28

    def next_batch(self, mb_size):
        st = self.idx * mb_size
        et = (self.idx + 1) * mb_size
        if et > self.num:
            self.idx = 0
            return self.next_batch(mb_size)
        else:
            self.idx += 1
        a = np.array(self.x_train[st:et], dtype="float32")
        a = a.reshape(-1, self.x_dim)
        x = Variable(torch.from_numpy(a))
        return x


def train():
    mb_size = 64
    z_dim = 100
    x_dim = 28 * 28
    h_dim = 128
    lr = 1e-3
    d_steps = 5
    top_g_steps = 20000
    num_epochs = 1000000
    data = DataGen(4)
    path = "output4"

    # define GAN loss
    G = Generator(input_size=z_dim, hidden_size=h_dim, output_size=x_dim)
    D = Discriminator(input_size=x_dim, hidden_size=h_dim, output_size=1)

    D_solver = optim.Adam(D.parameters(), lr=lr, betas=(0.9, 0.999))
    G_solver = optim.Adam(G.parameters(), lr=lr, betas=(0.9, 0.999))

    ones_label = Variable(torch.ones(mb_size))
    zeros_label = Variable(torch.zeros(mb_size))
    c = 0
    for epoch in range(num_epochs):
        # 优化生成网络
        for d_index in range(d_steps):
            D.zero_grad()
            # 训练真实的样本
            x = data.next_batch(mb_size)
            z = Variable(torch.randn(mb_size, z_dim))
            s = G(z)
            D_real = D(x)
            D_fake = D(s)

            D_loss_real = F.binary_cross_entropy(D_real, ones_label)
            D_loss_fake = F.binary_cross_entropy(D_fake, zeros_label)
            D_loss = D_loss_real + D_loss_fake
            D_loss.backward()  # backward只是反馈了, 但是参数仍未更新
            D_solver.step()    # 进行参数更新

        # 优化生成网络
        G.zero_grad()
        G.zero_grad()
        z = Variable(torch.randn(mb_size, z_dim))
        s = G(z)
        D_fake = D(s)
        # 初始化阶段, 对loss少做改变
        if epoch < top_g_steps:
            G_loss = F.binary_cross_entropy(D_fake, zeros_label)
        else:
            G_loss = F.binary_cross_entropy(D_fake, ones_label)
        G_loss.backward()
        G_solver.step()
        # Print and plot every now and then
        if epoch % 1000 == 0:
            loss_d = D_loss.data.numpy()
            loss_g = G_loss.data.numpy()
            print('iter={}, D_loss={}, G_loss={}'.format(epoch, loss_d, loss_g))
            # 随机取生成的前16张进行可视化
            samples = G(z).data.numpy()[:16]

            fig = plt.figure(figsize=(4, 4))
            gs = gridspec.GridSpec(4, 4)
            gs.update(wspace=0.05, hspace=0.05)

            for i, sample in enumerate(samples):
                ax = plt.subplot(gs[i])
                plt.axis('off')
                ax.set_xticklabels([])
                ax.set_yticklabels([])
                ax.set_aspect('equal')
                plt.imshow(sample.reshape(28, 28), cmap='Greys_r')

            if not os.path.exists(path):
                os.makedirs(path)

            plt.savefig('{}/{}.png'.format(path, str(c).zfill(3)),
                        bbox_inches='tight')
            c += 1
            plt.close(fig)

if __name__ == '__main__':
    train()