import gc
import six
import numpy as np
from abc import ABCMeta, abstractmethod

from .base import BaseEstimator


class BaseSVM(six.with_metaclass(ABCMeta, BaseEstimator)):
    def __init__(self, kernel='linear', degree=3, C=1.0, epsilon=1e-3, max_iter=100):
        self.max_iter = max_iter
        self._kernel = kernel
        self.degree = degree
        self.C = C
        self.epsilon = epsilon

    # 显式的初始化变量
    def init(self, X: np.ndarray, y: np.ndarray):
        self.m, self.n = X.shape
        self.X = X
        self.y = y
        self.b = 0.0
        self.K = self.kernel_mat(X, X)  # 全部计算好的kernel
        self.alpha = np.zeros(self.m)
        self.g_v = self.g_vec()
        self.E = self.g_v - self.y  # 将Ei保存在一个列表里

    # g(x)，输入xi（X[i]）
    def g(self, i):
        in_sigma = self.alpha * self.y * self.K[i]
        self.g_v[i] = np.sum(in_sigma) + self.b
        return self.g_v[i]

    # vec结尾的是函数，v结尾的是缓存向量
    def g_vec(self):
        before_sigma = self.K * self.alpha * self.y
        return np.sum(before_sigma, axis=-1) + self.b

    # E（x）为g(x)对输入x的预测值和y的差
    def _e(self, i):
        return self.g(i) - self.y[i]

    # 核函数
    def kernel(self, x1: np.ndarray, x2: np.ndarray):
        if self._kernel == 'linear':
            return np.sum(x1 * x2)
        elif self._kernel == 'poly':
            return (np.sum(x1 * x2) + 1) ** self.degree
        return 0

    # 向量化kernel，一次获得xi对所有x的
    def kernel_vec(self, x: np.ndarray):  # todo: 待添加更多kernel
        if self._kernel == 'linear':
            return np.sum(self.X * x, axis=-1)
        elif self._kernel == 'poly':
            return (np.sum(self.X * x, axis=-1) + 1) ** self.degree
        return None

    # 直接计算所有kernel，以后直接调用
    def kernel_mat(self, X1: np.ndarray, X2: np.ndarray):
        x1 = X1[np.newaxis, ...]
        x2 = X2[:, np.newaxis, :]
        if self._kernel == 'linear':
            return np.sum(x1 * x2, axis=-1)  # 广播，对最后一个维度求和
        elif self._kernel == 'poly':
            return (np.sum(x1 * x2, axis=-1) + 1) ** self.degree
        return None

    # 选择α1 α2，返回index
    def select_alpha(self):
        _a = self.alpha
        # 得到mask
        con1, con2 = (_a > 0), (_a < self.C)
        # yi*g(xi)
        ygx = self.y * self.g_v
        # αi == 0 and yi*gxi ≥ 1-ε    KKT条件
        err1 = ygx - 1 + self.epsilon
        err1[(con1 & (err1 <= 0)) | (~con1 & (err1 > 0))] = 0  # 不在此类或符合置0
        # 0 < αi < C and abs(yi*gxi - 1) ≤ ε
        err2 = np.abs(ygx - 1) - self.epsilon
        err2[~con1 | ~con2] = 0  # 置 αi ≤ 0 and αi ≥ C 的为0
        # αi == C and yi*gxi ≤ 1+ε
        err3 = ygx - 1 - self.epsilon
        err3[(con2 & (err3 >= 0)) | (~con2 & (err3 < 0))] = 0
        # 计算总error，排序获得index
        err = err1 ** 2 + err2 ** 2 + err3 ** 2
        # α1为违反KKT条件最严重的点
        i1 = np.argmax(err)
        # 如果E1是+，选择最小的；如果E1是负的，选择最大的
        i2 = np.argmin(self.E) if self.E[i1] >= 0 else np.argmax(self.E)
        if i2 == i1:  # 如果相等，返回第二大or小的
            if self.E[i1] >= 0:
                i2 = np.argsort(self.E)[1]
            else:
                i2 = np.argsort(self.E)[-2]
        return i1, i2

    # 优化α1 α2，更新b
    def optimize(self, i1, i2):
        a1_old, a2_old, b_old = self.alpha[i1], self.alpha[i2], self.b
        _y, _K = self.y, self.K
        # 边界
        if _y[i1] == _y[i2]:
            L = max(0, a2_old + a1_old - self.C)
            H = min(self.C, a2_old + a1_old)
        else:
            L = max(0, a2_old - a1_old)
            H = min(self.C, self.C + a2_old - a1_old)

        E1, E2 = self.E[i1], self.E[i2]

        eta = _K[i1, i1] + _K[i2, i2] - 2 * _K[i1, i2]  # 7.107 η = K11 + K22 - 2K12
        if eta <= 0:
            print('eta <= 0')
            return

        a2_new_unc = a2_old + _y[i2] * (E1 - E2) / eta  # 7.106

        def cut_alpha(a, h, l):  # 7.108
            if a > h:
                return h
            elif a < l:
                return l
            else:
                return a

        a2_new = cut_alpha(a2_new_unc, H, L)

        a1_new = a1_old + _y[i1] * _y[i2] * (a2_old - a2_new)  # 7.109

        b1_new = -E1 - _y[i1] * _K[i1, i1] * (a1_new - a1_old) - \
                 _y[i2] * _K[i2, i1] * (a2_new - a2_old) + b_old  # 7.115
        b2_new = -E2 - _y[i1] * _K[i1, i2] * (a1_new - a1_old) - \
                 _y[i2] * _K[i2, i2] * (a2_new - a2_old) + b_old  # 7.116

        if 0 < a1_new < self.C:
            b_new = b1_new
        elif 0 < a2_new < self.C:
            b_new = b2_new
        else:  # 选择中点，否则b1=b2的
            b_new = (b1_new + b2_new) / 2

        def e_new(i):
            in_sigma = self.alpha * self.y * self.K[i]
            in_sigma[self.alpha <= 0] = 0
            in_sigma[self.alpha > self.C] = 0
            return np.sum(in_sigma) + self.b - self.y[i]

        # 更新参数
        self.alpha[i1], self.alpha[i2] = a1_new, a2_new
        self.b = b_new
        self.E[i1], self.E[i2] = e_new(i1), e_new(i2)

        return None

    # 停机条件
    def shutdown(self):
        c, eps = self.C, self.epsilon
        _a, _y = self.alpha, self.y

        if np.any(_a < 0) | np.any(_a > c):
            return False
        elif np.abs(np.sum(_a * _y)) > eps:
            return False
        else:
            ygx = _y * self.g_v
            if np.any(ygx[_a == 0] < 1 - eps):
                return False
            elif np.any(np.abs(ygx[(0 < _a) & (_a < c)] - 1) > eps):
                return False
            elif np.any(ygx[_a == c] > 1 + eps):
                return False

        return True

    def fit(self, X: np.ndarray, y: np.ndarray):
        self.init(X, y)
        for t in range(self.max_iter):
            i1, i2 = self.select_alpha()  # 选择优化变量
            self.optimize(i1, i2)  # 求解两个变量的最优化问题
            if self.shutdown():  # 停机条件
                # print('Early stop')
                break
        return

    def predict(self, X):
        fx = self.f(X)
        fx[fx >= 0] = 1
        fx[fx < 0] = -1
        return fx

    # 决策函数
    def f(self, X):
        k_mat = self.kernel_mat(self.X, X)  # shape: n_test, n_train
        in_sigma = k_mat * self.alpha * self.y  # 广播
        return np.sum(in_sigma, axis=-1) + self.b


