"""
Multi-output Support Vector Regression
"""
# Copyright (C) 2020 Xinze Zhang, Kaishuai Xu, Siyue Yang, Yukun Bao
# <xinze@hust.edu.cn>, <xu.kaishuai@gmail.com>, <siyue_yang@hust.edu.cn>, <yukunbao@hust.edu.cn>

# This program is free software: you can redistribute it and/or modify
# it under the terms of the Apache.

# This program 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
# Apache License for more details.


from task.util import toTorch
import numpy as np
from sklearn.metrics.pairwise import pairwise_kernels
from sklearn.metrics import mean_squared_error
import torch
from  task.TaskLoader import mlp_dataset
from task.TaskLoader import Opt
from task.metric import rmse

class MSVR():
    def __init__(self, opts=None, logger=None):
        self.opts = opts
        self.logger = logger
        self.opts.device = torch.device('cpu')
        for (arg, value) in opts.dict.items():
            self.logger.info("Argument %s: %r", arg, value)
        
        super(MSVR, self).__init__()
        
        
        self.kernel = opts.kernel if 'kernel' in opts.kernel else 'rbf'
        self.degree = opts.degree
        self.gamma = opts.gamma
        self.coef = opts.coef
        self.tol = opts.tol
        self.C = opts.C
        self.epsilon = opts.epsilon
        self.Beta = None
        self.NSV = None
        self.xTrain = None
        self.Lag_order  = opts.lag_order
        self.Output_dim = opts.H
        

    def data_loader(self, data,):
        '''
        Transform the numpy array data into the pytorch data_loader
        '''
        set_data = mlp_dataset(data, self.Output_dim, self.Lag_order)      
        # set_loader = torch_dataloader(set_data, batch_size= set_data.samples,cuda= False)
        return set_data
    
    # @profile
    def xfit(self, train_data, val_data):
        """fit the data to msvr
        
        Parameters
        ----------
        input : torch array-like shape, (n_samples, Input_dim)
            The data to be transformed.
        target : torch array-like shape, (n_samples, Output_dim)
            The data to be transformed.
            
        Returns
        -------
        self : returns an instance of self.
        """
        self.fit_info = Opt()
        
        train_loader = self.data_loader(train_data)
        val_loader = self.data_loader(val_data)
        
        with torch.no_grad():
            train_x, train_y = train_loader.data, train_loader.label
            # for batch_x, batch_y in train_loader:
            #     batch_x = batch_x.to(torch.float32).to(self.opts.device)
            #     batch_y = batch_y.to(torch.float32).to(self.opts.device)
            #     train_x, train_y = batch_x.detach().clone().numpy(), batch_y.detach().clone().numpy()

            val_x, val_y = val_loader.data, val_loader.label
            # for batch_x, batch_y in val_loader:
            #     batch_x = batch_x.to(torch.float32).to(self.opts.device)
            #     batch_y = batch_y.to(torch.float32).to(self.opts.device)
            #     val_x, val_y = batch_x.detach().clone().numpy(), batch_y.detach().clone().numpy()
                
            self.fit(train_x, train_y)
            pred = self.predict(train_x)
            vpred = self.predict(val_x)

            self.fit_info.trmse = rmse(train_y,pred)
            self.fit_info.vrmse = rmse(val_y, vpred)

        self.logger.critical(
            'Training RMSE: {:.4g} \t Validating RMSE: {:.4g}'.format(self.fit_info.trmse, self.fit_info.vrmse))
            
        return self.fit_info

    def fit(self, x, y):
        self.xTrain = x.copy()
        C = self.C
        epsi = self.epsilon
        tol = self.tol

        n_m = np.shape(x)[0]  # num of samples
        n_d = np.shape(x)[1]  # input data dimensionality
        n_k = np.shape(y)[1]  # output data dimensionality (output variables)

        # H = kernelmatrix(ker, x, x, par)
        H = pairwise_kernels(x, x, metric=self.kernel, filter_params=True,
                             degree=self.degree, gamma=self.gamma, coef=self.coef)

        self.Beta = np.zeros((n_m, n_k))

        #E = prediction error per output (n_m * n_k)
        E = y - np.dot(H, self.Beta)
        #RSE
        u = np.sqrt(np.sum(E**2, 1, keepdims=True))

        #RMSE
        RMSE = []
        RMSE_0 = np.sqrt(np.mean(u**2))
        RMSE.append(RMSE_0)

        #points for which prediction error is larger than epsilon
        i1 = np.where(u > epsi)[0]

        #set initial values of alphas a (n_m * 1)
        a = 2 * C * (u - epsi) / u

        #L (n_m * 1)
        L = np.zeros(u.shape)

        # we modify only entries for which  u > epsi. with the sq slack
        L[i1] = u[i1]**2 - 2 * epsi * u[i1] + epsi**2

        #Lp is the quantity to minimize (sq norm of parameters + slacks)
        Lp = []
        BetaH = np.dot(np.dot(self.Beta.T, H), self.Beta)
        Lp_0 = np.sum(np.diag(BetaH), 0) / 2 + C * np.sum(L)/2
        Lp.append(Lp_0)

        eta = 1
        k = 1
        hacer = 1
        # val = 1

        while(hacer):
            Beta_a = self.Beta.copy()
            E_a = E.copy()
            u_a = u.copy()
            i1_a = i1.copy()

            M1 = H[i1][:, i1] + \
                np.diagflat(1/a[i1]) + 1e-10 * np.eye(len(a[i1]))

            #compute betas
            #       sal1 = np.dot(np.linalg.pinv(M1),y[i1])  #求逆or广义逆（M-P逆）无法保证M1一定是可逆的？
            sal1 = np.dot(np.linalg.inv(M1), y[i1])

            eta = 1
            self.Beta = np.zeros(self.Beta.shape)
            self.Beta[i1] = sal1.copy()

            #error
            E = y - np.dot(H, self.Beta)
            #RSE
            u = np.sqrt(np.sum(E**2, 1)).reshape(n_m, 1)
            i1 = np.where(u >= epsi)[0]

            L = np.zeros(u.shape)
            L[i1] = u[i1]**2 - 2 * epsi * u[i1] + epsi**2

            #%recompute the loss function
            BetaH = np.dot(np.dot(self.Beta.T, H), self.Beta)
            Lp_k = np.sum(np.diag(BetaH), 0) / 2 + C * np.sum(L)/2
            Lp.append(Lp_k)

            #Loop where we keep alphas and modify betas
            while(Lp[k] > Lp[k-1]):
                eta = eta/10
                i1 = i1_a.copy()

                self.Beta = np.zeros(self.Beta.shape)
                #%the new betas are a combination of the current (sal1)
                #and of the previous iteration (Beta_a)
                self.Beta[i1] = eta*sal1 + (1-eta)*Beta_a[i1]

                E = y - np.dot(H, self.Beta)
                u = np.sqrt(np.sum(E**2, 1)).reshape(n_m, 1)

                i1 = np.where(u >= epsi)[0]

                L = np.zeros(u.shape)
                L[i1] = u[i1]**2 - 2 * epsi * u[i1] + epsi**2
                BetaH = np.dot(np.dot(self.Beta.T, H), self.Beta)
                Lp_k = np.sum(np.diag(BetaH), 0) / 2 + C * np.sum(L)/2
                Lp[k] = Lp_k

                #stopping criterion 1
                if(eta < 1e-16):
                    Lp[k] = Lp[k-1] - 1e-15
                    self.Beta = Beta_a.copy()

                    u = u_a.copy()
                    i1 = i1_a.copy()

                    hacer = 0

            #here we modify the alphas and keep betas
            # a_a = a.copy()
            a = 2 * C * (u - epsi) / u

            RMSE_k = np.sqrt(np.mean(u**2))
            RMSE.append(RMSE_k)

            if((Lp[k-1]-Lp[k])/Lp[k-1] < tol):
                hacer = 0

            k = k + 1

            #stopping criterion #algorithm does not converge. (val = -1)
            if(len(i1) == 0):
                hacer = 0
                self.Beta = np.zeros(self.Beta.shape)
                # val = -1

        self.NSV = len(i1)

    def predict(self, input):
        input = torch.tensor(input).float().to(self.opts.device).numpy()
        H = pairwise_kernels(input, self.xTrain, metric=self.kernel, filter_params=True,
                             degree=self.degree, gamma=self.gamma, coef=self.coef)
        yPred = np.dot(H, self.Beta)
        return yPred

    def loader_pred(self, data_loader):
        x, y = data_loader.data, data_loader.label
        pred = self.predict(x)
        return x, y, pred

    def task_pred(self, task_data):
        data_loader = self.data_loader(task_data)
        x, y, pred = self.loader_pred(data_loader)
        return x, y, pred