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

import time
import numpy as np
import pandas as pd
#import tensorflow as tf
import matplotlib.pyplot as plt


def DataGenerator(N_list, beta, Sigma, sigma = 0.5, model = 'Linear'):
    
    K = len(N_list)
    N = sum(N_list)
    p = beta.shape[0]
    beta = beta.reshape([p,1])
    
    shape = Sigma.shape[0]
    if shape!=p: 
        print('The feature dimension does not match in coeffcient and covariance !')
        return None
    
    
    X = np.random.normal(size = [N, p])
    L = np.linalg.cholesky(Sigma)
    wX = np.matmul(X, L.T)
    del X
    
    if model=='Linear':
        
        eps = sigma * np.random.normal(size = [N, 1])
        Y = np.matmul(wX,beta) + eps
        
    elif model=='Logit':
   
        u = np.matmul(wX, beta)
        prob = np.exp(u)
        prob = prob/(1+prob)
        y = np.random.uniform(size = [N,1])
        Y = 1. * (y<prob)
    
    elif model=='Poisson':
        
        u = np.matmul(wX, beta)
        prob = np.exp(u)
        Y = np.random.poisson(lam = prob.reshape(N,) ,size = [1,N]).reshape([N,1])
        
    else:
        print('Undefined Model Type!')
   

    cumlist = [0] + list(np.cumsum(N_list))
    x_list = []
    y_list = []
    for k in range(K):
        x = wX[cumlist[k]:cumlist[k+1],]
        x_list.append(x)
        y = Y[cumlist[k]:cumlist[k+1],]
        y_list.append(y)
    return x_list, y_list


def cal_OLS(X,Y):
    
    N,p = X.shape
    sxx = np.matmul(X.T, X)/N
    sxy = np.matmul(X.T, Y)/N
    b = np.matmul(np.linalg.inv(sxx),sxy)
    
    return b.reshape([p,1])


def cal_grad(x,y,beta, model = "Linear"):
    
    n,p = x.shape
    if model == 'Linear':
        grad = -np.matmul(x.T, y)/n + np.matmul(x.T, np.matmul(x,beta))/n
        return grad.reshape([p,1])  
    elif model == 'Logit':
        u = np.matmul(x, beta)
        prob = np.exp(u)
        prob = prob/(1+prob)
        resid = y-prob
        grad = - np.mean(x*resid, axis = 0).reshape([p,1])
        return grad.reshape([p,1])    
    
    elif model=="Poisson":
        
        u = np.clip(np.matmul(x, beta), a_min = -5, a_max = 100)
        prob = np.exp(u)
        resid = y-prob
        grad = -np.mean(x*resid, axis = 0).reshape([p,1])
        return grad.reshape([p,1])  
    else:
        print('Undefined Model Type!')
        return None

    
def cal_hess(x,y,beta, model = "Logit"):
    
    n,p = x.shape
    
    if model == 'Logit':
        u = np.matmul(x, beta)
        prob = np.exp(u)
        prob = prob/(1+prob)
        weight = np.sqrt(prob*(1-prob))
        wx = x*weight
        hess = np.matmul(wx.T, wx)/n
        return hess
    
    elif model=="Poisson":
        
        u = np.clip(np.matmul(x, beta), a_min = -5, a_max = 100)
        prob = np.exp(u)
        weight = np.sqrt(prob)
        wx = x*weight
        hess = np.matmul(wx.T, wx)/n
        return hess
    else:
        print('Undefined Model Type!')
        return None

    
#def cal_grad_Logistic(x,y, beta):
    
#    n,p = x.shape
#    u = np.matmul(x, beta)
#    prob = np.exp(u)
#    prob = prob/(1+prob)
#    resid = y-prob
#    grad = np.mean(x*resid, axis = 0).reshape([p,1])
    
#    return -grad


#def cal_hess_Logistic(x,y, beta):
    
#    n,p = x.shape
#    u = np.matmul(x, beta)
#    prob = np.exp(u)
#    prob = prob/(1+prob)
#    weight = np.sqrt(prob*(1-prob))
#    wx = x*weight
#    hess = np.matmul(wx.T, wx)/n
#    return hess


#def cal_grad_Poisson(x,y, beta):
    
#    n,p = x.shape
#    u = np.clip(np.matmul(x, beta), a_min = -5, a_max = 100)
#    prob = np.exp(u)
#    resid = y-prob
#    grad = np.mean(x*resid, axis = 0).reshape([p,1])
    
#    return -grad


#def cal_hess_Poisson(x,y, beta):
    
