#!/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, NbAFL



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
ratio_list = [1, 3, 5]
epsilon = 2
B = 80

for ratio in ratio_list:
    for n in n_list:
        result = np.zeros([B,4])
        for b in range(B):
            print(b)
            
            d = np.clip((n + (np.sqrt(n)*ratio)*np.random.normal(size = [int((K-1)/2)])),5,2*n-5).astype(int)
            e = 2*n - d
            f = np.array([n])
            N_list = np.hstack([d,e,f])
            
            X_list, Y_list = DataGenerator(N_list, beta, Sigma, sigma = 0.2, model = 'Linear')

            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')
            beta_nbafl = NbAFL(X_list, Y_list, T = 20, mu = 0.001, epsilon = epsilon, delta = 1e-5, C = 1)

            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)
            result[b,3] = np.linalg.norm(beta_nbafl - beta)

        df = pd.DataFrame(result)
        df.columns = ['fgd','na','ana','nbafl']
        df.to_csv('./linear_unbalance_n={}_ratio={}.csv'.format(n,ratio))

