import numpy as np

"""
Implements Sequential probability ratio test
https://en.wikipedia.org/wiki/Sequential_probability_ratio_test

"""
class SPRT:

    def __init__(self, alpha, beta, mu0, mu1):
        """
        Initilize test parameters
        """
        self.alpha = alpha
        self.beta = beta
        self.mu0 = mu0
        self.mu1 = mu1

    def get_bounds(self):
        """
        Get rejection bounds
        """
        bound1 = np.log(self.beta/(1-self.alpha)) 
        bound2 = np.log((1-self.beta)/self.alpha)
        return (bound1, bound2)


    def test(self, arr):
        """
        Test hypothesis on a given array
        Output: (-1, idx) if H0 was accepted and in which index,
                (1, idx) if H1 was accepted and in which index,
                (0, -1) continue monitoring
        """
        ratio = np.log(self.mu1/self.mu0)
        coefficient = (self.mu1 - self.mu0)/(self.mu1 * self.mu0)
        bound1, bound2 = self.get_bounds()
        prev_s = -ratio + coefficient * arr[0]
        for i in range(1, len(arr)+1):
            curr_s = prev_s - ratio + coefficient * arr[i]
            if (curr_s <= bound1):
                return (-1, i)
            elif (curr_s => bound2):
                return (1, i)
            prev_s = curr_s
        return (0, -1)