import torch
from Network.network_setting import get_network
from Learnware_package import *

def get_time():
    return str(time.strftime("[%Y-%m-%d %H:%M:%S]", time.localtime()))


class Lane_toy():

    def __init__(self):
        self.z = []
        self.label = []
        self.num_classes = 0
        self.batch_real = 256
        self.Iteration = 2000
        self.model = 'ConvNet_1D'

    def assign_indices_class(self, labels, n_classes):
        indices_class = [[] for c in range(n_classes)]
        for i, lab in enumerate(labels):
            indices_class[int(lab)].append(i)
        return indices_class

    def get_data_class_cov(self, dataset, n_classes, indices_class):
        data_class_cov = []
        dataset_new = dataset.cpu().detach().numpy()
        for i in range(n_classes):
            dataset_class = dataset_new[indices_class[i], :]
            data_class_cov.append(np.cov(dataset_class, rowvar=False))

        return data_class_cov

    def fit(self, dataset, labels, channel, num_classes, specification_size, device):
        self.ipc = specification_size

        indices_class = self.assign_indices_class(labels, num_classes)

        def get_data_shuffle(c, n): # get random n data from class c
            idx_shuffle = np.random.permutation(indices_class[c])[:n]
            return dataset[idx_shuffle]


        data_syn = torch.rand(size=(num_classes * self.ipc, channel),
                              dtype=torch.float32,
                              requires_grad=True,
                              device=device)
        label_syn = torch.tensor([np.ones(self.ipc) * i for i in range(num_classes)],
                                 dtype=torch.long,
                                 requires_grad=False,
                                 device=device).view(-1)
        for c in range(num_classes):
            data_syn.data[c*self.ipc:(c+1)*self.ipc] = get_data_shuffle(c, self.ipc).detach().data

        optimizer = torch.optim.SGD([data_syn, ],
                                    lr=0.01,
                                    momentum=0.5)
        record_loss = []
        best_loss = 1000

        for it in range(self.Iteration+1):
            net = get_network(self.model,
                              channel=channel,
                              num_classes=num_classes,
                              ).to(device)
            net.train()
            for param in list(net.parameters()):
                param.requires_grad = False
            embed = net.module.embed if torch.cuda.device_count() > 1 else net.embed  # for GPU parallel
            loss_avg = 0
            record = {}
            record_data = {}
            ''' update synthetic data '''
            real_all = []
            syn_all = []
            loss = torch.tensor(0.0).to(device)
            for c in range(num_classes):
                real = get_data_shuffle(c, self.batch_real)
                syn = data_syn[c * self.ipc:(c + 1) * self.ipc].reshape((self.ipc, channel))
                real_all.append(real)
                syn_all.append(syn)
            real_all = torch.cat(real_all, dim=0)
            syn_all = torch.cat(syn_all, dim=0)
            real_all = real_all.unsqueeze(0).permute(1, 2, 0)
            syn_all = syn_all.unsqueeze(0).permute(1, 2, 0)

            output_real = embed(real_all).detach()
            # syn_all_clone = syn_all.clone()
            # syn_all_clone = syn_all_clone.int()
            output_syn = embed(syn_all)
            real_loss = torch.mean(output_real.reshape(num_classes, int(output_real.shape[0] / num_classes), -1), dim=1)
            syn_loss = torch.mean(output_syn.reshape(num_classes, self.ipc, -1), dim=1)
            loss += torch.sum((real_loss - syn_loss) ** 2)
            # loss.requires_grad_(True)

            ''' Update '''
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            '''print objection loss'''
            loss_avg += loss.item()
            loss_avg /= (num_classes)
            record_loss.append(loss_avg)

            if it % 100 == 0:
                print('%s iter = %05d, loss = %.4f' % (get_time(), it, loss_avg))

            if loss_avg < best_loss:
                best_loss = loss_avg
                best_it = it
                self.z = output_syn
                self.label = label_syn
                self.syn_data = syn_all
                self.mean_data = syn_loss


