# coding: utf-8

import gzip
import os
import shutil

import numpy as np
import paddle.v2 as paddle


def extract_param(path):
    """
    Return the parameters of weight and bias
    """
    with gzip.open(path) as f:
        p = paddle.parameters.Parameters.from_tar(f)
    names = p.names()
    weight, bias = [], []
    for name in names:
        if name.endswith('bias'):
            bias.append(p.get(name))
        else:
            weight.append(p.get(name))
    return weight, bias


def param2header(param, path, dtype, var_name):
    """
    Trans the list of 2d np array into C++ header file
    """
    with open(path, 'w') as f:
        postfix = 1
        for matrix in param:
            rows = len(matrix)
            cols = len(matrix[0])
            beg = 'static {0} {1}{2}[{3}]={{'.format(dtype, var_name, postfix, rows * cols)
            end = '};\n'
            f.write(beg)
            for i in range(rows):
                for j in range(cols):
                    if i == 0 and j == 0:
                        f.write('{:.6f}'.format(matrix[i][j]))
                    else:
                        f.write(',{:.6f}'.format(matrix[i][j]))
                f.write('\n')
            f.write(end)
            postfix += 1
        # layers = len(param)
        # postfix = 1
        # beg = 'static {0} **{1} = {0}*[{2}]{{{1}{3}'.format(dtype, var_name, layers, postfix)
        # f.write(beg)
        # while postfix < layers:
        #     postfix += 1
        #     f.write(', {0}{1}'.format(var_name, postfix))
        # f.write(end)


def cp_gtsrb_img(src, dst):
    """
    Copy all GTSRB images from src to dst, just like:
    src/000xx/000yy_000zz.ppm -> dst/000xx_000yy_000zz.ppm
    """
    folders = os.listdir(src)
    for folder in folders:
        imgs = filter(lambda fname: fname.endswith('ppm'),
                      os.listdir(src + '/' + folder))
        for img in imgs:
            shutil.copy(src + '/' + folder + '/' + img,
                        dst + '/' + folder + '_' + img)


def generate_gtsrb_dict(dir):
    imgs = os.listdir(dir)
    gtsrb_dict = {}
    for img in imgs:
        gtsrb_dict[img] = int(img.split('_')[0])
    return gtsrb_dict


if __name__ == '__main__':
    gtsrb_img_path = 'data/GTSRB/Final_Training/Images'
    gtsrb_cp_dst = 'data/GTSRB/Final_Training/All'

    param_path = 'data/params_pass_1900.tar.gz'
    weight, bias = extract_param(param_path)

    for w in weight:
        print w.shape

    for b in bias:
        print b.shape

    param2header(weight, 'weight.h', 'float', 'weight')
    param2header(bias, 'bias.h', 'float', 'bias')

    # cp_gtsrb_img(gtsrb_img_path, gtsrb_cp_dst)
    # generate_gtsrb_dict(gtsrb_cp_dst)