import numpy as np
import torch

class DomainRandomizer:
    """Isaac Gym域随机化配置类"""

    def __init__(self, robot_name: str = "GR-1"):
        self.robot_name = robot_name
        self.params = self._get_default_params()

    def _get_default_params(self) -> dict:
        """默认随机化参数范围"""
        return {
            # 物理参数随机化
            "ground_friction": {"default": 0.7, "range": [0.3, 1.2]},
            "motor_torque_constant": {"default": 0.05, "range": [0.04, 0.06]},
            "motor_friction": {"default": 0.01, "range": [0.005, 0.02]},
            "motor_dead_zone": {"default": 0.0, "range": [0.0, 0.02]},
            # 观测噪声 (3σ覆盖)
            "joint_pos_noise": {"default": 0.0, "range": [0.0, 0.05]},
            "joint_vel_noise": {"default": 0.0, "range": [0.0, 0.1]},
            "base_vel_noise": {"default": 0.0, "range": [0.0, 0.15]},
            "imu_noise": {"default": 0.0, "range": [0.0, 0.05]},
            # 控制延迟 (步数)
            "action_delay": {"default": 0, "range": [0, 3]},
            "observation_delay": {"default": 0, "range": [0, 2]},
            # 附加负载
            "payload_mass": {"default": 0.0, "range": [0.0, 2.0]},
            # 重力扰动
            "gravity_perturbation": {"default": 0.0, "range": [-0.1, 0.1]},
            # 关节限位缩放
            "joint_limit_scale": {"default": 1.0, "range": [0.95, 1.05]},
        }

    def sample_randomized_params(self) -> dict:
        """采样一套随机化参数,均匀分布"""
        sampled = {}
        for name, cfg in self.params.items():
            low, high = cfg["range"]
            sampled[name] = np.random.uniform(low, high)
        return sampled

    def apply_to_sim(self, sim_env, randomized_params: dict):
        """将随机化参数应用到仿真环境"""
        if "ground_friction" in randomized_params:
            sim_env.set_ground_friction(randomized_params["ground_friction"])
        if "payload_mass" in randomized_params:
            sim_env.add_payload_mass(randomized_params["payload_mass"])
        if "motor_torque_constant" in randomized_params:
            sim_env.set_motor_torque_constant(
                randomized_params["motor_torque_constant"]
            )
        if "gravity_perturbation" in randomized_params:
            g_base = 9.81
            g_perturbed = g_base + randomized_params["gravity_perturbation"]
            sim_env.set_gravity([0, 0, -g_perturbed])
        return sim_env

# 使用示例: 每1万步重采样一次参数
randomizer = DomainRandomizer(robot_name="GR-1")
envs_per_gpu = 4096  # Isaac Gym典型并行数
reset_buf = torch.zeros(envs_per_gpu, dtype=torch.long, device="cuda:0")
reset_count = 0

for step in range(100000):
    if reset_count == 0 or step % 10000 == 0:
        randomized_params = randomizer.sample_randomized_params()
        randomizer.apply_to_sim(gym_env, randomized_params)
        reset_count = 50

    actions = policy.compute_actions(obs_dict)
    obs_dict, rewards, dones = gym_env.step(actions)
    reset_buf = torch.where(dones, torch.ones_like(reset_buf), reset_buf)
    reset_count -= 1

import numpy as np
from scipy.optimize import minimize

def motor_model(torque_cmd: np.ndarray,
                joint_vel: np.ndarray,
                k_t: float = 0.05,
                f_c: float = 0.01,
                f_v: float = 0.001) -> np.ndarray:
    """电机力矩模型: 含库仑+粘性摩擦"""
    friction = f_c * np.sign(joint_vel) + f_v * np.abs(joint_vel)
    actual_torque = k_t * torque_cmd - friction
    return actual_torque

