# -*- coding: utf-8 -*-
import os, sys, numpy as np
import pinocchio as pin
from omegaconf import OmegaConf


sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from core.interaction_estimater import InteractionEstimator


# —— 误差函数（更鲁棒的相对误差 + 绝对误差）——
def rel_err(a, b, eps=1e-9):
    """对称式相对误差：‖a-b‖ / max(0.5(‖a‖+‖b‖), eps)"""
    na, nb = np.linalg.norm(a), np.linalg.norm(b)
    denom = max(0.5 * (na + nb), eps)
    return np.linalg.norm(a - b) / denom


def abs_err(a, b):
    """绝对误差：‖a-b‖"""
    return np.linalg.norm(a - b)


# —— 固定范数的扳手采样（方向随机）——
def sample_wrench_fixed_norm(F_lin_norm=10.0, F_ang_norm=2.0, rng=None):
    """
    在线性部分和角部分分别固定范数，方向随机的 6D 扳手采样
    """
    rng = np.random.default_rng() if rng is None else rng
    v = rng.standard_normal(6)
    f = v[:3]
    m = v[3:]
    f = f / (np.linalg.norm(f) + 1e-12) * F_lin_norm
    m = m / (np.linalg.norm(m) + 1e-12) * F_ang_norm
    return np.hstack([f, m])


def load_model(urdf_path: str) -> pin.Model:
    assert os.path.exists(urdf_path), f"URDF not found: {urdf_path}"
    model = pin.buildModelFromUrdf(urdf_path)
    print(f"[URDF] model loaded: nq={model.nq}, nv={model.nv}, njoints={model.njoints}")
    return model


def oracle_test(model,
                chest_frame: str = 'slave_shoulder',
                ee_frame: str = 'slave_ee',
                n_trials: int = 300,
                noise_tau_std: float = 0.0,
                lambda_damp: float = 1e-6,
                seed: int = 11,
                sv_min_skip: float = 1e-6):
    """
    使用 InteractionEstimator 的“理想/带噪声”仿真：
    - 通过动力学模型生成 tau_model
    - 叠加胸腔系扳手 CJ^T F_true 和测量噪声得到 tau_meas
    - 调用 est.estimate 恢复 tau_int, F_est
    - 统计相对/绝对误差的均值、方差、p90
    """
    rng = np.random.default_rng(seed)
    est = InteractionEstimator(model, chest_frame, ee_frame, lambda_damp=lambda_damp)
    data = model.createData()

    errs_rel, errs_abs = [], []
    kept = 0
    for k in range(n_trials):
        # 随机状态
        q = pin.randomConfiguration(model)
        qd = np.random.randn(model.nv) * 0.1
        qdd = np.random.randn(model.nv) * 0.2

        # 动力学力矩 tau_model
        M = pin.crba(model, data, q)
        M = (M + M.T) - np.diag(M.diagonal())
        nle = pin.nonLinearEffects(model, data, q, qd)
        g = pin.computeGeneralizedGravity(model, data, q)
        Cqd = nle - g
        tau_model = M @ qdd + Cqd + g

        # 胸腔系雅可比 CJ
        CJ = est._chest_jacobian(q, qd)  # 6 x nv
        svals = np.linalg.svd(CJ, compute_uv=False)
        if svals[-1] < sv_min_skip:
            # 跳过奇异邻域样本
            continue

        # 固定范数扳手
        CF_true = sample_wrench_fixed_norm(10.0, 2.0, rng=rng)

        # 合成“测得力矩” tau_meas = tau_model + CJ^T F_true + noise
        tau_meas = tau_model + CJ.T @ CF_true + np.random.randn(model.nv) * noise_tau_std

        # 利用 estimator 做估计
        tau_int, CF_est, CJ_check = est.estimate(q, qd, qdd, tau_meas)

        # 误差统计
        errs_rel.append(rel_err(CF_est, CF_true))
        errs_abs.append(abs_err(CF_est, CF_true))
        kept += 1

        # 一致性：回投误差 & 功率一致性（抽样打印）
        if k % 50 == 0:
            back_err = rel_err(CJ_check.T @ CF_est, tau_int)
            v_C = CJ @ qd
            Pq = float(tau_int @ qd)
            Pv = float(CF_true @ v_C)
            print(
                f"[{k:03d}] rel={errs_rel[-1]:.3e}, abs={errs_abs[-1]:.3e}, "
                f"back={back_err:.3e}, power={abs(Pq - Pv):.3e}, sv_min={svals[-1]:.2e}"
            )

    errs_rel = np.array(errs_rel)
    errs_abs = np.array(errs_abs)
    if kept == 0:
        print("[Oracle] all samples skipped by sv_min filter; try smaller sv_min_skip.")
        # 返回空 summary
        return dict(
            errs_rel=errs_rel,
            errs_abs=errs_abs,
            kept=0,
            mean_rel=np.nan,
            std_rel=np.nan,
            p90_rel=np.nan,
            mean_abs=np.nan,
            std_abs=np.nan,
            p90_abs=np.nan,
        )

    mean_rel = float(errs_rel.mean())
    std_rel = float(errs_rel.std())
    p90_rel = float(np.percentile(errs_rel, 90))

    mean_abs = float(errs_abs.mean())
    std_abs = float(errs_abs.std())
    p90_abs = float(np.percentile(errs_abs, 90))

    print(
        f"[Oracle REL] mean={mean_rel:.3e}, std={std_rel:.3e}, "
        f"median={np.median(errs_rel):.3e}, p90={p90_rel:.3e}"
    )
    print(
        f"[Oracle ABS] mean={mean_abs:.3e}, std={std_abs:.3e}, "
        f"median={np.median(errs_abs):.3e}, p90={p90_abs:.3e}"
    )

    summary = dict(
        errs_rel=errs_rel,
        errs_abs=errs_abs,
        kept=kept,
        mean_rel=mean_rel,
        std_rel=std_rel,
        p90_rel=p90_rel,
        mean_abs=mean_abs,
        std_abs=std_abs,
        p90_abs=p90_abs,
    )
    return summary