class weight_estimation():
    def __init__(self, learnware):
        self.learnware_list = []
        for i, values in enumerate(learnware):
            self.learnware_list.append(values)
        self.c = len(learnware)
        self.beta = []
        self.cov = []
        self.class_num = []

    def eval(self, Phi, X, beta):
        Z = np.array(Phi)
        # beta = np.array(beta)
        v = np.zeros((Z.shape[0], X.shape[0]))
        for i in range(Z.shape[0]):
            z_i = Z[i].reshape(1, -1)
            # v[i, :] = beta[i] * rbf_kernel(z_i, X, gamma=0.01)
            v[i, :] = z_i @ X.T
        # v = np.sum(v, axis=0)
        return v

    def mixture_weight_estimate(self, X, device, specification_size):
        self.learnwares = []
        self.cov = []
        specification = []
        for i, Phi in enumerate(self.learnware_list):
            Z = Phi.syn_data.cpu()
            Z = Z.detach().numpy()
            Z = np.squeeze(Z)
            Z_tensor = torch.tensor(Z).to(device)
            self.learnwares.append(Z)
            specification.append(Z)

        # computer class_information & class's mean
        for i, Phi in enumerate(self.learnwares):
            div_flag = specification_size
            beta_i = np.ones(int(Phi.shape[0] / div_flag)) / int(Phi.shape[0] / div_flag)
            mean_data = []
            cov_data = []
            for j in range(0, Phi.shape[0], div_flag):
                mean_data.append(np.mean(Phi[j:j + div_flag], axis=0).reshape(1, -1))
                cov_data.append(np.cov(Phi[j:j + div_flag], rowvar=False))
            self.learnwares[i] = np.concatenate(mean_data, axis=0)
            self.beta.append(beta_i)
            self.cov.append(cov_data)

        C = []
        N = X.shape[0]
        for i, Phi in enumerate(self.learnwares):
            self.class_num.append(Phi.shape[0])
            v = self.eval(Phi, X, self.beta[i])
            C.append(np.sum(v, axis=1))
        C = np.concatenate(C, axis=0)
        C = C / N
        C = C.reshape(-1, 1)
        H = np.zeros((C.shape[0], C.shape[0]))
        class_all_Phi = np.concatenate(self.learnwares)
        count_i = 0
        for i, Phi in enumerate(self.learnwares):
            for j in range(Phi.shape[0]):
                Z_i = np.array(Phi[j]).reshape(1, -1)
                Z_j = np.array(class_all_Phi)
                K = Z_i @ Z_j.T
                H[count_i, :] = K
                count_i += 1
        # w_all = np.linalg.pinv(H) @ C
        # w_all = w_all / np.sum(w_all)
        # w_all = np.ravel(w_all)
        w_all, _ = solve_qp(H, C)
        w_all = np.array(w_all).reshape(-1)
        # print(w_all)
        self.w = []
        weight_estimate = np.zeros(self.c)
        count_i = 0
        for i in range(self.c):
            class_num = self.class_num[i]
            count = 0
            record = np.zeros((1, class_num))
            for j in range(class_num):
                count += w_all[count_i]
                record[0, j] = w_all[count_i]
                count_i += 1
            self.w.append(record)
            weight_estimate[i] = count
        print(weight_estimate)

        return weight_estimate, specification

    def weight_sampling_random(self, total_N, w_select):
        # count = 0
        # for i, values in enumerate(w_select):
        #     count += np.sum(values)
        # for i in range(len(w_select)):
        #     w_select[i] = w_select[i] / count
        # self.w = self.w / np.sum(self.w)
        w_select = w_select / np.sum(w_select)
        X = []
        y = []
        for i, Phi in enumerate(self.learnwares):
            if w_select[i] > 0:
                N = int(round(w_select[i] * total_N))
                if N >= 1:
                    herding_samples = self.herding_random(Phi, N, self.w[i], self.cov[i])
                    if herding_samples.shape[0] == 0:
                        continue
                    X.append(herding_samples)
                    y.extend([i] * herding_samples.shape[0])
        if len(X) > 1:
            X = np.concatenate(X, axis=0)
        elif len(X) == 1:
            X = X[0]
        y = np.array(y)
        if len(X) == 0:
            print(total_N, w_select, X)
        return (X, y)

    def herding_random(self, Phi, total_N, beta, cov):
        X = []
        beta = np.array(beta).reshape(-1, 1)
        # beta = beta / np.sum(beta)
        # beta = np.ones(Phi.shape[0]) / Phi.shape[0]
        for i in range(Phi.shape[0]):
            # N = int(beta[i] * total_N)
            N = int(total_N / len(beta))
            if N <= 0:
                continue
            X.append(np.random.multivariate_normal(Phi[i], cov[i], N))

        if len(X) > 1:
            X = np.concatenate(X, axis=0, dtype=np.float32)
        elif len(X) == 1:
            X = X[0]
        return np.asarray(X)


def solve_qp(K: np.ndarray, C: np.ndarray):
    """Solver for the following quadratic programming(QP) problem:
        - min   1/2 x^T K x - C^T x
        s.t     1^T x - 1 = 0
                - I x <= 0

    Parameters
    ----------
    K : np.ndarray
            Parameter in the quadratic term.
    C : np.ndarray
            Parameter in the linear term.

    Returns
    -------
    torch.tensor
            Solution to the QP problem.
    """
    if torch.is_tensor(K):
        K = K.cpu().numpy()
    if torch.is_tensor(C):
        C = C.cpu().numpy()
    C = C.reshape(-1)

    n = K.shape[0]
    P = np.array(K)
    P = scipy.sparse.csc_matrix(P)
    q = np.array(-C)
    G = np.array(-np.eye(n))
    G = scipy.sparse.csc_matrix(G)
    h = np.array(np.zeros((n, 1)))
    A = np.array(np.ones((1, n)))
    A = scipy.sparse.csc_matrix(A)
    b = np.array(np.ones((1, 1)))
    problem = Problem(P, q, G, h, A, b)
    solution = solve_problem(problem, solver="clarabel")
    w = solution.x
    w = torch.from_numpy(w).reshape(-1)
    return w, solution.obj
