text: #-----------------------------------------------------------------------------#
  #
  #  Two-Sample Test with  Nearest-Neighbors Density Ratio Estimation
  #
  #  MIT License
  #  Copyright (c) 2018: Andrea De Simone, Thomas Jacques
  #
  #-----------------------------------------------------------------------------#
  
  from __future__ import print_function
  from __future__ import division
  import numpy as np
  from matplotlib import pyplot as plt
  import seaborn as sns
  from random import randint
  from scipy import stats
  from sklearn.utils import shuffle
  from sklearn.neighbors import NearestNeighbors
  
  #-----------------------------------------------------------------------------#
  # Nearest Neighbors class
  class NearestNeighborsClass(object):
      """
      Nearest neighbor density ratio estimator.
      Used to compute the Test Statistic.
      """
      
      def __init__(self, n_neighbors=2):
          """
          Initialize the class instance.
          n_neighbors: number of nearest neighbors (must be >= 1)
          """
          self.n_neighbors = n_neighbors
      
      
      def fit(self, x_benchmark, x_trial):
          """
          Build kd-trees for both samples.
          x_benchmark: benchmark (control) sample
          x_trial: trial (test) sample
          """
          
          # build kd-trees for both samples
          self.nbrs_benchmark = NearestNeighbors(n_neighbors=self.n_neighbors,
                                                 algorithm='kd_tree').fit(x_benchmark)
          self.nbrs_trial     = NearestNeighbors(n_neighbors=self.n_neighbors+1,
                                                 algorithm='kd_tree').fit(x_trial)
  
  
      def compute_distances(self, x_q):
          """
          Compute distances of Kth-NN in B and T from query points x_q.
          x_q: sample of query data points at which we evaluate the density ratio.
          Return: array of distances of Kth-NN in Benchmark and Trial,
                  and array of coordinates (X_q) for each data point in the query sample.
          """
          # get distances of nearest neighbors in Benchmark from Query points
          distances_B, _ = self.nbrs_benchmark.kneighbors(x_q)
          radii_B = distances_B[:,-1] # distances of Kth-NN
                      
          # get distances of nearest neighbors in Trial from Query points
          distances_T, _ = self.nbrs_trial.kneighbors(x_q)
          radii_T = distances_T[:,-1] # distances of Kth-NN
                              
          coords = x_q
                                  
          return(radii_B, radii_T, coords)
  #*****************************************************************************#
  
  
  #-----------------------------------------------------------------------------#
  # Hypothesis Test class
  class HypothesisTest(object):
      """
      Class to perform hypothesis testing (two-sample test) with
      Nearest Neighbors Density Ratio Estimation of pdfs.
      """
      
      def __init__(self, x_benchmark, x_trial, K=5, n_perm=100):
          """
          Initialize the class instance.
          """
          self.x_benchmark = x_benchmark
          self.x_trial = x_trial
          self.K = K
          self.n_perm = n_perm
          
          # Number of points in the samples
          self.NB = self.x_benchmark.shape[0]
          self.NT = self.x_trial.shape[0]
          
          # Dimensionality of points
          self.D = self.x_benchmark.shape[1]
      
      
      def TestStatistic(self, x_b, x_t):
          """
          Compute the test statistic using nearest neighbor
          density ratio estimator.
          """
          # Instantiate NearestNeighbors class
          NN = NearestNeighborsClass(n_neighbors = self.K)
          
          # Build kd-trees with fixed K
          NN.fit(x_b, x_t)
          
          # Compute distances r_{j,B}, r_{j,T} of Kth-NN in B and T,
          # from x_j in Trial
          self.r_B, self.r_T, _ = NN.compute_distances(x_t)
          
          # Compute estimated density ratio on Trial points
          r_hat = np.power(np.divide(self.r_B, self.r_T), self.D) * (self.NB/float(self.NT-1))
          
          # Compute test statistic over Trial points
          TS =  np.mean( np.log(r_hat) )
          
          return(TS)
      
      
      def PermutationTest(self):
          """
          Permutation Test to reconstruct the distribution of the test statistic.
          Called in 'compute_pvalue' method.
          It stores in 'self.TS_tilde' the set of TS computed at each permutation.
          """
          # U = union of B and T
          union_sample = np.concatenate((self.x_benchmark, self.x_trial), axis=0)
          n_samples = self.NB + self.NT
          
          # Initialize array of test statistic values
          self.TS_tilde = np.zeros(self.n_perm, dtype=np.float)
          
          count=0
          print("Running {:d} Permutations... 0%".format(self.n_perm))
          
          # loop over different samplings
          for i in range(self.n_perm):
              
              # Print progress
              progress = int(round(((i+1)/self.n_perm)*100,0))
              progress_list = [25, 50, 75, 100]
              if count < len(progress_list) and progress == progress_list[count]:
                  count+=1
                  print("Running {:d} Permutations... {:d}%".format(self.n_perm, progress))
              
              # Random permutations of U (sampling without replacement)
              x_resampled = shuffle(union_sample)
              # Assign first NB elements to Benchmark
              B_resampled = x_resampled[:self.NB]
              # Assign remaning NT elements to Trial
              T_resampled = x_resampled[self.NB:]
      
              # Compute the test statistic
              self.TS_tilde[i]  = self.TestStatistic(B_resampled, T_resampled)
  
  
      def compute_pvalue(self):
          """
          Perform permutation test and compute the (two-sided) p-value
          of the observed value of the test statistic.
          If p-value from permutations is 0, compute the p-value
          from gaussian fit to distribution (pvalue_gaussian).
          """
          # Run permutation test
          self.PermutationTest()
          # TS obtained from the original B,T samples
          self.TS_obs = self.TestStatistic(self.x_benchmark, self.x_trial)
                      
          # Mean and std of the TS distribution
          self.mu = np.mean(self.TS_tilde)
          self.sigma = np.std(self.TS_tilde)
                  
          # Standardized test statistic (zero mean, unit variance)
          self.TS_prime = (self.TS_tilde - self.mu)/self.sigma
          self.TS_prime_obs = (self.TS_obs - self.mu)/self.sigma
                          
          # Two-sided p-value from TS' distribution
          self.p_value = 2*(1 - 0.01 * stats.percentileofscore(self.TS_prime,
                                                              abs(self.TS_prime_obs)) )
                              
          # if 0, compute it from standard normal
          if self.p_value == 0.0:
              self.p_value = self.pvalue_gaussian()
                  
          print("")
          print("p-value  = {:e}".format(self.p_value))
  
  
      def pvalue_gaussian(self):
          """
          Two-sided p-value from stanard normal distribution (zero-mean unit-variance)
          of the standardized TS disitribution from permutations.
          The 'compute_pvalue' method needs to be evaluated first.
          """
              
          pv = 2 * stats.norm.sf(abs(self.TS_prime_obs), loc=0, scale=1)
          return(pv)
  
  
      def pvalue_test(self, alpha=0.01):
          """
          Compare p-value with alpha.
          The 'compute_pvalue' method needs to be evaluated first.
          """
          CL = int((1-alpha)*100)  # confidence level
                  
          if self.p_value < alpha:
              print("Null hypothesis rejected at {:d}%CL => distributions are different".format(CL))
          else:
              print("Null hypothesis NOT rejected => distributions are the same")
                          
  
      def significance(self):
          """
          Convert two-sided p-value to equivalent significance Z 
          of a standard normal distribution: Z=InverseCDF(1-p/2),
          or equivalently Z=-InvserCDF(p/2).
          The 'compute_pvalue' method needs to be evaluated first.
          """
              
          Z = -1*stats.norm.ppf(0.5*self.p_value)
          print("Z = {:e} sigma".format(Z))
  #*****************************************************************************#
  
  #-----------------------------------------------------------------------------#
  # Gaussian Data class
  class GaussianData(object):
      """
      Class to generate 2-dimensional Benchmark and Trial
      samples from gaussian distributions
      """
      
      def __init__(self, mean_b, cov_b, mean_t, cov_t, n_points=1000):
          """
          Initialize the class instance.
          Input: mean vectors and covariance matrices for B and T samples.
          n_points: number of sample points to generate.
          """
          self.mean_b = mean_b
          self.cov_b  = cov_b
          self.mean_t = mean_t
          self.cov_t  = cov_t
          self.n_points = n_points
      
      
      def generate_data(self):
          """
          Generate toy data from two-dim Gaussian distributions
          """
          np.random.seed(0)
          
          L_b = np.linalg.cholesky(self.cov_b)
          L_t = np.linalg.cholesky(self.cov_t)
          
          self.x_benchmark = np.dot(L_b,(np.random.randn(self.n_points,2) + self.mean_b).T).T
          self.x_trial = np.dot(L_t,(np.random.randn(self.n_points,2) + self.mean_t).T).T
      
      
      def scatter_plot(self):
          """
          Generate plots of benchmark/trial samples in 2D feature space
          """
          sns.set_style('whitegrid')
          
          fig, ax = plt.subplots()
          cmap = sns.cubehelix_palette(8, start=.5, rot=-.75, as_cmap=True)
          
          
          plt.title('Benchmark and Trial Samples', fontsize=16)
          
          ax.xaxis.set_tick_params(labelsize=16, direction='inout', length=6, width=1, color='gray')
          ax.yaxis.set_tick_params(labelsize=16, direction='inout', length=6, width=1, color='gray')
          
          ax.scatter(self.x_benchmark[:,0], self.x_benchmark[:,1], c='magenta',
                     alpha=0.5, marker='x',label='B sample')
          ax.scatter(self.x_trial[:,0],self.x_trial[:,1], c='blue',
                     alpha=0.2, marker='s',label='T sample')
                     
          plt.grid(True)
          plt.legend(loc='upper left', fontsize=14)
          plt.show()
  #*****************************************************************************#
  
  #-----------------------------------------------------------------------------#
  def main():
  
      # Generate gaussian random samples
      # Benchmark sample
      mean_b = np.array([1.0,1.0])
      cov_b  = np.array([[1,0.],[0.,1]])
      # Trial sample
      mean_t = np.array([1.15,1.15])
      cov_t  = np.array([[1.0,0.0],[0.0,1.0]])
  
      gdata = GaussianData(mean_b, cov_b, mean_t, cov_t, n_points=20000)
      gdata.generate_data()
      #gdata.scatter_plot()
  
      X_benchmark = gdata.x_benchmark
      X_trial    = gdata.x_trial
  
      print("Benchmark sample: {:d} points".format(X_benchmark.shape[0]))
      print("Trial sample: {:d} points\n".format(X_trial.shape[0]))
      
      
      
      
      # Compute Test Statistic and p-value
      K=5
      Nperm=100
  
      HT = HypothesisTest(X_benchmark, X_trial, K=K, n_perm=Nperm)
  
      print("TS_obs (K={:d})  = {:.5f}\n".format(K, HT.TestStatistic(X_benchmark, X_trial)) )
  
  
      HT.compute_pvalue()
  
      HT.significance()
      
      HT.pvalue_test(alpha=0.01)
      
  #*****************************************************************************#
  
  
  
  #-----------------------------------------------------------------------------#
  if __name__ == '__main__':
      main()
      print("")
  #*****************************************************************************#
  