class SVC(BaseSVM):
    def __init__(self, kernel='linear', degree=3, C=1.0, epsilon=1e-3,
                 max_iter=100, decision_function_shape='ovr'):
        super(SVC, self).__init__(kernel=kernel, degree=degree, C=C,
                                  epsilon=epsilon, max_iter=max_iter)
        self.decision_function_shape = decision_function_shape

    # 生成二分类的数据
    def get_xyi(self, X, y, i, j=0):
        xi, yi = X.copy(), y.copy()
        if self.decision_function_shape == 'ovr':
            yi[yi != i] = -1
            yi[yi == i] = 1
        elif self.decision_function_shape == 'ovo':
            pi, pj = yi == i, yi == j
            xi, yi = xi[pi | pj], yi[pi | pj]
            yi[yi == j] = -1
            yi[yi == i] = 1
        return xi, yi

    def fit(self, X: np.ndarray, y: np.ndarray):
        self.n_classes = len(set(y.tolist()))
        if self.decision_function_shape == 'ovr':
            self.K = self.kernel_mat(X, X)
            self.svms = list()
            for i in range(self.n_classes):
                svm = BaseSVM(self._kernel, self.degree, self.C, self.epsilon, self.max_iter)
                svm.fit(*self.get_xyi(X, y, i))
                del svm.K
                gc.collect()
                svm.K = self.K
                self.svms.append(svm)
        elif self.decision_function_shape == 'ovo':
            self.svmd = dict()
            for i in range(self.n_classes):
                for j in range(i + 1, self.n_classes):
                    svm = BaseSVM(self._kernel, self.degree, self.C, self.epsilon, self.max_iter)
                    xij, yij = self.get_xyi(X, y, i, j)
                    svm.fit(xij, yij)
                    self.svmd[f"{i}-{j}"] = svm
        else:
            raise ValueError("不支持的策略!")

    # 多分类决策函数
    def decision_function(self, predictions: np.ndarray) -> np.ndarray:
        pred_y = np.zeros(predictions.shape[0])
        if self.decision_function_shape == 'ovr':
            pred_y = predictions.argmax(axis=-1)
        elif self.decision_function_shape == 'ovo':
            pred_y = predictions.argmax(axis=-1)
        return pred_y

    def predict(self, X):
        n_samples = X.shape[0]
        predictions = np.zeros([n_samples, self.n_classes])
        if self.decision_function_shape == 'ovr':
            for i, svm in enumerate(self.svms):
                f = svm.f(X)  # todo 函数值尺度不一，有偏见
                f_max, f_min = f.max(), f.min()
                predictions[:, i] = f - f_min / (f_max - f_min)
        elif self.decision_function_shape == 'ovo':
            for k, svm in self.svmd.items():
                i, j = int(k[0]), int(k[2])
                pred = svm.predict(X)
                pred[pred == -1] = 0
                predictions[:, i] += pred
                predictions[:, j] += -pred + 1
        return self.decision_function(predictions)