def fit_motor_parameters(
    measured_current: np.ndarray,
    measured_torque: np.ndarray,
    measured_vel: np.ndarray
) -> dict:
    """通过实际测量数据拟合电机参数(最小二乘)"""
    def objective(x):
        k_t, f_c, f_v = x
        pred = motor_model(measured_current, measured_vel, k_t, f_c, f_v)
        return np.sum((measured_torque - pred) ** 2)

    # 初值: 标称参数
    x0 = [0.05, 0.01, 0.001]
    # 边界约束
    bounds = [
        (0.03, 0.08),   # k_t: 合理范围
        (0.001, 0.05),  # f_c
        (0.0001, 0.01)  # f_v
    ]

    result = minimize(objective, x0, method='L-BFGS-B', bounds=bounds)
    k_t, f_c, f_v = result.x

    print(f"拟合结果: k_t={k_t:.5f}, f_c={f_c:.5f}, f_v={f_v:.5f}")
    print(f"残差RMS: {np.sqrt(result.fun/len(measured_torque)):.5f} N·m")

    return {"k_t": k_t, "f_c": f_c, "f_v": f_v}

# 示例: 使用虚拟数据验证
np.random.seed(42)
I = np.random.uniform(0, 1, 1000)        # 0~1A电流
v = np.random.uniform(-3, 3, 1000)       # -3~3 rad/s速度
# 虚拟真实力矩 (带噪声的真机数据)
true_tau = motor_model(I, v, k_t=0.048, f_c=0.012, f_v=0.0015) \
           + np.random.normal(0, 0.005, 1000)

fitted = fit_motor_parameters(I, true_tau, v)
# 输出: k_t≈0.048, f_c≈0.012, f_v≈0.0015 (拟合成功)

import torch
import torch.nn as nn

class RobustActorCritic(nn.Module):
    """
    鲁棒Actor-Critic网络
    改进点:
    1. LayerNorm观察归一化
    2. 动作平滑 (Exponential Moving Average)
    3. 安全动作裁剪
    4. 隐层域适应 (Domain-Adaptive Layer)
    """

    def __init__(
        self,
        num_obs: int,
        num_actions: int,
        hidden_sizes: list = [256, 256, 128],
        action_scale: float = 1.0
    ):
        super().__init__()
        self.action_scale = action_scale
        self.num_obs = num_obs
        self.num_actions = num_actions

        # 观察归一化层
        self.obs_norm = nn.LayerNorm(num_obs)

        # Actor网络 (策略)
        actor_layers = []
        in_dim = num_obs
        for h_dim in hidden_sizes:
            actor_layers.extend([
                nn.Linear(in_dim, h_dim),
                nn.LayerNorm(h_dim),
                nn.Tanh()
            ])
            in_dim = h_dim
        actor_layers.append(nn.Linear(in_dim, num_actions))
        actor_layers.append(nn.Tanh())  # 输出归一化到[-1, 1]
        self.actor = nn.Sequential(*actor_layers)

        # Critic网络 (价值)
        critic_layers = []
        in_dim = num_obs
        for h_dim in hidden_sizes:
            critic_layers.extend([
                nn.Linear(in_dim, h_dim),
                nn.LayerNorm(h_dim),
                nn.Tanh()
            ])
            in_dim = h_dim
        critic_layers.append(nn.Linear(in_dim, 1))
        self.critic = nn.Sequential(*critic_layers)

        # 动作EMA平滑器
        self.action_ema = None
        self.ema_alpha = 0.7

        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.orthogonal_(m.weight, gain=np.sqrt(2))
                nn.init.constant_(m.bias, 0.0)

    def forward(self, obs: torch.Tensor) -> tuple:
        """前向传播,返回(动作, 价值)"""
        obs_norm = self.obs_norm(obs)
        actions = self.actor(obs_norm)
        values = self.critic(obs_norm).squeeze(-1)
        return actions, values

    def compute_actions(self,
                        obs: torch.Tensor,
                        deterministic: bool = False) -> torch.Tensor:
        """计算动作 (含EMA平滑和安全裁剪)"""
        actions, _ = self.forward(obs)

        # EMA动作平滑
        if self.action_ema is None:
            self.action_ema = actions.detach().clone()
        else:
            self.action_ema = self.ema_alpha * self.action_ema + \
                             (1 - self.ema_alpha) * actions.detach()

        # 推理时用EMA平滑后的动作
        if deterministic:
            actions = self.action_ema

        # 安全裁剪: 限制关节角速度变化率
        clipped_actions = torch.clamp(
            actions,
            min=-self.action_scale,
            max=self.action_scale
        )
        return clipped_actions