#    n,p = x.shape
#    u = np.clip(np.matmul(x, beta), a_min = -5, a_max = 100)
#    prob = np.exp(u)
#    weight = np.sqrt(prob)
#    wx = x*weight
#    hess = np.matmul(wx.T, wx)/n
#    return hess


def cal_MLE(X,Y, tol = 1e-8, model = 'Logit'):
    
    n,p = X.shape
    dist = 1.0
    beta_hat = np.zeros([p,1])
    t = 0
    while (dist>tol):
        t+=1
        grad = cal_grad(X,Y,beta_hat, model)
        hess = cal_hess(X,Y,beta_hat, model)
        
        if grad == None:
            break
        update = - np.matmul(np.linalg.inv(hess),grad)
        beta_hat = beta_hat + update
        dist = np.sum((update)**2)/p
        if (t%10)==0:
            print('Current Loop: ', t, 'Dist: ', dist)
    return beta_hat


def cal_loss_Linear(X,Y,beta):
    N, p = X.shape
    loss = (Y-np.matmul(X,beta))**2
    loss = np.mean(loss)/2
    return loss

#def cal_grad_Linear(X, Y, beta):
    
    # record sample size and feature dimension
#    N, p = X.shape
    
    # calculate noise-added gradient
#    grad = np.matmul(X.T, np.matmul(X,beta))/N - np.matmul(X.T, Y)/N
#    return grad  

def GenerateToplez(a, p):
    Sigma = np.zeros([p,p])
    for i in range(p):
        for j in range(p):
            if (a**(np.abs(i-j)))>1e-3:              
                Sigma[i,j] = a**(np.abs(i-j))
            else:
                continue
    return Sigma

def GenerateBalancedSample(N,K):
    
    averaged_n = int(np.ceil(N/K))
    
    n_list = np.zeros([K])
    for i in range(K-1):
        n_list[i] = averaged_n
    n_list[-1] = N- (K-1)*averaged_n
    
    return n_list


def GenerateNoise(p, K, n_list,c):
    
    noise = np.random.normal(size = [p,K])
    sigma = c*np.sqrt(p)/n_list
    noise_agg = np.sum(sigma*noise*(n_list/N), axis = 1).reshape([p,1])
    
    return noise_agg
    

def GenerateNoise_eq(p,K,N,c,T):
    
    noise = np.random.normal(size = [T,p])
    sigma = c*np.sqrt(p*K)/N
    noise_agg = sigma*noise
    
    return noise_agg    


def FGD(X_list,Y_list, alpha, T, model = 'Linear'):
    
    K = len(X_list)
    N_list = np.array([X_list[k].shape[0] for k in range(K)])
    weight = N_list/np.sum(N_list)
    
    p = X_list[0].shape[1]
    beta = np.zeros([p,1])
   
    for t in range(T):
        grads = np.zeros([p,1])
        for k in range(K):
            grads_k = cal_grad(X_list[k], Y_list[k], beta, model)
            # aggregation
            grads = grads + grads_k * weight[k]            
        beta = beta - alpha*grads
        if ((t+1)%5000)==0:
            print('Iteration:',t+1,'-', 'dist: ',np.linalg.norm(alpha*grads))
    return beta


def FGD_Noise_Adding(X_list,Y_list, alpha, T, epsilon, delta, C,model = 'Linear'):
    '''
    epsilon: privacy budget
    delta: privacy relaxation
    C: gradient clipping norm
    '''
    K = len(X_list)
    N_list = np.array([X_list[k].shape[0] for k in range(K)])
    weight = N_list/np.sum(N_list)
    nn = min(N_list)
    p = X_list[0].shape[1]
    beta = np.zeros([p,1])
   
    for t in range(T):
        grads = np.zeros([p,1])
        for k in range(K):
            grads_k = cal_grad(X_list[k], Y_list[k], beta, model)
            gamma_k = min(1, C/np.linalg.norm(grads_k))
            noise_level = 0.438*np.sqrt(T)*np.sqrt(np.log(1/delta))*C/(nn*epsilon)
            noise = np.random.normal(size = [p,1]) * noise_level
            # aggregation
            grads = grads + (grads_k + noise) * weight[k]            
        beta = beta - alpha*grads
        if ((t+1)%5000)==0:
            print('Iteration:',t+1,'-', 'dist: ',np.linalg.norm(alpha*grads))
    return beta

