#!/usr/bin/env python2

import math
import numpy as np
from scipy.misc import imresize
from matplotlib import pyplot as plt

def weight_image(weights, weight_name):
    weights = weights[weight_name]

    w, h, c, cnt = weights.shape
    grid = int(math.ceil(math.sqrt(cnt)))
    grid = min(8, grid)
    cnt = min(cnt, grid*grid)

    fig = plt.figure(figsize=(10, 10))
    for j in range(c):
        for i in range(cnt):
            ax = fig.add_subplot(grid, grid, i + 1)

            ax.tick_params(labelcolor='none', top='off', bottom='off', 
                            left='off', right='off')
            img = weights[:, :, j, i]
            img = imresize(img, (32, 32))
            ax.imshow(img)
        plt.subplots_adjust(wspace=0.001, hspace=0.001)
        plt.savefig(weight_name + '-' + str(j+1) + '.png', bbox_inches='tight')

    print(weight_name, weights.shape)

def weight_hist(weights):
    keys = sorted(weights.keys())

    grid = int(math.ceil(math.sqrt(len(keys))))

    i = 0
    fig = plt.figure(figsize=(14, 8))
    for k in keys:
        i += 1
        w = weights[k].flatten()

        plt.subplot(grid, grid, i)
        plt.title(k)
        plt.hist(w, bins=20)

    plt.tight_layout()
    plt.savefig('weight_hist.png', bbox_inches='tight')

def weight_sum_hist(weights):
    keys = sorted(weights.keys())
    keys = [k for k in keys if 'conv' in k and 'W' in k]

    fig = plt.figure(figsize=(14, 8))

    grid = int(math.ceil(math.sqrt(len(keys))))
    idx = 0
    for k in keys:
        idx += 1

        weight = weights[k]
        weight_sum = []
        w, h, c, cnt = weight.shape
        for i in range(c):
            for j in range(cnt):
                weight_sum.append(weight[:, :, i, j].sum())

        plt.subplot(grid, grid, idx)
        plt.hist(weight_sum, bins=20)
        plt.title(k)

    plt.tight_layout()
    plt.savefig('weight_sum_hist.png', bbox_inches='tight')

if __name__ == '__main__':
    weight_file_name = './vgg16_weights.npz'
    weights = np.load(weight_file_name)

    keys = sorted(weights.keys())

    #for k in keys:
    #    print(k, weights[k].shape)
    
    #weight_image(weights, 'conv1_1_W')
    #visualize_weight(weights, 'conv1_2_W')
    #visualize_weight(weights, 'conv4_1_W')

    #weight_hist(weights)
    weight_sum_hist(weights)
