import numpy as np

class Sobol:
    def __init__(self, dim: int, seed: int=None):
        self._dim = None

        self._atmost = 2**self.log_max - 1
        self._maxcol = self._i4_bit_hi1(self.atmost)

        self.dim = dim
        self.seed = -1
    
    @property
    def dim(self):
        return self._dim
    
    @dim.setter
    def dim(self, dim):
        self.update_dim(dim)

    @property
    def dim_max(self):
        return 40
    
    @property
    def log_max(self):
        return 30

    @property
    def atmost(self):
        return self._atmost

    @property
    def maxcol(self):
        return self._maxcol
    
    @property
    def poly(self):
        return self._poly
    
    def update_dim(self, dim):
        if not hasattr(self, '_dim') \
             or (self._dim is None) \
             or (self._dim != dim):
             
            self._init_poly()
            self._init_v()
            self._update_v(dim)
            self._dim = dim

    def update_seed(self, seed):
        v = self._v
        dim_num = self.dim

        if seed <= 0:
            l = 1
            self.lastq = np.zeros(dim_num)

        elif seed == self.seed + 1:
            #
            #  Find the position of the right-hand zero in SEED.
            #
            l = self._i4_bit_lo0(seed)

        elif seed <= self.seed:
            self.seed = 0
            l = 1
            self.lastq = np.zeros(dim_num)

            for seed_temp in range(int(self.seed), int(seed)):
                l = self._i4_bit_lo0(seed_temp)
                for i in range(1, dim_num + 1):
                    self.lastq[i - 1] = np.bitwise_xor(int(self.lastq[i - 1]), int(v[i - 1, l - 1]))

            l = self._i4_bit_lo0(seed)

        elif self.seed + 1 < seed:
            
            for seed_temp in range(int(self.seed + 1), int(seed)):
                l = self._i4_bit_lo0(seed_temp)
                for i in range(1, dim_num + 1):
                    self.lastq[i - 1] = np.bitwise_xor(int(self.lastq[i - 1]), int(v[i - 1, l - 1]))

            l = self._i4_bit_lo0(seed)
        
        self.seed = seed

        #
        #  Check that the user is not calling too many times!
        #
        if self.maxcol < l:
            raise RuntimeError("\n".join([
                f"i4_sobol - Fatal error!",
                f" Too many calls!",
                f" MAXCOL = {self.maxcol}\n",
                f" L =   {l}\n"
            ]))

        return l
        

    def tau(self, dim: int=None):

        # *****************************************************************************80
        #
        ## tau_sobol() defines favorable starting seeds for Sobol sequences.
        #
        #  Discussion:
        #
        #    For spatial dimensions 1 through 13, this routine returns
        #    a "favorable" value TAU by which an appropriate starting point
        #    in the Sobol sequence can be determined.
        #
        #    These starting points have the form N = 2**K, where
        #    for integration problems, it is desirable that
        #      TAU + DIM_NUM - 1 <= K
        #    while for optimization problems, it is desirable that
        #      TAU < K.
        #
        #  Licensing:
        #
        #    This code is distributed under the MIT license.
        #
        #  Modified:
        #
        #    22 February 2011
        #
        #  Author:
        #
        #    Original FORTRAN77 version by Bennett Fox.
        #    MATLAB version by John Burkardt.
        #    Python version by Corrado Chisari
        #
        #  Reference:
        #
        #    IA Antonov, VM Saleev,
        #    USSR Computational Mathematics and Mathematical Physics,
        #    Volume 19, 1980, pages 252 - 256.
        #
        #    Paul Bratley, Bennett Fox,
        #    Algorithm 659:
        #    Implementing Sobol's Quasirandom Sequence Generator,
        #    ACM Transactions on Mathematical Software,
        #    Volume 14, Number 1, pages 88-100, 1988.
        #
        #    Bennett Fox,
        #    Algorithm 647:
        #    Implementation and Relative Efficiency of Quasirandom
        #    Sequence Generators,
        #    ACM Transactions on Mathematical Software,
        #    Volume 12, Number 4, pages 362-376, 1986.
        #
        #    Stephen Joe, Frances Kuo
        #    Remark on Algorithm 659:
        #    Implementing Sobol's Quasirandom Sequence Generator,
        #    ACM Transactions on Mathematical Software,
        #    Volume 29, Number 1, pages 49-57, March 2003.
        #
        #    Ilya Sobol,
        #    USSR Computational Mathematics and Mathematical Physics,
        #    Volume 16, pages 236-242, 1977.
        #
        #    Ilya Sobol, YL Levitan,
        #    The Production of Points Uniformly Distributed in a Multidimensional
        #    Cube (in Russian),
        #    Preprint IPM Akad. Nauk SSSR,
        #    Number 40, Moscow 1976.
        #
        #  Input:
        #
        #    integer DIM_NUM, the spatial dimension. Only values
        #    of 1 through 13 will result in useful responses.
        #
        #  Output:
        #
        #    integer TAU, the value TAU.
        #
        dim = dim if dim is not None else self.dim

        dim_max = 13
        tau_table = [0, 0, 1, 3, 5, 8, 11, 15, 19, 23, 27, 31, 35]

        if 1 <= dim and dim <= dim_max:
            tau = tau_table[dim]
        else:
            tau = -1

        return tau

    def _i4_bit_hi1(self, n: int):

        # *****************************************************************************80
        #
        ## i4_bit_hi1() returns the position of the high 1 bit base 2 in an I4.
        #
        #  Discussion:
        #
        #    An I4 is an integer ( kind = 4 ) value.
        #
        #  Example:
        #
        #       N    Binary    Hi 1
        #    ----    --------  ----
        #       0           0     0
        #       1           1     1
        #       2          10     2
        #       3          11     2
        #       4         100     3
        #       5         101     3
        #       6         110     3
        #       7         111     3
        #       8        1000     4
        #       9        1001     4
        #      10        1010     4
        #      11        1011     4
        #      12        1100     4
        #      13        1101     4
        #      14        1110     4
        #      15        1111     4
        #      16       10000     5
        #      17       10001     5
        #    1023  1111111111    10
        #    1024 10000000000    11
        #    1025 10000000001    11
        #
        #  Licensing:
        #
        #    This code is distributed under the GNU LGPL license.
        #
        #  Modified:
        #
        #    26 October 2014
        #
        #  Author:
        #
        #    John Burkardt
        #
        #  Input:
        #
        #    integer N, the integer to be measured.
        #    N should be nonnegative.  If N is nonpositive, the function
        #    will always be 0.
        #
        #  Output:
        #
        #    integer BIT, the position of the highest bit.
        #
        i = n
        bit = 0

        while True:

            if i <= 0:
                break

            bit += 1
            i = i // 2

        return bit

    def _i4_bit_lo0(self, n: int):

        # *****************************************************************************80
        #
        ## i4_bit_lo0() returns the position of the low 0 bit base 2 in an I4.
        #
        #  Discussion:
        #
        #    An I4 is an integer ( kind = 4 ) value.
        #
        #  Example:
        #
        #       N    Binary    Lo 0
        #    ----    --------  ----
        #       0           0     1
        #       1           1     2
        #       2          10     1
        #       3          11     3
        #       4         100     1
        #       5         101     2
        #       6         110     1
        #       7         111     4
        #       8        1000     1
        #       9        1001     2
        #      10        1010     1
        #      11        1011     3
        #      12        1100     1
        #      13        1101     2
        #      14        1110     1
        #      15        1111     5
        #      16       10000     1
        #      17       10001     2
        #    1023  1111111111    11
        #    1024 10000000000     1
        #    1025 10000000001     2
        #
        #  Licensing:
        #
        #    This code is distributed under the GNU LGPL license.
        #
        #  Modified:
        #
        #    08 February 2018
        #
        #  Author:
        #
        #    John Burkardt
        #
        #  Input:
        #
        #    integer N, the integer to be measured.
        #    N should be nonnegative.
        #
        #  Output:
        #
        #    integer BIT, the position of the low 1 bit.
        #
        bit = 0
        i = n

        while True:
            bit += 1
            i2 = i // 2

            if i == 2 * i2:
                break

            i = i2

        return bit

    def _init_params(self):
        self._dim = None

    def _init_v(self):
        #
        #  Initialize (part of) V.
        #

        v = np.zeros((self.dim_max, self.log_max))

        v[0:40, 0] = np.transpose([
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        ])

        v[2:40, 1] = np.transpose([
            1, 3, 1, 3, 1, 3, 3, 1,
            3, 1, 3, 1, 3, 1, 1, 3, 1, 3,
            1, 3, 1, 3, 3, 1, 3, 1, 3, 1,
            3, 1, 1, 3, 1, 3, 1, 3, 1, 3,
        ])

        v[3:40, 2] = np.transpose([
            7, 5, 1, 3, 3, 7, 5,
            5, 7, 7, 1, 3, 3, 7, 5, 1, 1,
            5, 3, 3, 1, 7, 5, 1, 3, 3, 7,
            5, 1, 1, 5, 7, 7, 5, 1, 3, 3,
        ])

        v[5:40, 3] = np.transpose([
            1,  7,  9, 13, 11,
            1,  3,  7,  9,  5, 13, 13, 11,  3, 15,
            5,  3, 15,  7,  9, 13,  9,  1, 11,  7,
            5, 15,  1, 15, 11,  5,  3,  1,  7,  9,
        ])

        v[7:40, 4] = np.transpose([
             9,  3, 27,
            15, 29, 21, 23, 19, 11, 25,  7, 13, 17,
             1, 25, 29,  3, 31, 11,  5, 23, 27, 19,
            21,  5,  1, 17, 13,  7, 15,  9, 31,  9,
        ])

        v[13:40, 5] = np.transpose([
            37, 33,  7,  5, 11, 39, 63,
            27, 17, 15, 23, 29,  3, 21, 13, 31, 25,
             9, 49, 33, 19, 29, 11, 19, 27, 15, 25,
        ])

        v[19:40, 6] = np.transpose([
            13,
            33, 115, 41, 79, 17,  29, 119, 75, 73, 105,
             7,  59, 65, 21,  3, 113,  61, 89, 45, 107,
        ])

        v[37:40, 7] = np.transpose([
            7, 23, 39,
        ])

        #
        #  Find the number of bits in ATMOST.
        #
        #maxcol = self._i4_bit_hi1(self.atmost)
        #
        #  Initialize row 1 of V.
        #
        v[0, 0:self.maxcol] = 1

        self._v = v

        return self._v

    def _init_poly(self):
        self._poly = np.array([
              1,   3,   7,  11,  13,  19,  25,  37,  59,  47,
             61,  55,  41,  67,  97,  91, 109, 103, 115, 131,
            193, 137, 145, 143, 241, 157, 185, 167, 229, 171,
            213, 191, 253, 203, 211, 239, 247, 285, 369, 299,
        ])
        return self._poly
    
    def _update_v(self, dim_num):
        v = self._v
        #
        #  Check parameters.
        #
        if dim_num < 1 or self.dim_max < dim_num:
            raise RuntimeError("\n".join([
                f"i4_sobol - Fatal error!",
                f" The spatial dimension DIM_NUM should satisfy:",
                f"  1 <= DIM_NUM <= {self.dim_max}",
                f" But this input value is DIM_NUM = {dim_num}",
            ]))
        #
        #  Initialize the remaining rows of V.
        #
        for i in range(2, dim_num + 1):
            #
            #  The bits of the integer POLY(I) gives the form of polynomial I.
            #
            #  Find the degree of polynomial I from binary encoding.
            #
            j = self.poly[i - 1]
            m = 0

            while True:
                #j = np.floor(j / 2.0)
                j //= 2
                if j <= 0:
                    break
                m = m + 1
            #
            #  Expand this bit pattern to separate components of the logical array INCLUD.
            #
            j = self.poly[i - 1]
            includ = np.zeros(m)
            for k in range(m, 0, -1):
                #j2 = np.floor(j / 2.0)
                j2 = j // 2
                includ[k - 1] = j != (2 * j2)
                j = j2
            #
            #  Calculate the remaining elements of row I as explained
            #  in Bratley and Fox, section 2.
            #
            for j in range(m + 1, self.maxcol + 1):
                newv = v[i - 1, j - m - 1]
                l = 1
                for k in range(1, m + 1):
                    l = 2 * l
                    if includ[k - 1]:
                        newv = np.bitwise_xor(int(newv), int(l * v[i - 1, j - k - 1]))
                v[i - 1, j - 1] = newv
        #
        #  Multiply columns of V by appropriate power of 2.
        #
        l = 1
        for j in range(self.maxcol - 1, 0, -1):
            l = 2 * l
            v[0:dim_num, j - 1] = v[0:dim_num, j - 1] * l
        #
        #  RECIPD is 1/(common denominator of the elements in V).
        #
        self.recipd = 1.0 / (2 * l)
        self.lastq = np.zeros(dim_num)

        self._v = v

        return self._v

    def _i4_sobol(self, dim=None, seed=None):

        # *****************************************************************************80
        #
        ## i4_sobol() generates a new quasirandom Sobol vector with each call.
        #
        #  Discussion:
        #
        #    The routine adapts the ideas of Antonov and Saleev.
        #
        #  Licensing:
        #
        #    This code is distributed under the MIT license.
        #
        #  Modified:
        #
        #    22 February 2011
        #
        #  Author:
        #
        #    Original FORTRAN77 version by Bennett Fox.
        #    MATLAB version by John Burkardt.
        #    Python version by Corrado Chisari
        #
        #  Reference:
        #
        #    Antonov, Saleev,
        #    USSR Computational Mathematics and Mathematical Physics,
        #    olume 19, 1980, pages 252 - 256.
        #
        #    Paul Bratley, Bennett Fox,
        #    Algorithm 659:
        #    Implementing Sobol's Quasirandom Sequence Generator,
        #    ACM Transactions on Mathematical Software,
        #    Volume 14, Number 1, pages 88-100, 1988.
        #
        #    Bennett Fox,
        #    Algorithm 647:
        #    Implementation and Relative Efficiency of Quasirandom
        #    Sequence Generators,
        #    ACM Transactions on Mathematical Software,
        #    Volume 12, Number 4, pages 362-376, 1986.
        #
        #    Ilya Sobol,
        #    USSR Computational Mathematics and Mathematical Physics,
        #    Volume 16, pages 236-242, 1977.
        #
        #    Ilya Sobol, Levitan,
        #    The Production of Points Uniformly Distributed in a Multidimensional
        #    Cube (in Russian),
        #    Preprint IPM Akad. Nauk SSSR,
        #    Number 40, Moscow 1976.
        #
        #  Input:
        #
        #    integer DIM_NUM, the number of spatial dimensions.
        #    DIM_NUM must satisfy 1 <= DIM_NUM <= 40.
        #
        #    integer SEED, the "seed" for the sequence.
        #    This is essentially the index in the sequence of the quasirandom
        #    value to be generated.
        #
        #  Output:
        #
        #    real QUASI(DIM_NUM), the next quasirandom vector.
        #
        #    integer SEED, an updated seed.
        #

        dim_num = dim if dim is not None else self.dim

        #
        #  Things to do only if the dimension changed.
        #
        if dim_num != self.dim:
            warnings.warn(f"dimension has been changed from {self.dim} to {dim_num}.")
            self.update_dim(dim_num)

        seed = int(np.floor(seed))

        l = self.update_seed(seed)
        
        #
        #  Calculate the new components of QUASI.
        #
        v = self._v
        quasi = np.zeros(dim_num)
        for i in range(1, dim_num + 1):
            quasi[i - 1] = self.lastq[i - 1] * self.recipd
            self.lastq[i - 1] = np.bitwise_xor(int(self.lastq[i - 1]), int(v[i - 1, l - 1]))

        next_seed = self.seed + 1

        return [quasi, next_seed]


    def generate(self, n: int, skip: int=None):

        # *****************************************************************************80
        #
        ## i4_sobol_generate() generates a Sobol dataset.
        #
        #  Licensing:
        #
        #    This code is distributed under the MIT license.
        #
        #  Modified:
        #
        #    22 February 2011
        #
        #  Author:
        #
        #    Original MATLAB version by John Burkardt.
        #    Python version by Corrado Chisari
        #
        #  Input:
        #
        #    integer M, the spatial dimension.
        #
        #    integer N, the number of points to generate.
        #
        #    integer SKIP, the number of initial points to skip.
        #
        #  Output:
        #
        #    real R(M,N), the points.
        #
        m = self.dim
        skip = skip if skip is not None else self.tau()

        r = np.zeros((m, n))
        seed = skip - 1
        for j in range(n):
            r[0:m, j], seed = self._i4_sobol(seed=seed)


        return r