"""
Any MDP Task Sampler
"""
from numpy import random
from copy import deepcopy
from xenoverse.utils import pseudo_random_seed, weights_and_biases

def LinearDSSampler(state_dim:int=16,
                 action_dim:int=8,
                 observation_dim:int=8,
                 seed=None,
                 verbose=False):
    # Sampling Transition Matrix and Reward Matrix based on Irwin-Hall Distribution and Gaussian Distribution
    # Task:
    # mode: static goal or moving goal
    # ndim: number of inner dimensions
    # born_loc: born location and noise
    # sgoal_loc: static goal location, range of sink, and reward
    # pf_loc: pitfall location, range of sink, and reward
    # line_potential_energy: line potential energy specified by direction and detal_V
    # point_potential_energy: point potential energy specified by location and detal_V

    if(seed is not None):
        pseudo_random_seed(seed)
    else:
        pseudo_random_seed()

    task = dict()
    
    # static: static goal, one-step reward with reset
    # dynamic: moving goal, continuous reward
    # universal: an random reward field generated by a neural network

    task["state_dim"] = state_dim
    task["observation_dim"] = observation_dim
    task["action_dim"] = action_dim

    task["max_steps"] = random.randint(100, 1000) # At most 10-dimensional space

    AB, X = weights_and_biases(state_dim + action_dim, state_dim, need_bias=True)
    C, Y = weights_and_biases(state_dim, observation_dim, need_bias=False)

    task["ld_A"] = AB[:, :state_dim] * random.choice([0.01, 0.02, 0.05, 0.1, 0.2])
    task["ld_B"] = AB[:, state_dim:]
    task["ld_C"] = C
    task["ld_X"] = X
    task["ld_Y"] = Y

    # Sample rewards
    task["action_cost"] = max(random.uniform(-0.05, 0.05), 0.0)
    task["reward_base"] = random.exponential(scale=0.10) + task["action_cost"]
    task["terminate_punish"] = random.uniform(1.0, 100.0)


    # probability without any procedural reward
    task["reward_factor"] = min(random.exponential(scale=0.01), 0.1 * task["reward_base"])

    born_loc = int(max(random.exponential(scale=1.0), 1))
    task["initial_states"] = [random.randn(state_dim) for _ in range(born_loc)]


    return task