import torch
import sys
import abganlibs.models.dcgan_gp_fz as dcgan
import numpy as np

# Check the parameter is legal or not
if 1 == len(sys.argv) or ".pth" != sys.argv[1][-4:]:
    saved_model = "saves/WGAN_GP_5/netG_epoch_2100_32.pth"
    #exit('USE: decoder.py MODEL.pth')
else:
    saved_model = sys.argv[1]

# Set number of levels
batch_size = 10

# Set the parameters same as the used in training process
map_size = 20
nz = 32
ngf = 32
ngpu = 1
n_extra_layers = 0
z_dims = 3

threshold = .5

# Load Generator
netG = dcgan.DCGAN_G(map_size, nz, z_dims, ngf, ngpu, n_extra_layers)
netG.load_state_dict(torch.load(saved_model))
noise = torch.Tensor(batch_size, nz, 1, 1).normal_(0, 1)

# Generate
results = netG(noise).data

for i, result in enumerate(results, 1):
    temp = np.zeros((3, map_size*map_size), dtype=float)
    print(result.shape)

    for x in range(map_size):
        for y in range(map_size):
            temp_binary_type = ''
            reliablity_one = []
            reliablity_zero = []

            temp_max = result[:, x, y].max()
            temp_result = result[:, x, y] / temp_max

            for j in temp_result:
                if j > threshold:
                    temp_binary_type += '1'
                    reliablity_one.append(j)
                else:
                    temp_binary_type += '0'
                    reliablity_zero.append(j)

            channel = int(temp_binary_type, 2)
            if channel < 3 and channel >= 0:
                if len(reliablity_one) != 0:
                    one_avg = sum(reliablity_one) / len(reliablity_one)
                else:
                    one_avg = 0

                if len(reliablity_zero) != 0:
                    zero_avg = sum(reliablity_zero) / len(reliablity_zero)
                else:
                    zero_avg = 0

                channel_value = one_avg - zero_avg
                index = x * map_size + y
                temp[channel][index] = channel_value * temp_max

    np.savetxt('data/generated_data_wgan_gp_5_fz/from_net_{0:02d}.txt'.format(i),
               temp, delimiter=",", fmt='%.2f', encoding='utf8')
