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

import numpy as np
import pandas as pd
from FLDP_utils import DataGenerator, cal_grad, cal_OLS, FGD, FGD_Noise_Adding, FGD_Averaged_Noise_Adding



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 = 15
eps_list = [0.5,0.8,1] 
B = 50

for epsilon in eps_list:
    for n in n_list:
        result = np.zeros([B,3])
        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 = 'Poisson')


            beta_fgd = FGD(X_list,Y_list, 0.01, 10000, model = 'Poisson')
            beta_na = FGD_Noise_Adding(X_list,Y_list, 0.01, 10000, epsilon = epsilon, delta = 1e-5, C = 1, model = 'Poisson')
            beta_ana = FGD_Averaged_Noise_Adding(X_list,Y_list, 0.01, 20000, epsilon = epsilon, delta = 1e-5, C = 1, model = 'Poisson')
            
            result[b,0] = np.linalg.norm(beta_fgd - beta)
            result[b,1] = np.linalg.norm(beta_na - beta)
            result[b,2] = np.linalg.norm(beta_ana - beta)
            
        df = pd.DataFrame(result)
        df.columns = ['fgd','na','ana']
        df.to_csv('./poisson_sparse_n={}_K={}_eps={}.csv'.format(n,K,epsilon))

