import numpy as np
import scipy.io as sio
from numpy.random import Generator, MT19937, PCG64, PCG64DXSM, Philox, SFC64

def rng_gen(seed = -1, random_type = 'MT19937'):
    r"""
    Obtain random number generator according to random number seed and random number generator type
    Parameters: 
    seed: int, optional
    Random number seed which be set 
    random_type: str, optional
    Selected random number generator type

    Returns:
    rng: Generator
    The random number generator obtained by specifying the random tree seed and generator type
    """
    if random_type == 'MT19937':
        if seed != -1:
            rng = Generator(MT19937(seed))

        else:
            rng = Generator(MT19937())
    elif random_type == 'PCG64':
        if seed != -1:
            rng = Generator(PCG64(seed))

        else:
            rng = Generator(PCG64())
    elif random_type == 'PCG64DXSM':
        if seed != -1:
            rng = Generator(PCG64DXSM(seed))

        else:
            rng = Generator(PCG64DXSM())

    elif random_type == 'Philox':
        if seed != -1:
            rng = Generator(Philox(seed))

        else:
            rng = Generator(Philox())

    elif random_type == 'SFC64':
        if seed != -1:
            rng = Generator(SFC64(seed))

        else:
            rng = Generator(SFC64())
    
    return rng

def sym_gen(num, base = 1, bit_load = 0, seed = -1, random_type = 'MT19937', **kwargs):
    r""" 
    Generate data sequence according to input conditions
    Parameters:
    num: int
    Length of data sequence to be generated
    base: int, optional
    Represents the maximum value of the generated data value. The default value is 1, representing '0-1' bits
    bit_load: bool, optional
    Boolean variable representing whether to read data from the saved data file
    seed: int, optional
    Random number seed which be set
    random_type: str, optional
    Selected random number generator type

    Returns:
    sym_seq: ndarray
    Generated data sequence
    """
    if bit_load:
        data_path = kwargs.get("data_path", '/home/ospan/code/Niu/data/random_num.mat')
        data = sio.loadmat(data_path)[random_type][0:num]
        sym_seq = data.astype(np.int64).reshape(-1)
        return sym_seq
    else:
        rng = rng_gen(seed = seed, random_type = random_type)
        sym_seq = rng.integers(low=0, high=base+1, size=num)
        return sym_seq
    