import gc
import six
import numpy as np
from abc import ABCMeta, abstractmethod

from .base import BaseEstimator


class BaseSVM(six.with_metaclass(ABCMeta, BaseEstimator)):
    def __init__(self, kernel='linear', degree=3, C=1.0, epsilon=1e-3, max_iter=100):
        self.max_iter = max_iter
        self._kernel = kernel
        self.degree = degree
        self.C = C
        self.epsilon = epsilon

    # 显式的初始化变量
    def init(self, X: np.ndarray, y: np.ndarray):
        self.m, self.n = X.shape
        self.X = X
        self.y = y
        self.b = 0.0
        self.K = self.kernel_mat(X, X)  # 全部计算好的kernel
        self.alpha = np.zeros(self.m)
        self.g_v = self.g_vec()
        self.E = self.g_v - self.y  # 将Ei保存在一个列表里

    # g(x)，输入xi（X[i]）
    def g(self, i):
        in_sigma = self.alpha * self.y * self.K[i]
        self.g_v[i] = np.sum(in_sigma) + self.b
        return self.g_v[i]

    # vec结尾的是函数，v结尾的是缓存向量
    def g_vec(self):
        before_sigma = self.K * self.alpha * self.y
        return np.sum(before_sigma, axis=-1) + self.b

    # E（x）为g(x)对输入x的预测值和y的差
    def _e(self, i):
        return self.g(i) - self.y[i]

    # 核函数
    def kernel(self, x1: np.ndarray, x2: np.ndarray):
        if self._kernel == 'linear':
            return np.sum(x1 * x2)
        elif self._kernel == 'poly':
            return (np.sum(x1 * x2) + 1) ** self.degree
        return 0

    # 向量化kernel，一次获得xi对所有x的
    def kernel_vec(self, x: np.ndarray):  # todo: 待添加更多kernel
        if self._kernel == 'linear':
            return np.sum(self.X * x, axis=-1)
        elif self._kernel == 'poly':
            return (np.sum(self.X * x, axis=-1) + 1) ** self.degree
        return None

    # 直接计算所有kernel，以后直接调用
    def kernel_mat(self, X1: np.ndarray, X2: np.ndarray):
        x1 = X1[np.newaxis, ...]
        x2 = X2[:, np.newaxis, :]
        if self._kernel == 'linear':
            return np.sum(x1 * x2, axis=-1)  # 广播，对最后一个维度求和
        elif self._kernel == 'poly':
            return (np.sum(x1 * x2, axis=-1) + 1) ** self.degree
        return None

    # 选择α1 α2，返回index
    def select_alpha(self):
        _a = self.alpha
        # 得到mask
        con1, con2 = (_a > 0), (_a < self.C)
        # yi*g(xi)
        ygx = self.y * self.g_v
        # αi == 0 and yi*gxi ≥ 1-ε    KKT条件
        err1 = ygx - 1 + self.epsilon
        err1[(con1 & (err1 <= 0)) | (~con1 & (err1 > 0))] = 0  # 不在此类或符合置0
        # 0 < αi < C and abs(yi*gxi - 1) ≤ ε
        err2 = np.abs(ygx - 1) - self.epsilon
        err2[~con1 | ~con2] = 0  # 置 αi ≤ 0 and αi ≥ C 的为0
        # αi == C and yi*gxi ≤ 1+ε
        err3 = ygx - 1 - self.epsilon
        err3[(con2 & (err3 >= 0)) | (~con2 & (err3 < 0))] = 0
        # 计算总error，排序获得index
        err = err1 ** 2 + err2 ** 2 + err3 ** 2
        # α1为违反KKT条件最严重的点
        i1 = np.argmax(err)
        # 如果E1是+，选择最小的；如果E1是负的，选择最大的
        i2 = np.argmin(self.E) if self.E[i1] >= 0 else np.argmax(self.E)
        if i2 == i1:  # 如果相等，返回第二大or小的
            if self.E[i1] >= 0:
                i2 = np.argsort(self.E)[1]
            else:
                i2 = np.argsort(self.E)[-2]
        return i1, i2

    # 优化α1 α2，更新b
    def optimize(self, i1, i2):
        a1_old, a2_old, b_old = self.alpha[i1], self.alpha[i2], self.b
        _y, _K = self.y, self.K
        # 边界
        if _y[i1] == _y[i2]:
            L = max(0, a2_old + a1_old - self.C)
            H = min(self.C, a2_old + a1_old)
        else:
            L = max(0, a2_old - a1_old)
            H = min(self.C, self.C + a2_old - a1_old)

        E1, E2 = self.E[i1], self.E[i2]

        eta = _K[i1, i1] + _K[i2, i2] - 2 * _K[i1, i2]  # 7.107 η = K11 + K22 - 2K12
        if eta <= 0:
            print('eta <= 0')
            return

        a2_new_unc = a2_old + _y[i2] * (E1 - E2) / eta  # 7.106

        def cut_alpha(a, h, l):  # 7.108
            if a > h:
                return h
            elif a < l:
                return l
            else:
                return a

        a2_new = cut_alpha(a2_new_unc, H, L)

        a1_new = a1_old + _y[i1] * _y[i2] * (a2_old - a2_new)  # 7.109

        b1_new = -E1 - _y[i1] * _K[i1, i1] * (a1_new - a1_old) - \
                 _y[i2] * _K[i2, i1] * (a2_new - a2_old) + b_old  # 7.115
        b2_new = -E2 - _y[i1] * _K[i1, i2] * (a1_new - a1_old) - \
                 _y[i2] * _K[i2, i2] * (a2_new - a2_old) + b_old  # 7.116

        if 0 < a1_new < self.C:
            b_new = b1_new
        elif 0 < a2_new < self.C:
            b_new = b2_new
        else:  # 选择中点，否则b1=b2的
            b_new = (b1_new + b2_new) / 2

        def e_new(i):
            in_sigma = self.alpha * self.y * self.K[i]
            in_sigma[self.alpha <= 0] = 0
            in_sigma[self.alpha > self.C] = 0
            return np.sum(in_sigma) + self.b - self.y[i]

        # 更新参数
        self.alpha[i1], self.alpha[i2] = a1_new, a2_new
        self.b = b_new
        self.E[i1], self.E[i2] = e_new(i1), e_new(i2)

        return None

    # 停机条件
    def shutdown(self):
        c, eps = self.C, self.epsilon
        _a, _y = self.alpha, self.y

        if np.any(_a < 0) | np.any(_a > c):
            return False
        elif np.abs(np.sum(_a * _y)) > eps:
            return False
        else:
            ygx = _y * self.g_v
            if np.any(ygx[_a == 0] < 1 - eps):
                return False
            elif np.any(np.abs(ygx[(0 < _a) & (_a < c)] - 1) > eps):
                return False
            elif np.any(ygx[_a == c] > 1 + eps):
                return False

        return True

    def fit(self, X: np.ndarray, y: np.ndarray):
        self.init(X, y)
        for t in range(self.max_iter):
            i1, i2 = self.select_alpha()  # 选择优化变量
            self.optimize(i1, i2)  # 求解两个变量的最优化问题
            if self.shutdown():  # 停机条件
                # print('Early stop')
                break
        return

    def predict(self, X):
        fx = self.f(X)
        fx[fx >= 0] = 1
        fx[fx < 0] = -1
        return fx

    # 决策函数
    def f(self, X):
        k_mat = self.kernel_mat(self.X, X)  # shape: n_test, n_train
        in_sigma = k_mat * self.alpha * self.y  # 广播
        return np.sum(in_sigma, axis=-1) + self.b


