from scipy.stats import ks_2samp
from IKS import IKS
import numpy as np
from time import time

from collections import deque

initial_A = np.random.normal(loc = 0, scale = 1, size = 500)
initial_B = np.random.normal(loc = 1, scale = 1, size = 500)

stream_A = np.random.normal(loc = 0, scale = 1, size = 5000)
stream_B = np.random.normal(loc = 1, scale = 1, size = 5000)

######################
## TEST IKS
######################

start = time()

iks_statistics = [] # collect statistics generated by IKS
iks = IKS()
sliding_A = deque(initial_A) # sliding window
sliding_B = deque(initial_B) # sliding window
for a, b in zip(initial_A, initial_B):
  iks.Add(a, 0)
  iks.Add(b, 1)

# process sliding window
for a, b in zip(stream_A, stream_B):
  iks.Remove(sliding_A.popleft(), 0)
  iks.Remove(sliding_B.popleft(), 1)

  sliding_A.append(a)
  sliding_B.append(b)

  iks.Add(a, 0)
  iks.Add(b, 1)

  iks_statistics.append(iks.KS())

finish = time()
print(f'Elapsed time for IKS to process stream: {round(finish - start, 2)} sec')

######################
## TEST ks_2samp
######################

start = time()

ks_2samp_statistics = [] # gather all statistics generated by ks_2samp
sliding_A = deque(initial_A) # sliding window
sliding_B = deque(initial_B) # sliding window

for a, b in zip(stream_A, stream_B):
  sliding_A.popleft()
  sliding_B.popleft()
  sliding_A.append(a)
  sliding_B.append(b)

  ks_2samp_statistics.append(ks_2samp(sliding_A, sliding_B).statistic)

finish = time()
print(f'Elapsed time for ks_2samp to process stream: {round(finish - start, 2)} sec')

max_diff = np.max(np.abs(np.array(iks_statistics) - np.array(ks_2samp_statistics)))

print(f'Maximum difference between IKS and ks_2samp: {max_diff}')