#!/usr/bin/env python
# coding: utf-8


import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
from opacus import PrivacyEngine
from torch.nn import MSELoss
from FLDP_utils import DataGenerator, cal_grad, cal_OLS, FGD, FGD_Noise_Adding, FGD_Averaged_Noise_Adding, NbAFL


def train(model, train_loader, optimizer, epoch, device, delta):
    model.train()
    criterion = torch.nn.MSELoss()
    losses = []    
    for (data, target) in train_loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()       
    return loss.item()


class TensorDataset(Dataset):
    
    def __init__(self, X,Y):
        self.X = X
        self.Y = Y
        
    def __getitem__(self, index):
        return self.X[index], self.Y[index]
    
    
    def __len__(self):
        return self.X.size(0)


p = 200
beta = np.zeros([p,1])
beta[0:10] = 1/np.sqrt(10)
#beta = np.random.uniform(size = [p,1])
#beta = beta/(np.linalg.norm(beta))
Sigma = np.eye(p)

n_list = [150,300,500,750]
K_list = [15]
eps_list = [0.8,2,5]
B = 10

for epsilon in eps_list:
    for K in K_list:
        for n in n_list:
            result = np.zeros([B,4])
            for b in range(B):
                print(b)
                N_list = [n for k in range(K)]
                X_list, Y_list = DataGenerator(N_list, beta, Sigma, sigma = 0.2, model = 'Linear')

                X = torch.Tensor(np.vstack(X_list))
                Y = torch.Tensor(np.vstack(Y_list))
                
                train_dataset = TensorDataset(X,Y)
                train_loader = DataLoader(train_dataset,
                                           batch_size = n*K,
                                           shuffle = False,
                                           num_workers = 0)

                model = torch.nn.Sequential(torch.nn.Linear(200, 1, bias = False))
                optimizer = torch.optim.SGD(model.parameters(), lr=0.015)
                privacy_engine = PrivacyEngine(accountant = 'rdp')
                model, optimizer, train_loader = privacy_engine.make_private_with_epsilon(module=model,
                                                                                           optimizer=optimizer,
                                                                                           data_loader=train_loader,
                                                                                           target_epsilon = epsilon,
                                                                                           target_delta = 1e-5,
                                                                                           epochs = 8000,
                                                                                           max_grad_norm=1.0)
                
                for epoch in range(8000):
                    loss = train(model, train_loader, optimizer, epoch, device="cpu", delta=1e-5)
                    if ((epoch+1)%1000)==0:
                        print('Iteration:',epoch+1,'-', 'loss: ', loss)
                
                beta_opa = next(model.parameters())
                beta_opa = beta_opa.data.reshape(200,1)
                
                
                
                beta_nbafl = NbAFL(X_list, Y_list, T = 20, mu = 0.001, epsilon = epsilon, delta = 1e-5, C = 1)
                
                beta_fgd = FGD(X_list,Y_list, 0.015, 5000, model = 'Linear')
                #beta_na = FGD_Noise_Adding(X_list,Y_list, 0.015, 8000, epsilon = epsilon, delta = 1e-5, C = 1, model = 'Linear')
                beta_ana = FGD_Averaged_Noise_Adding(X_list,Y_list, 0.015, 8000, epsilon = epsilon, delta = 1e-5, C = 1, model = 'Linear')

                result[b,0] = np.linalg.norm(beta_fgd - beta)
                #result[b,1] = np.linalg.norm(beta_na - beta)
                result[b,1] = np.linalg.norm(beta_ana - beta)
                result[b,2] = torch.norm(beta_opa - beta).data
                result[b,3] = np.linalg.norm(beta_nbafl - beta)
                
                
            df = pd.DataFrame(result)
            df.columns = ['fgd','ana','opacus','nbafl']
            df.to_csv('./linear_comparison_n={}_K={}_eps={}.csv'.format(n,K,epsilon))