class SVC(BaseSVM):
    def __init__(self, kernel='linear', degree=3, C=1.0, epsilon=1e-3,
                 max_iter=100, decision_function_shape='ovr'):
        super(SVC, self).__init__(kernel=kernel, degree=degree, C=C,
                                  epsilon=epsilon, max_iter=max_iter)
        self.decision_function_shape = decision_function_shape

    # 生成二分类的数据
    def get_xyi(self, X, y, i, j=0):
        xi, yi = X.copy(), y.copy()
        if self.decision_function_shape == 'ovr':
            yi[yi != i] = -1
            yi[yi == i] = 1
        elif self.decision_function_shape == 'ovo':
            pi, pj = yi == i, yi == j
            xi, yi = xi[pi | pj], yi[pi | pj]
            yi[yi == j] = -1
            yi[yi == i] = 1
        return xi, yi

    def fit(self, X: np.ndarray, y: np.ndarray):
        self.n_classes = len(set(y.tolist()))
        if self.decision_function_shape == 'ovr':
            self.K = self.kernel_mat(X, X)
            self.svms = list()
            for i in range(self.n_classes):
                svm = BaseSVM(self._kernel, self.degree, self.C, self.epsilon, self.max_iter)
                svm.fit(*self.get_xyi(X, y, i))
                del svm.K
                gc.collect()
                svm.K = self.K
                self.svms.append(svm)
        elif self.decision_function_shape == 'ovo':
            self.svmd = dict()
            for i in range(self.n_classes):
                for j in range(i + 1, self.n_classes):
                    svm = BaseSVM(self._kernel, self.degree, self.C, self.epsilon, self.max_iter)
                    xij, yij = self.get_xyi(X, y, i, j)
                    svm.fit(xij, yij)
                    self.svmd[f"{i}-{j}"] = svm
        else:
            raise ValueError("不支持的策略!")

    # 多分类决策函数
    def decision_function(self, predictions: np.ndarray) -> np.ndarray:
        pred_y = np.zeros(predictions.shape[0])
        if self.decision_function_shape == 'ovr':
            pred_y = predictions.argmax(axis=-1)
        elif self.decision_function_shape == 'ovo':
            pred_y = predictions.argmax(axis=-1)
        return pred_y

    def predict(self, X):
        n_samples = X.shape[0]
        predictions = np.zeros([n_samples, self.n_classes])
        if self.decision_function_shape == 'ovr':
            for i, svm in enumerate(self.svms):
                f = svm.f(X)  # todo 函数值尺度不一，有偏见
                f_max, f_min = f.max(), f.min()
                predictions[:, i] = f - f_min / (f_max - f_min)
        elif self.decision_function_shape == 'ovo':
            for k, svm in self.svmd.items():
                i, j = int(k[0]), int(k[2])
                pred = svm.predict(X)
                pred[pred == -1] = 0
                predictions[:, i] += pred
                predictions[:, j] += -pred + 1
        return self.decision_function(predictions)

