from Learnware_package import *
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from plot_utils import plot_density, Figure_Manager, plot_boundary, plot_region, plot_circle
from Lane_utils import *



class Oracle():
    # The oracle binary classifier
    # Predicts 1 inside the circle, 0 outside the circle.
    def __init__(self, Radius):
        self.radius = Radius

    def predict(self, X):
        n = X.shape[0]
        y = np.zeros(n)
        for i in range(n):
            x = X[i, :]
            # positive inside circle
            if np.linalg.norm(x) <= self.radius:
                y[i] = 1
        return y

""" This function is for generating the toy data.

    Args:
        Radius: the radius of the decision boundary.

    Returns:
        X: the training data of local tasks.
        Xt_task: the testing data under task-recurrent assumption.
        Xt_inst: the testing data under instance-recurrent assumption.
        Task_inst: the superclasses of the testing data.

    Reference:
        [1] Xi-Zhu Wu, Wenkai Xu, Song Liu, and Zhi-Hua Zhou. Model Reuse with Reduced Kernel Mean Embedding Speciﬁcation. In arXiv, 2020.
"""


def generate_toy_data(Radius):
    h = Radius * np.sqrt(2) / 2
    center_list = [[Radius, 0], [h, h], [0, Radius], [-h, h], [-Radius, 0], [-h, -h], [0, -Radius], [h, -h]]
    X_list = []
    sigma = 0.2  # default: 0.2
    N = 80
    for center in center_list:
        mu = np.array(center)
        X_list.append(np.random.normal(loc=mu, scale=sigma, size=(N, 2)))

    X = [0, 0, 0, 0]
    X[0] = np.concatenate((X_list[0], X_list[1]), axis=0)
    X[1] = np.concatenate((X_list[2], X_list[3]), axis=0)
    X[2] = np.concatenate((X_list[4], X_list[5]), axis=0)
    X[3] = np.concatenate((X_list[6], X_list[7]), axis=0)

    test_N = 1000
    weight = [0.7, 0.7, 0.3, 0.3, 0.0, 0.0, 0.0, 0.0]
    inst_list = []
    Yinst_list = []
    task_list = []
    for Y_num, (i, j) in enumerate([(0, 1), (2, 3), (4, 5), (6, 7)]):
        nb = int(round(test_N * weight[i] / 2))
        X_i = np.random.normal(loc=center_list[i], scale=sigma, size=(nb, 2))
        X_j = np.random.normal(loc=center_list[j], scale=sigma, size=(nb, 2))
        inst_list.append(np.concatenate((X_i, X_j), axis=0))
        Yinst_list.append([Y_num for _ in range(nb*2)])
        if i <= 4:
            task_list.append(np.ones(nb * 2) * int(i / 2))
        else:
            task_list.append(np.ones(nb * 2) * 10000)  # new class

    Xt_inst = np.concatenate(inst_list, axis=0)
    Yt_inst = np.concatenate(Yinst_list, axis=0)
    Task_list = np.concatenate(task_list, axis=0)

    test_n = int(round(test_N / 2))
    X1 = np.random.normal(loc=center_list[2], scale=sigma, size=(test_n, 2))
    X2 = np.random.normal(loc=center_list[3], scale=sigma, size=(test_n, 2))
    Xt_task = np.concatenate((X1, X2), axis=0)

    for i in range(len(X)):
        np.random.shuffle(X[i])
    np.random.shuffle(Xt_task)

    # shuffle_index = np.arange(Xt_inst.shape[0])
    # np.random.shuffle(shuffle_index)
    # Xt_inst = Xt_inst[shuffle_index, :]
    # Yt_inst = Yt_inst[shuffle_index, :]
    Task_inst = Task_list

    return X, Xt_task, Xt_inst, Task_inst, Yt_inst