def FGD_Noise_Adding_v2(X_list,Y_list, alpha, T, epsilon, delta, C,model = 'Linear'):
    '''
    epsilon: privacy budget
    delta: privacy relaxation
    C: gradient clipping norm
    '''
    K = len(X_list)
    N_list = np.array([X_list[k].shape[0] for k in range(K)])
    weight = N_list/np.sum(N_list)
    nn = min(N_list)
    p = X_list[0].shape[1]
    beta = np.zeros([p,1])
   
    for t in range(T):
        grads = np.zeros([p,1])
        for k in range(K):
            grads_k = cal_grad(X_list[k], Y_list[k], beta, model)
            gamma_k = min(1, C/np.linalg.norm(grads_k,ord = 1))
            noise_level = np.sqrt(T)*np.sqrt(np.log(1/delta))*C/(nn*epsilon)
            noise = np.random.laplace(size = [p,1]) * noise_level
            # aggregation
            grads = grads + (grads_k + noise) * weight[k]            
        beta = beta - alpha*grads
        if ((t+1)%5000)==0:
            print('Iteration:',t+1,'-', 'dist: ',np.linalg.norm(alpha*grads))
    return beta

def FGD_Averaged_Noise_Adding(X_list,Y_list, alpha, T, epsilon, delta, C,model = 'Linear'):
    '''
    epsilon: privacy budget
    delta: privacy relaxation
    C: gradient clipping norm
    '''
    K = len(X_list)
    N_list = np.array([X_list[k].shape[0] for k in range(K)])
    weight = N_list/np.sum(N_list)
    nn = min(N_list)
    p = X_list[0].shape[1]
    beta = np.zeros([p,1])
    beta_avg = np.zeros([p,1])
    for t in range(T):
        grads = np.zeros([p,1])
        for k in range(K):
            grads_k = cal_grad(X_list[k], Y_list[k], beta, model)
            gamma_k = min(1, C/np.linalg.norm(grads_k))
            noise_level = 0.438*np.sqrt(T)*np.sqrt(np.log(1/delta))*C/(nn*epsilon)
            noise = np.random.normal(size = [p,1]) * noise_level
            # aggregation
            grads = grads + (grads_k + noise) * weight[k]            
        beta = beta - alpha*grads
        beta_avg = (beta_avg*t + beta)/(t+1)
        if ((t+1)%5000)==0:
            print('Iteration:',t+1,'-', 'dist: ',np.linalg.norm(alpha*grads))
    return beta_avg

def FGD_Averaged_Noise_Adding_v2(X_list,Y_list, alpha, T, epsilon, delta, C,model = 'Linear'):
    '''
    epsilon: privacy budget
    delta: privacy relaxation
    C: gradient clipping norm
    '''
    K = len(X_list)
    N_list = np.array([X_list[k].shape[0] for k in range(K)])
    weight = N_list/np.sum(N_list)
    nn = min(N_list)
    p = X_list[0].shape[1]
    beta = np.zeros([p,1])
    beta_avg = np.zeros([p,1])
    for t in range(T):
        grads = np.zeros([p,1])
        for k in range(K):
            grads_k = cal_grad(X_list[k], Y_list[k], beta, model)
            gamma_k = min(1, C/np.linalg.norm(grads_k, ord = 1))
            noise_level = np.sqrt(T)*np.sqrt(np.log(1/delta))*C/(nn*epsilon)
            noise = np.random.laplace(size = [p,1]) * noise_level
            # aggregation
            grads = grads + (grads_k + noise) * weight[k]            
        beta = beta - alpha*grads
        beta_avg = (beta_avg*t + beta)/(t+1)
        if ((t+1)%5000)==0:
            print('Iteration:',t+1,'-', 'dist: ',np.linalg.norm(alpha*grads))
    return beta_avg

def NbAFL(X_list, Y_list, T, mu = 0.01, epsilon = 2, delta = 1e-4, C = 1):
    
    K = len(X_list)
    N_list = np.array([X_list[k].shape[0] for k in range(K)])
    weight = N_list/np.sum(N_list)
    
    p = X_list[0].shape[1]
    beta = np.zeros([p,1])
    
    nn = min(N_list)
    for t in range(T):
        
        
        betas = np.zeros([p,1])
        for k in range(K):
            x = X_list[k]
            y = Y_list[k]
            n = N_list[k]
            sxx = np.matmul(x.T, x)/n + mu * np.eye(p)
            sxy = np.matmul(x.T, y)/n + mu * beta
            b = np.matmul(np.linalg.inv(sxx),sxy)
            b = C*b/np.linalg.norm(b)  
            noise_level = np.sqrt(2)*np.sqrt(np.log(1.25/delta))*T*C/(nn*epsilon)
            noise = np.random.normal(size = [p,1]) * noise_level
            betas = betas + b * weight[k]     
        if ((t+1)%10)==0:
            print('Iteration:',t+1,'-', 'dist: ',np.linalg.norm(beta-betas))
        beta = betas
        
    return beta
