"""Training configs for Go1 go2style environments.

Importing triggers registration of: flat, stairs, combined.
"""
from dataclasses import dataclass

import motrix_envs.locomotion.go1.go2style_flat     # noqa: F401
import motrix_envs.locomotion.go1.go2style_stairs   # noqa: F401
import motrix_envs.locomotion.go1.go2style_combined # noqa: F401
import motrix_envs.locomotion.go1.go2style_cts      # noqa: F401
import motrix_envs.locomotion.go1.orig_combined     # noqa: F401

from motrix_rl.registry import rlcfg
from motrix_rl.skrl.config import SkrlCfg


def _apply_go2style_ppo(cfg: SkrlCfg):
    """Apply go2_rl_gym-style PPO params to config."""
    cfg.runner.models.policy.hiddens = [512, 256, 128]
    cfg.runner.models.value.hiddens = [512, 256, 128]
    cfg.runner.agent.rollouts = 24
    cfg.runner.agent.learning_epochs = 5
    cfg.runner.agent.mini_batches = 4
    cfg.runner.agent.learning_rate = 1e-3
    cfg.runner.trainer.timesteps = 60000


def _apply_orig_ppo(cfg: SkrlCfg):
    """Match original Go1 stairs config: [512,256,128], 30k steps."""
    cfg.runner.models.policy.hiddens = [512, 256, 128]
    cfg.runner.models.value.hiddens = [512, 256, 128]
    cfg.runner.agent.rollouts = 24
    cfg.runner.agent.learning_epochs = 5
    cfg.runner.agent.mini_batches = 4
    cfg.runner.agent.learning_rate = 1e-3
    cfg.runner.trainer.timesteps = 30000  # same as original stairs


class skrl:
    @rlcfg("go1-flat-terrain-walk-go2style")
    @dataclass
    class Flat(SkrlCfg):
        def __post_init__(self): _apply_go2style_ppo(self)

    @rlcfg("go1-stairs-terrain-walk-go2style")
    @dataclass
    class Stairs(SkrlCfg):
        def __post_init__(self): _apply_go2style_ppo(self)

    @rlcfg("go1-combined-terrain-walk-go2style")
    @dataclass
    class Combined(SkrlCfg):
        def __post_init__(self): _apply_go2style_ppo(self)

    @rlcfg("go1-cts-flat-walk-go2style")
    @dataclass
    class CTS(SkrlCfg):
        def __post_init__(self): _apply_go2style_ppo(self)

    @rlcfg("go1-orig-combined-walk")
    @dataclass
    class OrigCombined(SkrlCfg):
        """Original Go1 (PD 80, 45-dim, [256,128,64]) on flat+stairs."""
        def __post_init__(self): _apply_orig_ppo(self)