if __name__ == '__main__':
    save_path = '../Result/toy.pt'
    np.random.seed(1)
    Radius = 1

    ''' The Upload Phase '''
    # Generating task models,  task datasets, and requirement dataset
    if not os.path.exists(save_path):
        # Generate toy data
        X, Xt_task, Xt_inst, Task_inst, Ytinst = generate_toy_data(Radius)
        X = [X[0], X[1], X[2], X[3]]

        c = len(X)
        oracle = Oracle(Radius)
        Y = [None, None, None, None]
        for i in range(c):
            Y[i] = oracle.predict(X[i])
        Yt_task = oracle.predict(Xt_task)
        Yt_inst = oracle.predict(Xt_inst)
        for i in range(Xt_inst.shape[0]):
            if Task_inst[i] == 10000:
                Yt_inst[i] = 10000

        # Pre-trained models
        f = []
        for i in range(c):
            f_local = SVC(kernel='rbf', gamma=2, probability=True,
                          decision_function_shape='ovr')
            f_local.fit(X[i], Y[i])
            f.append(f_local)

        # Specification generation and submitting to learnware market
        K = 5
        Phi = []
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        for i in range(c):
            rs = Lane_toy()
            task_data = torch.from_numpy(X[i]).to(device)
            task_data = task_data.type(torch.float32)
            rs.fit(task_data, Y[i], task_data.shape[1], 2, K, device)
            Phi.append(rs)


        ''' The Deployment Phase '''
        # Mixture  weight estimation
        w_estimator = weight_estimation(Phi)
        estimate_weight, specification = w_estimator.mixture_weight_estimate(Xt_inst, device, K)
        print('True weight: (0.7,0.3,0.0,0.0)')
        print('Estimated weight: (%.3f,%.3f,%.3f,%.3f)' % (estimate_weight[0], estimate_weight[1], estimate_weight[2], estimate_weight[3]))

        # Generating mimic dataset
        n = Xt_inst.shape[0]
        X_mimic, Y_mimic = w_estimator.weight_sampling_random(n, estimate_weight)

        # Training learnware index selector
        selector_learnware = SVC(kernel='rbf', gamma=2, probability=True, decision_function_shape='ovr')
        selector_learnware.fit(X_mimic, Y_mimic)

        # Predicting requirement
        selected_idx_learnware = selector_learnware.predict(Xt_inst).tolist()
        Ytinst = Ytinst.tolist()
        label = [int(y) for y in Ytinst]
        acc = accuracy_score(label, selected_idx_learnware)

        print('Accuracy on Instance-recurrent: %.4f' % (acc))

        torch.save({'oracle': oracle, 'X': X, 'specification': specification, 'label': label, 'Xt_inst': Xt_inst, 'acc': acc, 'weights': estimate_weight,
                    'X_mimic': X_mimic, 'Y_mimic': Y_mimic, 'selector_learnware': selector_learnware, 'selected_idx_learnware': selected_idx_learnware}, save_path)

    else:
        infor = torch.load(save_path)
        oracle = infor['oracle']
        X = infor['X']
        # acc = infor['acc']
        # estimate_weight = infor['weights']
        Xt_inst = infor['Xt_inst']
        specification = infor['specification']
        label = infor['label']
        X_mimic = infor['X_mimic']
        Y_mimic = infor['Y_mimic']
        selector_learnware = infor['selector_learnware']
        selected_idx_learnware = infor['selected_idx_learnware']


    ''' Visualization  '''

    # plt task dataset
    plt.figure(figsize=(5, 5))
    xlimit = 2 * Radius
    ylimit = 2 * Radius
    figure_manager = Figure_Manager(max_row=1, max_col=1, xlim=[-xlimit, xlimit], ylim=[-ylimit, ylimit])

    ax = figure_manager.next()
    markers = ['^', 'o', 's', '*']
    colors = ['#FFCC00', '#0066FF']
    custom_cmap = LinearSegmentedColormap.from_list('my_list', colors, N=5)
    m_size = [60, 60, 60, 60]
    l_width = 0.3
    for i, x in enumerate(X):
        plt.scatter(x[:, 0], x[:, 1], c=oracle.predict(x[:, :]),
                    s=m_size[i], linewidth=l_width, alpha=1,
                    cmap=custom_cmap, edgecolors='k',
                    marker=markers[i], label='$X_{%d}$' % (i + 1))
    legend = plt.legend(["Task 1", "Task 2", "Task 3", "Task 4"], loc='lower right')
    plot_circle(ax, Radius)


    # plt specification
    plt.figure(figsize=(5, 5))
    figure_manager = Figure_Manager(max_row=1, max_col=1, xlim=[-xlimit, xlimit], ylim=[-ylimit, ylimit])
    figure_manager.next()
    # bound_square = [-2, 2, -2, 2]
    # cmap_strs = ['Greens', 'Greens', 'Greens', 'Greens']

    for i, Z in enumerate(specification):
            plt.scatter(Z[:, 0], Z[:, 1], edgecolors='k', s=m_size[i]*2, linewidths=l_width, alpha=1,
                        marker=markers[i], label='$Z_{%d}$' % (i + 1))
    legend = plt.legend(["Task 1", "Task 2", "Task 3", "Task 4"], loc='lower right')


    # plt requirement dataset
    plt.figure(figsize=(5, 5))
    figure_manager = Figure_Manager(max_row=1, max_col=1, xlim=[-xlimit, xlimit], ylim=[-ylimit, ylimit])
    ax = figure_manager.next()
    plt.scatter(Xt_inst[:, 0], Xt_inst[:, 1],
                edgecolors='#000000', facecolor='Gray', s=m_size * 250,
                marker='H', linewidth=l_width, alpha=1)
    legend = plt.legend(["Requirement data"], loc='lower right')


    # plt mimic dataset
    plt.figure(figsize=(5, 5))
    figure_manager = Figure_Manager(max_row=1, max_col=1, xlim=[-xlimit, xlimit], ylim=[-ylimit, ylimit])
    ax = figure_manager.next()

    class_num = np.unique(Y_mimic).shape[0]
    indices_class = [[] for c in range(class_num)]
    for i, lab in enumerate(Y_mimic):
        indices_class[lab].append(i)

    plt_mimic_data = []

    for ind in indices_class:
        plt_mimic_data.append(X_mimic[ind])
        # print(plt_mimic_data)
    for i, mimic in enumerate(plt_mimic_data):
        plt.scatter(mimic[:, 0], mimic[:, 1], edgecolors='#000000', facecolor='Gray', s=m_size[i],
                marker=markers[i], linewidth=l_width, alpha=1, label='$Xmimic_{%d}$' % (i + 1))
    legend = plt.legend(["Mimic data", "Mimic data", "Mimic data", "Mimic data"], loc='lower right')

    # plt margin
    plt.figure(figsize=(5, 5))
    figure_manager = Figure_Manager(max_row=1, max_col=1, xlim=[-xlimit, xlimit], ylim=[-ylimit, ylimit])
    ax = figure_manager.next()

    plot_boundary(selector_learnware, bound_square=[-2, 2, -2, 2])

    # plt learnware result
    plt.figure(figsize=(5, 5))
    figure_manager = Figure_Manager(max_row=1, max_col=1, xlim=[-xlimit, xlimit], ylim=[-ylimit, ylimit])
    ax = figure_manager.next()

    plot_boundary(selector_learnware, bound_square=[-2, 2, -2, 2])

    indices_class = [[] for c in range(class_num)]
    for i, lab in enumerate(selected_idx_learnware):
        indices_class[lab].append(i)

    data_Xt_inst = []
    for ind in indices_class:
        data_Xt_inst.append(Xt_inst[ind])
    for i, data in enumerate(data_Xt_inst):
        plt.scatter(data[:, 0], data[:, 1], c=oracle.predict(data[:, :]),
                    edgecolors='#000000', cmap=custom_cmap, s=m_size[i],
                    marker=markers[i], linewidth=l_width, alpha=1)

    plot_circle(ax, Radius)
    plt.show()








