"""
 Copyright 2018 Johns Hopkins University  (Author: Jesus Villalba)
 Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
"""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division
from six.moves import xrange

import logging
import numpy as np
import h5py

import scipy.linalg as la
from scipy.special import erfinv

from ..hyp_defs import float_cpu
from ..hyp_model import HypModel


class Gaussianizer(HypModel):
    """Class to make i-vector distribution standard Normal.
    """
    def __init__(self, max_vectors=None, r=None, **kwargs):
        super(Gaussianizer, self).__init__(**kwargs)
        self.max_vectors = max_vectors
        self.r = r

        

    def predict(self, x):
        px_cum = np.linspace(0, 1, self.r.shape[0]+2)[1:-1]
        y_map = erfinv(2*px_cum-1)*np.sqrt(2)

        r=self.r[1:]
        y = np.zeros_like(x)
        for i in xrange(x.shape[1]):
            y_index = np.searchsorted(r[:,i], x[:,i])
            logging.debug(y_index)
            y[:,i] = y_map[y_index]
            
        return y


    
    def fit(self, x):

        r = np.sort(x, axis=0, kind='heapsort')

        x = np.zeros((1, x.shape[-1]), dtype=float_cpu())

        if r.shape[0] > self.max_vectors:
            index = np.round(np.linspace(0, r.shape[0]-1, self.max_vectors, dtype=float)).astype(int)
            r = r[index, :]

        self.r = np.vstack((x, r))
        


    def get_config(self):
        config = {'max_vectors': self.max_vectors}

        base_config = super(Gaussianizer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


        
    def save_params(self, f):
        params = {'r': self.r }
        self._save_params_from_dict(f, params)


        
    @classmethod
    def load_params(cls, f, config):
        param_list = ['r']
        params = cls._load_params_to_dict(f, config['name'], param_list)
        return cls(r=params['r'], max_vectors=config['max_vectors'], name=config['name'])

    
    
    @classmethod
    def load_mat(cls, file_path):
        with h5py.File(file_path, 'r') as f:
            r = np.asarray(f['r'], dtype='float32')
            return cls(r=r)

        
        
    def save_mat(self, file_path):
        with h5py.File(file_path, 'w') as f:
            f.create_dataset('r', data=self.r)



    @staticmethod
    def filter_args(prefix=None, **kwargs):
        if prefix is None:
            p = ''
        else:
            p = prefix + '_'
            
        valid_args = ('max_vectors', 'name')
        return dict((k, kwargs[p+k])
                    for k in valid_args if p+k in kwargs)



    @staticmethod
    def add_argparse_args(parser, prefix=None):
        if prefix is None:
            p1 = '--'
            p2 = ''
        else:
            p1 = '--' + prefix + '-'
            p2 = prefix + '_'
            
        parser.add_argument(p1+'max-vectors', dest=(p2+'max_vectors'), default=None,
                            type=int,
                            help=('maximum number of background vectors'))

        parser.add_argument('--name', dest='name', default='gauss')
