import torch, copy, pandas as pd, numpy as np
from torch import nn
from sklearn.utils import shuffle
import numpy as np
from torch.utils.data import TensorDataset, DataLoader
from sklearn.metrics import accuracy_score
import warnings
warnings.filterwarnings('ignore')

#its-og-login1.syr.edu ALTERNATE
# def get_partition(x):
#   """
#
#   Partition one dataset into three partitions: private, public, and test
#   """
#   np.random.seed(0)
#   rnd = np.random.rand()
#   if rnd<=0.8:
#     return 'private'
#   elif rnd<=0.9:
#     return 'public'
#   else:
#     return 'test'



def get_partition(pd00, num_public):
    pd00 = shuffle(pd00)
    pd00['index']  = pd00.index
    public_idx = np.random.choice(range(len(pd00)), num_public, replace = False)
    public_pd = pd00[pd00['index'].isin(public_idx)]
    public_pd['partition'] = 'public'
    res_pd = pd00[~pd00['index'].isin(public_idx)]
    res_pd['partition'] = res_pd['index'].apply(lambda x:  'test' if np.random.rand()>=0.8  else 'private')
    pd00 = shuffle( pd.concat([res_pd, public_pd])).drop(columns = ['index'])
    return pd00


def rand_response(z, num_z, eps):

  """
  Implement randomized response mechanism 
  """
  main_prob = np.exp(eps)/(num_z -1 + np.exp(eps))
  aux_prob = (1 - main_prob)/float(num_z -1)
  pr = [aux_prob]  * num_z
  pr[int(z)] = main_prob
  return np.random.choice(list(range(num_z)), p= pr)