def oracle_test_baseline(model,
                         chest_frame: str = 'slave_shoulder',
                         ee_frame: str = 'slave_ee',
                         n_trials: int = 300,
                         noise_tau_std: float = 0.05,
                         seed: int = 17,
                         sv_min_skip: float = 1e-6):
    """
    基线：不调用 InteractionEstimator 的残差/阻尼公式，
    直接用未阻尼伪逆解胸腔扳手：

        tau_int_true = tau_meas - tau_model
        F_est = (CJ CJ^T)^{-1} CJ tau_int_true

    用来对比在同样噪声/配置下的数值稳定性。
    """
    rng = np.random.default_rng(seed)
    est_tmp = InteractionEstimator(model, chest_frame, ee_frame, lambda_damp=0.0)
    data = model.createData()

    errs_rel, errs_abs = [], []
    kept = 0
    for k in range(n_trials):
        # 随机状态
        q = pin.randomConfiguration(model)
        qd = np.random.randn(model.nv) * 0.1
        qdd = np.random.randn(model.nv) * 0.2

        # 动力学力矩 tau_model
        M = pin.crba(model, data, q)
        M = (M + M.T) - np.diag(M.diagonal())
        nle = pin.nonLinearEffects(model, data, q, qd)
        g = pin.computeGeneralizedGravity(model, data, q)
        Cqd = nle - g
        tau_model = M @ qdd + Cqd + g

        # 胸腔雅可比 CJ
        CJ = est_tmp._chest_jacobian(q, qd)
        svals = np.linalg.svd(CJ, compute_uv=False)
        if svals[-1] < sv_min_skip:
            continue

        # 固定范数扳手（与 oracle_test 相同范数）
        CF_true = sample_wrench_fixed_norm(10.0, 2.0, rng=rng)

        # 测得力矩（含噪声）
        tau_meas = tau_model + CJ.T @ CF_true + np.random.randn(model.nv) * noise_tau_std

        # “真实”关节残差力矩（在仿真中我们知道 tau_model）
        tau_int_true = tau_meas - tau_model

        # 未阻尼伪逆 (CJ CJ^T)^{-1} CJ tau_int_true
        JJt = CJ @ CJ.T
        try:
            F_est = np.linalg.solve(JJt, CJ @ tau_int_true)
        except np.linalg.LinAlgError:
            # 数值奇异，跳过该样本
            continue

        errs_rel.append(rel_err(F_est, CF_true))
        errs_abs.append(abs_err(F_est, CF_true))
        kept += 1

        if k % 50 == 0:
            back_err = rel_err(CJ.T @ F_est, tau_int_true)
            print(
                f"[BASE {k:03d}] rel={errs_rel[-1]:.3e}, abs={errs_abs[-1]:.3e}, "
                f"back={back_err:.3e}, sv_min={svals[-1]:.2e}"
            )

    errs_rel = np.array(errs_rel)
    errs_abs = np.array(errs_abs)
    if kept == 0:
        print("[Baseline] all samples skipped by sv_min filter; try smaller sv_min_skip.")
        return dict(
            errs_rel=errs_rel,
            errs_abs=errs_abs,
            kept=0,
            mean_rel=np.nan,
            std_rel=np.nan,
            p90_rel=np.nan,
            mean_abs=np.nan,
            std_abs=np.nan,
            p90_abs=np.nan,
        )

    mean_rel = float(errs_rel.mean())
    std_rel = float(errs_rel.std())
    p90_rel = float(np.percentile(errs_rel, 90))

    mean_abs = float(errs_abs.mean())
    std_abs = float(errs_abs.std())
    p90_abs = float(np.percentile(errs_abs, 90))

    print(
        f"[BASELINE REL] mean={mean_rel:.3e}, std={std_rel:.3e}, "
        f"median={np.median(errs_rel):.3e}, p90={p90_rel:.3e}"
    )
    print(
        f"[BASELINE ABS] mean={mean_abs:.3e}, std={std_abs:.3e}, "
        f"median={np.median(errs_abs):.3e}, p90={p90_abs:.3e}"
    )

    summary = dict(
        errs_rel=errs_rel,
        errs_abs=errs_abs,
        kept=kept,
        mean_rel=mean_rel,
        std_rel=std_rel,
        p90_rel=p90_rel,
        mean_abs=mean_abs,
        std_abs=std_abs,
        p90_abs=p90_abs,
    )
    return summary


