import os
import sys
import json
import math
import numpy as np


from mountaintop import loggerx


def update(means, variances, count):
    for i in range(len(means)):
        means[i] /= count
        tmp_v = variances[i] / count - means[i] * means[i]
        tmp_v = max(tmp_v, 1.0e-20)
        variances[i] = 1.0 / math.sqrt(tmp_v)
    means = np.array(means)
    variances = np.array(variances)
    return means, variances


def load_json_cmvn(json_cmvn_file):
    """ Load the json format cmvn stats file and calculate cmvn

    Args:
        json_cmvn_file: cmvn stats file in json format

    Returns:
        a numpy array of [means, variances]
    """
    with open(json_cmvn_file) as f:
        cmvn_stats = json.load(f)

    means = cmvn_stats['mean_stats']
    variance = cmvn_stats['var_stats']
    count = cmvn_stats['num_frames']
    return update(means, variance, count)


def load_kaldi_cmvn(kaldi_cmvn_file):
    """ Load the kaldi format cmvn stats file and calculate cmvn

    Args:
        kaldi_cmvn_file:  kaldi text style global cmvn file, which
           is generated by:
           compute-cmvn-stats --binary=false scp:feats.scp global_cmvn

    Returns:
        a numpy array of [means, variances]
    """
    means = []
    variance = []
    with open(kaldi_cmvn_file, 'r') as fid:
        # kaldi binary file start with '\0B'
        if fid.read(2) == '\0B':
            loggerx.error('kaldi cmvn binary file is not supported, please '
                          'recompute it by: compute-cmvn-stats --binary=false '
                          ' scp:feats.scp global_cmvn')
            sys.exit(1)
        fid.seek(0)
        arr = fid.read().split()
        assert (arr[0] == '[')
        assert (arr[-2:] == '0]')
        feat_dim = int((len(arr) - 4) / 2)
        
        means = map(float, arr[1 : feat_dim+1])
        count = float(arr[feat_dim+1])
        variance = map(float, arr[feat_dim+2 : 2*feat_dim+2])
        
    return update(means, variance, count)


def load_cmvn(cmvn_file, cmvn_type="json"):
    means, variance = None, None
    if not os.path.exists(cmvn_file):
        loggerx.warning("CMVN file (%s) not exists!"%(str(cmvn_file)))
        return means, variance
    if cmvn_type == "kaldi":
        means, variance = load_kaldi_cmvn(cmvn_file)
    elif cmvn_type == "json":
        means, variance = load_json_cmvn(cmvn_file)
    
    return means, variance

