from types import SimpleNamespace
from dataclasses import asdict, dataclass


@dataclass
class DefaultCfg:
    num_objective: int = 2
    safe_obj_list: tuple = (False, False)
    time_steps: int = 500000
    weight_bc_loss: float = 200.0
    tau: float = 0.005
    policy_noise: float = 0.2
    noise_clip: float = 0.5
    eval_episodes: int = 10 #5 
    max_episode_len: int = None
    layer_N_critic: int = 2
    layer_N_actor: int = 2
    hidden_size: int = 400
    cuda: bool = True
    w_step_size_final_eval: float = 0.01
    w_step_size: float = 0.05

@dataclass
class D4MORL_HalfCheetah_Cfg(DefaultCfg):
    max_episode_len: int = 500

@dataclass
class D4MORL_Walker2d_Cfg(DefaultCfg):
    max_episode_len: int = 500

@dataclass
class D4MORL_Swimmer_Cfg(DefaultCfg):
    max_episode_len: int = 500
    weight_bc_loss: float = 1000.0

@dataclass
class D4MORL_Hopper_Cfg(DefaultCfg):
    max_episode_len: int = 500

@dataclass
class D4MORL_Ant_Cfg(DefaultCfg):
    max_episode_len: int = 500

@dataclass
class DSRL_Cfg(DefaultCfg):
    safe_obj_list: tuple = (False, True)
    eval_episodes: int = 10
    weight_bc_loss: float = 20.0

@dataclass
class DSRL_RUN_Cfg(DefaultCfg):
    safe_obj_list: tuple = (False, True)
    eval_episodes: int = 10
    weight_bc_loss: float = 20.0
    time_steps: int = 500000
    tau: float = 0.001

@dataclass
class CMO_Cfg(DefaultCfg):
    safe_obj_list: tuple = (False, False, True)
    num_objective: int = 3
    w_step_size_final_eval: float = 0.1
    w_step_size: float = 0.1

HYPERPARAMS = {
    'MO-Hopper-v2': D4MORL_Hopper_Cfg,
    'MO-HalfCheetah-v2': D4MORL_HalfCheetah_Cfg,
    'MO-Swimmer-v2': D4MORL_Swimmer_Cfg,
    'MO-Ant-v2': D4MORL_Ant_Cfg,
    'MO-Walker2d-v2': D4MORL_Walker2d_Cfg,

    'OfflineBallCircle-v0': DSRL_Cfg,  
    'OfflineAntCircle-v0': DSRL_Cfg,    
    'OfflineCarCircle-v0': DSRL_Cfg,
    'OfflineDroneCircle-v0': DSRL_Cfg,

    'OfflineBallRun-v0': DSRL_RUN_Cfg,    
    'OfflineCarRun-v0': DSRL_RUN_Cfg,
    'OfflineDroneRun-v0': DSRL_RUN_Cfg,
    'OfflineAntRun-v0': DSRL_RUN_Cfg,

    'OfflinePointGoal1Gymnasium-v0': DSRL_Cfg,
    'OfflineCarGoal1Gymnasium-v0': DSRL_Cfg,
    'OfflineAntVelocityGymnasium-v1': DSRL_Cfg,
    'OfflineHalfCheetahVelocityGymnasium-v1': DSRL_Cfg,
    'OfflineHopperVelocityGymnasium-v1': DSRL_Cfg,
    'OfflineSwimmerVelocityGymnasium-v1': DSRL_Cfg,
    'OfflineWalker2dVelocityGymnasium-v1': DSRL_Cfg,

    'CMO-Hopper-v2': CMO_Cfg,
    'CMO-Walker2d-v2': CMO_Cfg,
    'CMO-HalfCheetah-v2': CMO_Cfg,
    'CMO-Swimmer-v2': CMO_Cfg,
    'CMO-Ant-v2': CMO_Cfg,
}