# Copyright (c) 2020, Anthony Nouy, Erwan Grelier
# This file is part of tensap (tensor approximation package).

# tensap is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# tensap is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.

# You should have received a copy of the GNU Lesser General Public License
# along with tensap.  If not, see <https://www.gnu.org/licenses/>.

'''
Module linear_model_learning_density_l2.

'''

import numpy as np
import tensap


class LinearModelLearningDensityL2(tensap.LinearModelLearning):
    '''
    Class LinearModelLearningDensityL2.

    Attributes
    ----------
    is_basis_orthonormal : bool, optional
        Boolean indicating if the basis used to compute the approximation
        is orthonormal according to some reference measure. The default is
        True.

    '''

    def __init__(self):
        '''
        Constructor for the class LinearModelLearningDensityL2.

        Returns
        -------
        None.

        '''
        super().__init__(tensap.SquareLossFunction())

        self.is_basis_orthonormal = True

    def solve(self):
        '''
        Solution (Ordinary or Regularized) of the minimization problem and
        cross-validation procedure.

        Returns
        -------
        sol : numpy.ndarray  or tensap.FunctionalBasisArray
            The solution of the minimization problem.
        output : dict
            Outputs of the algorithm.

        '''
        self.initialize()

        if self.basis_adaptation:
            sol, output = self._solve_basis_adaptation()
        elif self.regularization:
            sol, output = self._solve_regularized()
        else:
            sol, output = self._solve_standard()

        if self.test_error:
            f_eval = np.matmul(self.basis_eval_test, sol)
            output['test_error'] = self.loss_function.test_error(
                f_eval, self.test_data, sol.norm()**2)

        if self.basis is not None:
            sol = tensap.FunctionalBasisArray(sol, self.basis)

        return sol, output

    def _solve_standard(self):
        '''
        Solution of the minimization problem and cross-validation procedure
        (if requested).

        Returns
        -------
        sol : numpy.ndarray
            The solution of the minimization problem.
        output : dict
            Outputs of the algorithm.

        '''
        assert self.is_basis_orthonormal, \
            'Only implemented for orthonormal bases.'

        A = self.basis_eval
        if np.ndim(A) == 3:
            A = np.squeeze(A, axis=2)

        sol = np.mean(A, axis=0)
        if isinstance(self.training_data, list) and \
            len(self.training_data) == 2 and \
                np.size(self.training_data[1]):
            b = np.reshape(self.training_data[1], sol.shape)

            sol -= b
        else:
            b = []

        output = {}
        if self.error_estimation and self.error_estimation_type == 'leave_out':
            N = A.shape[0]
            if np.size(b) == 0:
                output['error'] = (-N**2)/(1-N)**2*np.linalg.norm(sol)**2 + \
                    (2*N-1)/(N*(N-1)**2)*np.sum(np.ravel(A)**2)
            else:
                output['error'] = (N**2-2*N)/(N-1)**2*np.linalg.norm(sol)**2 +\
                    1/(N-1)**2*np.linalg.norm(b)**2 - \
                    2/(N-1)**2*np.sum(np.matmul(A, b)) + \
                    (2*N-1)/(N*(N-1)**2)*np.sum(np.ravel(A)**2) - \
                    2/(N-1)*np.sum(np.matmul(A, sol)) + \
                    2*np.sum(sol*b)

        return sol, output

    def _solve_regularized(self):
        '''
        Solution of the regularized minimization problem and cross-validation
        procedure.

        Returns
        -------
        sol : numpy.ndarray
            The solution of the minimization problem.
        output : dict
            Outputs of the algorithm.

        '''
        sol_standard, _ = self._solve_standard()

        A = self.basis_eval
        if np.ndim(A) == 3:
            A = np.squeeze(A, axis=2)
        N = A.shape[0]

        if isinstance(self.training_data, list) and \
            len(self.training_data) == 2 and \
                np.size(self.training_data[1]):
            b = np.reshape(self.training_data[1], sol_standard.shape)
        else:
            b = []

        A_square_sum = np.sum(A**2, 0)

        list_sort = np.argsort(-np.abs(sol_standard))

        if 'included_coefficients' in self.regularization_options:
            incl_coef = self.regularization_options['included_coefficients']
            list_sort = list_sort[np.logical_not(np.in1d(list_sort,
                                                         incl_coef))]

            sol_incl_coef = np.array(sol_standard)
            sol_incl_coef[np.setdiff1d(range(len(sol_incl_coef)),
                                       incl_coef)] = 0

            if np.size(b) == 0:
                err_incl_coef = \
                    (-N**2)/(1-N)**2*np.linalg.norm(sol_incl_coef)**2 + \
                    (2*N-1)/(N*(N-1)**2)*np.sum(
                        np.ravel(A_square_sum)[incl_coef])
            else:
                b_incl_coef = b[incl_coef]
                err_incl_coef = \
                    (N**2-2*N)/(N-1)**2*np.linalg.norm(sol_incl_coef)**2 +\
                    1/(N-1)**2*np.linalg.norm(b_incl_coef)**2 - \
                    2/(N-1)**2*np.sum(
                        np.matmul(A[:, incl_coef], b_incl_coef)) + \
                    (2*N-1)/(N*(N-1)**2)*np.sum(
                        np.ravel(A_square_sum)[incl_coef]) -\
                    2/(N-1)*np.sum(np.matmul(A, sol_incl_coef)) + \
                    2*np.sum(sol_incl_coef*b_incl_coef)

        if list_sort.size == 0:
            return sol_standard, {'error': np.nan}

        n = list_sort.size
        sol = np.zeros((sol_standard.shape[0], n))
        err = np.zeros(n)

        for i in range(n):
            ind = list_sort[:i+1]
            if 'included_coefficients' in self.regularization_options:
                ind = np.concatenate((incl_coef, ind))
            sol_red = sol_standard[ind]

            if np.size(b) == 0:
                err[i] = \
                    (-N**2)/(1-N)**2*np.linalg.norm(sol_red)**2 + \
                    (2*N-1)/(N*(N-1)**2)*np.sum(
                        np.ravel(A_square_sum)[ind])
            else:
                err[i] = \
                    (N**2-2*N)/(N-1)**2*np.linalg.norm(sol_red)**2 +\
                    1/(N-1)**2*np.linalg.norm(b[ind])**2 - \
                    2/(N-1)**2*np.sum(np.matmul(A[:, ind], b[ind])) + \
                    (2*N-1)/(N*(N-1)**2)*np.sum(np.ravel(A_square_sum)[ind]) -\
                    2/(N-1)*np.sum(np.matmul(A[:, ind], sol_red)) + \
                    2*np.sum(sol_red*b[ind])

            sol[ind, i] = np.squeeze(sol_red)

        if 'included_coefficients' in self.regularization_options:
            sol = np.hstack((np.reshape(sol_incl_coef, (-1, 1)), sol))
            err = np.hstack((err_incl_coef, err))

        ind = np.argmin(err)
        pattern = sol != 0

        output = {}
        output['error'] = err[ind]
        output['error_path'] = err
        output['ind'] = ind
        output['pattern'] = pattern[:, ind]
        output['pattern_path'] = pattern
        output['optimal_solution'] = sol[:, ind]
        output['solution_path'] = sol

        return sol[:, ind], output

    def _solve_basis_adaptation(self):
        '''
        Solution of the minimization problem with working-set and
        cross-validation procedure.

        Returns
        -------
        sol : numpy.ndarray
            The solution of the minimization problem.
        output : dict
            Outputs of the algorithm.

        '''
        if self.basis_adaptation_path is None:
            solpath = np.triu(np.full([self.basis_eval.shape[1]]*2, True))
        else:
            solpath = self.basis_adaptation_path

        sol, output = self._select_optimal_path(solpath)
        output['flag'] = 2
        return sol, output

    def _select_optimal_path(self, solpath):
        '''
        Selection of a solution using leave-one-out cross-validation error.

        The dictionnary output contains:
            - error: the leave-one-out cross-validation error estimate,
            - error_path: the regularization path of leave-one-out
            cross-validation error estimates,
            - ind: the index of the optimal pattern,
            - pattern: the optimal sparsity pattern,
            - pattern_path: the regularization path of sparsity patterns,
            - solution_path: the regularization path of the solutions,
            - optimal_solution: the optimal solution.

        Parameters
        ----------
        solpath : numpy.ndarray
            Array of shape (P, m) whose columns give m potential solutions with
            different sparsity patterns.

        Returns
        -------
        numpy.ndarray
            The optimal solution of the minimization problem.
        dict
            Outputs of the algorithm.

        '''
        sol_standard, output = self._solve_standard()

        A = self.basis_eval
        if np.ndim(A) == 3:
            A = np.squeeze(A, axis=2)

        if isinstance(self.training_data, list) and \
            len(self.training_data) == 2 and \
                np.size(self.training_data[1]):
            b = np.reshape(self.training_data[1], sol_standard.shape)
        else:
            b = []

        A_square_sum = np.sum(A**2, 0)
        if np.size(b) != 0:
            A_sum = np.sum(A, 0)

        pattern = solpath != 0
        rep = np.any(pattern, 0)
        pattern = pattern[:, rep]

        _, ind = np.unique(pattern, return_index=True, axis=1)
        ind = np.sort(ind)
        pattern = pattern[:, ind]

        ind = np.sum(pattern != 0, 0) <= A.shape[0]
        pattern = pattern[:, ind]

        if np.linalg.norm(pattern) == 0:
            return sol_standard, output

        sol = np.zeros(pattern.shape)
        err = np.zeros(pattern.shape[1])

        N = A.shape[0]

        for i in range(pattern.shape[1]):
            ind = pattern[:, i]
            sol_red = sol_standard[ind]

            if np.size(b) == 0:
                err[i] = (-N**2)/(1-N)**2*np.linalg.norm(sol_red)**2 + \
                    (2*N-1)/(N*(N-1)**2)*np.sum(np.ravel(A_square_sum)[ind])
            else:
                b_red = b[ind]
                err[i] = (N**2-2*N)/(N-1)**2*np.linalg.norm(sol_red)**2 +\
                    1/(N-1)**2*np.linalg.norm(b_red)**2 - \
                    2/(N-1)**2*np.sum(A_sum[ind]*b_red) + \
                    (2*N-1)/(N*(N-1)**2)*np.sum(np.ravel(A_square_sum)[ind]) -\
                    2/(N-1)*np.sum(A_sum[ind]*sol_red) + \
                    2*np.sum(sol_red*b_red)
            sol[ind, i] = np.squeeze(sol_red)

            if i > 1 and \
                self.model_selection_options['stop_if_error_increase'] and \
                    err[i] > 2*err[i-1]:
                print('stop_if_error_increase')
                err[i+1:] = np.inf

        ind = np.argmin(err)

        output['error'] = err[ind]
        output['error_path'] = err
        output['ind'] = ind
        output['pattern'] = pattern[:, ind]
        output['pattern_path'] = pattern
        output['optimal_solution'] = sol[:, ind]
        output['solution_path'] = sol

        return sol[:, ind], output