# 使用示例 (GR-1人形机器人)
policy = RobustActorCritic(
    num_obs=48,      # base_lin_vel(3)+base_ang_vel(3)+projected_gravity(3)
                     # +commands(3)+joint_pos(19)+joint_vel(17)
    num_actions=19,  # 下肢关节数
    hidden_sizes=[256, 256, 128]
)
print(f"策略网络参数量: {sum(p.numel() for p in policy.parameters()) / 1e6:.2f}M")

import torch

class UnifiedRewardFunction:
    """
    统一奖励函数
    兼容站立+行走两种运动模式
    奖励 = w1*站立奖励 + w2*行走奖励 + w3*安全惩罚 + w4*平滑惩罚
    """

    def __init__(
        self,
        w_standing: float = 0.4,
        w_walking: float = 0.4,
        w_safety: float = 0.1,
        w_smooth: float = 0.1
    ):
        self.w_standing = w_standing
        self.w_walking = w_walking
        self.w_safety = w_safety
        self.w_smooth = w_smooth

    def compute(self,
               obs_dict: dict,
               actions: torch.Tensor,
               prev_actions: torch.Tensor = None) -> torch.Tensor:
        """计算奖励(obs_dict含base_height/gravity/vel/joint_pos等)"""
        r_total = torch.zeros_like(obs_dict["base_height"])

        # 站立基础奖励
        r_standing = self._reward_standing(
            base_height=obs_dict["base_height"],
            gravity_proj=obs_dict["projected_gravity"],
            joint_vel=obs_dict["joint_vel"]
        )
        r_total += self.w_standing * r_standing

        # 行走奖励
        r_walking = self._reward_walking(
            base_vel=obs_dict["base_lin_vel"],
            commands=obs_dict["commands"]
        )
        r_total += self.w_walking * r_walking

        # 安全惩罚
        r_safety = self._reward_safety(
            joint_pos=obs_dict["joint_pos"],
            base_height=obs_dict["base_height"],
            gravity_proj=obs_dict["projected_gravity"]
        )
        r_total += self.w_safety * r_safety

        # 动作平滑惩罚
        if prev_actions is not None:
            r_smooth = -torch.sum((actions - prev_actions) ** 2, dim=-1) * 0.5
            r_total += self.w_smooth * r_smooth

        return r_total

    def _reward_standing(self, base_height, gravity_proj, joint_vel):
        r_height = torch.exp(-10 * torch.abs(base_height - 0.85) ** 2)
        r_upright = torch.exp(-5 * torch.abs(gravity_proj[:, 2] - 1.0) ** 2)
        r_calm = torch.exp(-0.1 * torch.sum(joint_vel ** 2, dim=-1))
        return (r_height + r_upright + r_calm) / 3.0

    def _reward_walking(self, base_vel, commands):
        vx_cmd, vy_cmd, vyaw_cmd = commands[:, 0], commands[:, 1], commands[:, 2]
        r_vx = torch.exp(-5 * torch.abs(base_vel[:, 0] - vx_cmd) ** 2)
        r_vy = torch.exp(-10 * torch.abs(base_vel[:, 1] - vy_cmd) ** 2)
        r_yaw = torch.exp(-10 * torch.abs(base_vel[:, 2] / 0.3 - vyaw_cmd) ** 2)
        return (r_vx + r_vy + r_yaw) / 3.0

    def _reward_safety(self, joint_pos, base_height, gravity_proj):
        height_pen = torch.where(base_height < 0.5, -2.0, 0.0)
        fall_pen = torch.where(gravity_proj[:, 2] < 0.3, -5.0, 0.0)
        return height_pen + fall_pen