'''
Author: Dr. John T. Hwang <hwangjt@umich.edu>

This package is distributed under New BSD license.
'''

from __future__ import division

import numpy as np
import scipy.sparse
from six.moves import range
from numbers import Integral

from smt.utils.linear_solvers import get_solver, LinearSolver, VALID_SOLVERS
from smt.utils.line_search import get_line_search_class, LineSearch, VALID_LINE_SEARCHES
from smt.utils.caching import cached_operation
from smt.surrogate_models.surrogate_model import SurrogateModel


class RMTS(SurrogateModel):
    """
    Regularized Minimal-energy Tensor-product Spline interpolant base class for RMTC and RMTB.
    """

    def _initialize(self):
        super(RMTS, self)._initialize()
        declare = self.options.declare
        supports = self.supports

        declare('xlimits', types=np.ndarray,
                desc='Lower/upper bounds in each dimension - ndarray [nx, 2]')
        declare('smoothness', 1.0, types=(Integral, float, tuple, list, np.ndarray),
                desc='Smoothness parameter in each dimension - length nx. None implies uniform')
        declare('regularization_weight', 1e-14, types=(Integral, float),
                desc='Weight of the term penalizing the norm of the spline coefficients.' +
                    ' This is useful as an alternative to energy minimization ' +
                    ' when energy minimization makes the training time too long.')
        declare('energy_weight', 1e-4, types=(Integral, float),
                desc='The weight of the energy minimization terms')
        declare('extrapolate', False, types=bool,
                desc='Whether to perform linear extrapolation for external evaluation points')
        declare('min_energy', True, types=bool,
                desc='Whether to perform energy minimization')
        declare('approx_order', 4, types=Integral,
                desc='Exponent in the approximation term')
        declare('solver', 'krylov', values=VALID_SOLVERS, types=LinearSolver,
                desc='Linear solver')
        declare('derivative_solver', 'krylov', values=VALID_SOLVERS, types=LinearSolver,
                desc='Linear solver used for computing output derivatives (dy_dyt)')
        declare('grad_weight', 0.5, types=(Integral, float),
                desc='Weight on gradient training data')
        declare('solver_tolerance', 1e-12, types=(Integral, float),
                desc='Convergence tolerance for the nonlinear solver')
        declare('nonlinear_maxiter', 10, types=Integral,
                desc='Maximum number of nonlinear solver iterations')
        declare('line_search', 'backtracking', values=VALID_LINE_SEARCHES, types=LineSearch,
                desc='Line search algorithm')
        declare('save_energy_terms', False, types=bool,
                desc='Whether to cache energy terms in the data_dir directory')
        declare('data_dir', None, values=(None,), types=str,
                desc='Directory for loading / saving cached data; None means do not save or load')
        declare('max_print_depth', 5, types=Integral,
                desc='Maximum depth (level of nesting) to print operation descriptions and times')

        supports['training_derivatives'] = True
        supports['derivatives'] = True
        supports['output_derivatives'] = True

    def _setup_hessian(self):
        diag = np.ones(self.num['dof'])
        arange = np.arange(self.num['dof'])
        full_hess = scipy.sparse.csc_matrix((diag, (arange, arange)))
        return full_hess

    def _compute_jac(self, ix1, ix2, x):
        data, rows, cols = self._compute_jac_raw(ix1, ix2, x)
        n = x.shape[0]
        full_jac = scipy.sparse.csc_matrix((data, (rows, cols)), shape=(n, self.num['coeff']))
        if self.full_dof2coeff is not None:
            full_jac = full_jac * self.full_dof2coeff
        return full_jac

    def _compute_approx_terms(self):
        # This computes the approximation terms for the training points.
        # We loop over kx: 0 is for values and kx>0 represents.
        # the 1-based index of the derivative given by the training point data.
        num = self.num
        xlimits = self.options['xlimits']

        full_jac_dict = {}
        for kx in self.training_points[None]:
            xt, yt = self.training_points[None][kx]

            xmin = np.min(xt, axis=0)
            xmax = np.max(xt, axis=0)
            assert np.all(xlimits[:, 0] <= xmin), 'Training points below min for %s' % kx
            assert np.all(xlimits[:, 1] >= xmax), 'Training points above max for %s' % kx

            if kx == 0:
                c = 1.0
            else:
                self.options['grad_weight'] / xlimits.shape[0]

            full_jac = self._compute_jac(kx, 0, xt)
            full_jac_dict[kx] = (full_jac, full_jac.T.tocsc(), c)

        return full_jac_dict

    def _compute_energy_terms(self):
        # This computes the energy terms that are to be minimized.
        # The quadrature points are the centroids of the multi-dimensional elements.
        num = self.num
        xlimits = self.options['xlimits']

        inputs = {}
        inputs['nx'] = xlimits.shape[0]
        inputs['elem_list'] = num['elem_list']
        if self.__class__.__name__ == 'RMTB':
            inputs['num_ctrl_list'] = num['ctrl_list']
            inputs['order_list'] = num['order_list']

        if self.options['save_energy_terms']:
            cache_dir = self.options['data_dir']
        else:
            cache_dir = None
        with cached_operation(inputs, cache_dir) as outputs:
            if outputs:
                sq_mtx = outputs['sq_mtx']
            else:
                n = np.prod(2 * num['elem_list'])
                x = np.empty(n * num['x'])
                self.rmtsc.compute_quadrature_points(
                    n, np.array(2 * num['elem_list'], dtype=np.int32), x)
                x = x.reshape((n, num['x']))

                sq_mtx = [None] * num['x']
                for kx in range(num['x']):
                    mtx = self._compute_jac(kx+1, kx+1, x)
                    sq_mtx[kx] = mtx.T.tocsc() * mtx * (xlimits[kx, 1] - xlimits[kx, 0]) ** 4

                outputs['sq_mtx'] = sq_mtx

        elem_vol = np.prod((xlimits[:, 1] - xlimits[:, 0]) / (2 * num['elem_list']))
        total_vol = np.prod(xlimits[:, 1] - xlimits[:, 0])

        full_hess = scipy.sparse.csc_matrix((num['dof'], num['dof']))
        for kx in range(num['x']):
            full_hess += sq_mtx[kx] * (elem_vol / total_vol * self.options['smoothness'][kx] \
                / (xlimits[kx, 1] - xlimits[kx, 0]) ** 4)

        return full_hess

    def _opt_func(self, sol, p, yt_dict):
        full_hess = self.full_hess
        full_jac_dict = self.full_jac_dict

        func = 0.5 * np.dot(sol, full_hess * sol)
        for kx in self.training_points[None]:
            full_jac, full_jac_T, c = full_jac_dict[kx]
            yt = yt_dict[kx]
            func += 0.5 * c * np.sum((full_jac * sol - yt) ** p)

        return func

    def _opt_grad(self, sol, p, yt_dict):
        full_hess = self.full_hess
        full_jac_dict = self.full_jac_dict

        grad = full_hess * sol
        for kx in self.training_points[None]:
            full_jac, full_jac_T, c = full_jac_dict[kx]
            yt = yt_dict[kx]
            grad += 0.5 * c * full_jac_T * p * (full_jac * sol - yt) ** (p - 1)

        return grad

    def _opt_dgrad_dyt(self, sol, p, yt_dict, kx):
        full_hess = self.full_hess
        full_jac_dict = self.full_jac_dict

        full_jac, full_jac_T, c = full_jac_dict[kx]
        yt = yt_dict[kx]

        diag_vec = p * (p - 1) * (full_jac * sol - yt) ** (p - 2)
        diag_mtx = scipy.sparse.diags(diag_vec, format='csc')
        mtx = 0.5 * c * full_jac_T.dot(diag_mtx)

        return -mtx.todense()

    def _opt_hess(self, sol, p, yt_dict):
        full_hess = self.full_hess
        full_jac_dict = self.full_jac_dict

        hess = scipy.sparse.csc_matrix(full_hess)
        for kx in self.training_points[None]:
            full_jac, full_jac_T, c = full_jac_dict[kx]
            yt = yt_dict[kx]

            diag_vec = p * (p - 1) * (full_jac * sol - yt) ** (p - 2)
            diag_mtx = scipy.sparse.diags(diag_vec, format='csc')
            hess += 0.5 * c * full_jac_T * diag_mtx * full_jac

        return hess

    def _opt_norm(self, sol, p, yt_dict):
        full_hess = self.full_hess
        full_jac_dict = self.full_jac_dict

        grad = self._opt_grad(sol, p, yt_dict)
        return np.linalg.norm(grad)

    def _get_yt_dict(self, ind_y):
        yt_dict = {}
        for kx in self.training_points[None]:
            xt, yt = self.training_points[None][kx]
            yt_dict[kx] = yt[:, ind_y]
        return yt_dict

    def _run_newton_solver(self, sol):
        num = self.num
        options = self.options

        solver = get_solver(options['solver'])
        ls_class = get_line_search_class(options['line_search'])

        total_size = int(num['dof'])
        rhs = np.zeros((total_size, num['y']))
        d_sol = np.zeros((total_size, num['y']))

        p = self.options['approx_order']
        for ind_y in range(rhs.shape[1]):

            with self.printer._timed_context('Solving for output %i' % ind_y):

                yt_dict = self._get_yt_dict(ind_y)

                norm = self._opt_norm(sol[:, ind_y], p, yt_dict)
                fval = self._opt_func(sol[:, ind_y], p, yt_dict)
                self.printer(
                    'Iteration (num., iy, grad. norm, func.) : %3i %3i %15.9e %15.9e'
                    % (0, ind_y, norm, fval))

                iter_count = 0
                while iter_count < options['nonlinear_maxiter'] and norm > options['solver_tolerance']:
                    with self.printer._timed_context():
                        with self.printer._timed_context('Assembling linear system'):
                            mtx = self._opt_hess(sol[:, ind_y], p, yt_dict)
                            rhs[:, ind_y] = -self._opt_grad(sol[:, ind_y], p, yt_dict)

                        with self.printer._timed_context('Initializing linear solver'):
                            solver._setup(mtx, self.printer)

                        with self.printer._timed_context('Solving linear system'):
                            solver._solve(rhs[:, ind_y], d_sol[:, ind_y], ind_y=ind_y)

                        func = lambda x: self._opt_func(x, p, yt_dict)
                        grad = lambda x: self._opt_grad(x, p, yt_dict)

                        # sol[:, ind_y] += d_sol[:, ind_y]

                        ls = ls_class(sol[:, ind_y], d_sol[:, ind_y], func, grad)
                        with self.printer._timed_context('Performing line search'):
                            sol[:, ind_y] = ls(1.0)

                    norm = self._opt_norm(sol[:, ind_y], p, yt_dict)
                    fval = self._opt_func(sol[:, ind_y], p, yt_dict)
                    self.printer(
                        'Iteration (num., iy, grad. norm, func.) : %3i %3i %15.9e %15.9e'
                        % (iter_count, ind_y, norm, fval))

                    self.mtx = mtx

                    iter_count += 1

    def _solve(self):
        num = self.num
        options = self.options

        solver = get_solver(options['solver'])
        ls_class = get_line_search_class(options['line_search'])

        total_size = int(num['dof'])
        rhs = np.zeros((total_size, num['y']))
        sol = np.zeros((total_size, num['y']))
        d_sol = np.zeros((total_size, num['y']))

        with self.printer._timed_context('Solving initial startup problem (n=%i)' % total_size):

            approx_order = options['approx_order']
            nonlinear_maxiter = options['nonlinear_maxiter']
            options['approx_order'] = 2
            options['nonlinear_maxiter'] = 1

            self._run_newton_solver(sol)

            options['approx_order'] = approx_order
            options['nonlinear_maxiter'] = nonlinear_maxiter

        with self.printer._timed_context('Solving nonlinear problem (n=%i)' % total_size):

            self._run_newton_solver(sol)

        return sol

    def _new_train(self):
        """
        Train the model
        """
        with self.printer._timed_context('Pre-computing matrices', 'assembly'):

            with self.printer._timed_context('Computing dof2coeff', 'dof2coeff'):
                self.full_dof2coeff = self._compute_dof2coeff()

            with self.printer._timed_context('Initializing Hessian', 'init_hess'):
                self.full_hess = self._setup_hessian() * self.options['regularization_weight']

            if self.options['min_energy']:
                with self.printer._timed_context('Computing energy terms', 'energy'):
                    self.full_hess += self._compute_energy_terms() * self.options['energy_weight']

            with self.printer._timed_context('Computing approximation terms', 'approx'):
                self.full_jac_dict = self._compute_approx_terms()

        with self.printer._timed_context('Solving for degrees of freedom', 'total_solution'):
            self.sol = self._solve()

        if self.full_dof2coeff is not None:
            self.sol_coeff = self.full_dof2coeff * self.sol
        else:
            self.sol_coeff = self.sol

    def _train(self):
        """
        Train the model
        """
        self._setup()

        tmp = self.rmtsc
        self.rmtsc = None

        inputs = {'self': self}
        with cached_operation(inputs, self.options['data_dir']) as outputs:
            self.rmtsc = tmp

            if outputs:
                self.sol_coeff = outputs['sol_coeff']
                self.sol = outputs['sol']
                self.mtx = outputs['mtx']
                self.full_dof2coeff = outputs['full_dof2coeff']
                self.full_hess = outputs['full_hess']
                self.full_jac_dict = outputs['full_jac_dict']
            else:
                self._new_train()
                outputs['sol_coeff'] = self.sol_coeff
                outputs['sol'] = self.sol
                outputs['mtx'] = self.mtx
                outputs['full_dof2coeff'] = self.full_dof2coeff
                outputs['full_hess'] = self.full_hess
                outputs['full_jac_dict'] = self.full_jac_dict

    def _predict_values(self, x):
        """
        Evaluates the model at a set of points.

        Arguments
        ---------
        x : np.ndarray [n_evals, dim]
            Evaluation point input variable values

        Returns
        -------
        y : np.ndarray
            Evaluation point output variable values
        """
        mtx = self._compute_prediction_mtx(x, 0)
        y = mtx.dot(self.sol_coeff)

        return y

    def _predict_derivatives(self, x, kx):
        """
        Evaluates the derivatives at a set of points.

        Arguments
        ---------
        x : np.ndarray [n_evals, dim]
            Evaluation point input variable values
        kx : int
            The 0-based index of the input variable with respect to which derivatives are desired.

        Returns
        -------
        y : np.ndarray
            Derivative values.
        """
        mtx = self._compute_prediction_mtx(x, kx + 1)
        y = mtx.dot(self.sol_coeff)

        return y

    def _compute_prediction_mtx(self, x, kx):
        n = x.shape[0]

        num = self.num
        options = self.options

        data, rows, cols = self._compute_jac_raw(kx, 0, x)

        # In the explanation below, n is the number of dimensions, and
        # a_k and b_k are the lower and upper bounds for x_k.
        #
        # A C1 extrapolation can get very tricky, so we implement a simple C0
        # extrapolation. We basically linarly extrapolate from the nearest
        # domain point. For example, if n = 4 and x2 > b2 and x3 > b3:
        #    f(x1,x2,x3,x4) = f(x1,b2,b3,x4) + dfdx2 (x2-b2) + dfdx3 (x3-b3)
        #    where the derivatives are evaluated at x1,b2,b3,x4 (called b) and
        #    dfdx1|x = dfdx1|b + d2fdx1dx2|b (x2-b2) + d2fdx1dx3|b (x3-b3)
        #    dfdx2|x = dfdx2|b.
        # The dfdx2|x derivative is what it is because f and all derivatives
        # evaluated at x1,b2,b3,x4 are constant with respect to changes in x2.
        # On the other hand, the dfdx1|x derivative is what it is because
        # f and all derivatives evaluated at x1,b2,b3,x4 change with x1.
        # The extrapolation function is non-differentiable at boundaries:
        # i.e., where x_k = a_k or x_k = b_k for at least one k.
        if options['extrapolate']:

            # First we evaluate the vector pointing to each evaluation points
            # from the nearest point on the domain, in a matrix called dx.
            # If the ith evaluation point is not external, dx[i, :] = 0.
            dx = np.empty(n * num['support'] * num['x'])
            self.rmtsc.compute_ext_dist(n, num['support'], x.flatten(), dx)
            dx = dx.reshape((n * num['support'], num['x']))

            isexternal = np.array(np.array(dx, bool), float)

            for ix in range(num['x']):
                # Now we compute the first order term where we have a
                # derivative times (x_k - b_k) or (x_k - a_k).
                data_tmp, rows, cols = self._compute_jac_raw(kx, ix+1, x)
                data_tmp *= dx[:, ix]

                # If we are evaluating a derivative (with index kx),
                # we zero the first order terms for which dx_k = 0.
                if kx != 0:
                    data_tmp *= 1 - isexternal[:, kx-1]

                data += data_tmp

        mtx = scipy.sparse.csc_matrix((data, (rows, cols)), shape=(n, num['coeff']))
        return mtx

    def _predict_output_derivatives(self, x):
        # dy_dyt = dy_dw * (dR_dw)^{-1} * dR_dyt

        n = x.shape[0]
        nw = self.mtx.shape[0]
        nx = x.shape[1]
        ny = self.sol.shape[1]

        p = self.options['approx_order']

        dy_dw = self._compute_prediction_mtx(x, 0)
        if self.full_dof2coeff is not None:
            dy_dw = dy_dw * self.full_dof2coeff
        dy_dw = dy_dw.todense()

        dR_dw = self.mtx

        dy_dyt = {}
        for kx in self.training_points[None]:
            nt = self.training_points[None][kx][0].shape[0]

            dR_dyt = np.zeros((nw, nt, ny))
            for ind_y in range(ny):
                yt_dict = self._get_yt_dict(ind_y)
                dR_dyt[:, :, ind_y] = self._opt_dgrad_dyt(self.sol[:, ind_y], p, yt_dict, kx)

            solver = get_solver(self.options['derivative_solver'])
            solver._setup(dR_dw, self.printer)

            dw_dyt = np.zeros((nw, nt, ny))
            for ind_t in range(nt):
                for ind_y in range(ny):
                    solver._solve(
                        dR_dyt[:, ind_t, ind_y],
                        dw_dyt[:, ind_t, ind_y],
                        ind_y=ind_y)
                    dw_dyt[:, ind_t, ind_y] *= -1.

            if kx == 0:
                dy_dyt[None] = np.einsum('ij,jkl->ikl', dy_dw, dw_dyt)
            else:
                dy_dyt[kx-1] = np.einsum('ij,jkl->ikl', dy_dw, dw_dyt)

        return dy_dyt