# 显式的初始化变量 # 全部计算好的kernel # 将Ei保存在一个列表里 # g(x)，输入xi（X[i]） # vec结尾的是函数，v结尾的是缓存向量 # E（x）为g(x)对输入x的预测值和y的差 # 核函数 # 向量化kernel，一次获得xi对所有x的 # todo: 待添加更多kernel # 直接计算所有kernel，以后直接调用 # 广播，对最后一个维度求和 # 选择α1 α2，返回index # 得到mask # yi*g(xi) # αi == 0 and yi*gxi ≥ 1-ε    KKT条件 # 不在此类或符合置0 # 0 < αi < C and abs(yi*gxi - 1) ≤ ε # 置 αi ≤ 0 and αi ≥ C 的为0 # αi == C and yi*gxi ≤ 1+ε # 计算总error，排序获得index # α1为违反KKT条件最严重的点 # 如果E1是+，选择最小的；如果E1是负的，选择最大的 # 如果相等，返回第二大or小的 # 优化α1 α2，更新b # 边界 # 7.107 η = K11 + K22 - 2K12 # 7.106 # 7.108 # 7.109 # 7.115 # 7.116 # 选择中点，否则b1=b2的 # 更新参数 # 停机条件 # 选择优化变量 # 求解两个变量的最优化问题 # 停机条件 # print('Early stop') # 决策函数 # shape: n_test, n_train # 广播 # 生成二分类的数据 # 多分类决策函数 # todo 函数值尺度不一，有偏见