#! /usr/bin/env python3

import re
import sys

'''
usage: python3 compact_json.py 720 ${FN_JSON_OLD} ${FN_JSON_NEW}
'''

def compact_array(str_array):
    """ function to remove \n \t from given string.
    """
    a = str_array.group().replace('\n', '').replace('\t', '')
    return a

def compact_file(fn_in, fn_out, mode):
    """ re-print the json in a compact mode for better view.
    
    The knerex c++ code will print vector with each element in a line.
    It is difficult for user to read.

    This function will print [a, b, c] format instead.

    TODO: make it work for multiple level bracket [[a. b, c]].

    TODO: knerex 720 still has bug which export key of kdp520.
    This script will fix the bug. but supposed to be removed after knerex 720 bug fixed.
    """
    with open(fn_in, 'r') as  f:
        j = f.read()

    j = re.sub(r'\[.*?\]', compact_array, j, flags=re.DOTALL)
    j = re.sub(r':[ \n\t]*\[', ': [', j, flags=re.DOTALL)

    if mode in ['kdp720', '720']:
        j = re.sub(r'\"\bbn_shift_gamma\b\"', '"pool_shift"', j)
        # j = re.sub(r'\"\badd_coarse_shift\b\"', '"conv_coarse_shift"', j)
        # j = re.sub(r'\"\badd_fine_shift\b\"', '"add_coarse_shift"', j)
        j = re.sub(r'\"\b\S+_weight\b\"', '"weight_radix"', j)
        j = re.sub(r'\"\b\S+_gamma\b\"', '"weight_radix"', j)
        j = re.sub(r'\"\b\S+_bias\b\"', '"bias_radix"', j)
        j = re.sub(r'\"\b\S+_beta\b\"', '"bias_radix"', j)
        j = re.sub(r'\"\bweight_fl\b\"', '"gap_mul_fl"', j)
        j = re.sub(r'\"\bweight_fp\b\"', '"gap_mul_fp"', j)

    with open(fn_out, 'w') as  f:
        f.write(j)

if __name__ == "__main__":
    mode = sys.argv[1]
    fn_in = sys.argv[2]
    if len(sys.argv) < 4:
        fn_out = sys.argv[2]
    else:
        fn_out = sys.argv[3]

    compact_file(fn_in, fn_out, mode)