def consistency_sweep(model,
                      chest_frame: str = 'slave_shoulder',
                      ee_frame: str = 'slave_ee'):
    """
    对不同阻尼系数 lambda 进行一个小 sweep，查看：
    - 估计扳手范数
    - 回投误差
    - cond(JJ^T)
    """
    est_ref = InteractionEstimator(model, chest_frame, ee_frame, lambda_damp=1e-3)
    q = pin.randomConfiguration(model)
    qd = np.random.randn(model.nv) * 0.05
    qdd = np.zeros(model.nv)

    CJ = est_ref._chest_jacobian(q, qd)
    CF_true = np.array([5.0, -3.0, 8.0, 0.5, 0.2, -0.1])

    M = pin.crba(model, est_ref.data, q)
    M = (M + M.T) - np.diag(M.diagonal())
    nle = pin.nonLinearEffects(model, est_ref.data, q, qd)
    g = pin.computeGeneralizedGravity(model, est_ref.data, q)
    Cqd = nle - g
    tau_model = M @ qdd + Cqd + g
    tau_meas = tau_model + CJ.T @ CF_true

    lambdas = [1e-6, 3e-6, 1e-5, 3e-5, 1e-4, 3e-4, 1e-3, 3e-3, 1e-2]
    print("\n[Lambda sweep]  λ,  ‖F_est‖,  back_err,  cond(JJᵀ)")
    for lam in lambdas:
        est = InteractionEstimator(model, chest_frame, ee_frame, lambda_damp=lam)
        tau_int, CF_est, CJ_use = est.estimate(q, qd, qdd, tau_meas)
        back_err = rel_err(CJ_use.T @ CF_est, tau_int)
        JJt = CJ_use @ CJ_use.T
        cond = np.linalg.cond(JJt) if np.linalg.matrix_rank(JJt) == 6 else np.inf
        print(f"{lam:8.1e}  {np.linalg.norm(CF_est):8.3f}  {back_err:8.2e}  {cond:8.2e}")


def main():
    conf = OmegaConf.load("./config/config.yaml")
    urdf_path = getattr(conf, "slave_urdf", None)
    model = load_model(urdf_path)

    # 注意：chest_frame 在论文中记为 C 框架，这里保持一致
    chest_frame = "slave_base"   # 或 "slave_shoulder"，视你的 URDF 定义而定
    ee_frame = "slave_ee"

    print("\n=== ORACLE (no noise, λ=1e-6) ===")
    summary_ideal = oracle_test(
        model, chest_frame, ee_frame,
        n_trials=300, noise_tau_std=0.0, lambda_damp=1e-6,
        seed=11, sv_min_skip=1e-6
    )

    print("\n=== PROPOSED (tau noise 0.05 N·m, λ=1e-3) ===")
    summary_proposed = oracle_test(
        model, chest_frame, ee_frame,
        n_trials=300, noise_tau_std=0.05, lambda_damp=1e-3,
        seed=13, sv_min_skip=1e-6
    )

    print("\n=== BASELINE (tau noise 0.05 N·m, undamped pseudoinverse) ===")
    summary_baseline = oracle_test_baseline(
        model, chest_frame, ee_frame,
        n_trials=300, noise_tau_std=0.05,
        seed=17, sv_min_skip=1e-6
    )

    # 打印一个 markdown 风格的小表格，方便直接贴到论文
    def to_percent(x):
        return 100.0 * x if x is not None and not np.isnan(x) else float("nan")

    print("\n=== SUMMARY (relative wrench error e_F) ===")
    print("| Method                 | mean e_F [%] | std e_F [%] | p90 e_F [%] | Kept |")
    print("|------------------------|-------------:|------------:|------------:|-----:|")
    print(
        f"| Ideal (no noise)      | {to_percent(summary_ideal['mean_rel']):11.2f} | "
        f"{to_percent(summary_ideal['std_rel']):11.2f} | "
        f"{to_percent(summary_ideal['p90_rel']):11.2f} | "
        f"{summary_ideal['kept']:4d} |"
    )
    print(
        f"| Proposed (damped)     | {to_percent(summary_proposed['mean_rel']):11.2f} | "
        f"{to_percent(summary_proposed['std_rel']):11.2f} | "
        f"{to_percent(summary_proposed['p90_rel']):11.2f} | "
        f"{summary_proposed['kept']:4d} |"
    )
    print(
        f"| Baseline (PI, undamp) | {to_percent(summary_baseline['mean_rel']):11.2f} | "
        f"{to_percent(summary_baseline['std_rel']):11.2f} | "
        f"{to_percent(summary_baseline['p90_rel']):11.2f} | "
        f"{summary_baseline['kept']:4d} |"
    )

    # 可选：扫 lambda，看条件数/回投误差（和论文里 Fig. 6(a)(b) 对应）
    consistency_sweep(model, chest_frame, ee_frame)


if __name__ == "__main__":
    main()