"""Export tuned parameters to URDF and robot.yaml files.

Reads the original files, injects the optimised parameter values,
and writes ``rotary_cartpole_tuned.urdf`` + ``robot_tuned.yaml``
alongside the originals in the robot asset directory.
"""

from __future__ import annotations

import copy
import xml.etree.ElementTree as ET
from pathlib import Path

import structlog
import yaml

from src.sysid._urdf import patch_link_inertials

log = structlog.get_logger()


def export_tuned_files(
    robot_path: str | Path,
    params: dict[str, float],
    motor_params: dict[str, float] | None = None,
) -> tuple[Path, Path]:
    """Write tuned URDF and robot.yaml files.

    Parameters
    ----------
    robot_path : robot asset directory (contains robot.yaml + *.urdf)
    params : dict of parameter name → tuned value (the optimised set)
    motor_params : locked motor parameters merged underneath ``params``
        (``params`` wins on conflicts) so the exported YAML always has a
        complete motor model

    Returns
    -------
    (tuned_urdf_path, tuned_robot_yaml_path)
    """
    robot_path = Path(robot_path).resolve()
    if motor_params:
        params = {**motor_params, **params}

    # ── Load originals ───────────────────────────────────────────
    robot_yaml_path = robot_path / "robot.yaml"
    robot_cfg = yaml.safe_load(robot_yaml_path.read_text())
    urdf_path = robot_path / robot_cfg["urdf"]

    # ── Tune URDF ────────────────────────────────────────────────
    tree = ET.parse(urdf_path)
    patch_link_inertials(tree.getroot(), params)

    # Write tuned URDF.
    tuned_urdf_name = urdf_path.stem + "_tuned" + urdf_path.suffix
    tuned_urdf_path = robot_path / tuned_urdf_name

    # Preserve the XML declaration and original formatting as much as possible.
    ET.indent(tree, space="  ")
    tree.write(str(tuned_urdf_path), xml_declaration=True, encoding="unicode")
    log.info("tuned_urdf_written", path=str(tuned_urdf_path))

    # ── Tune robot.yaml ──────────────────────────────────────────
    tuned_cfg = copy.deepcopy(robot_cfg)

    # Point to the tuned URDF.
    tuned_cfg["urdf"] = tuned_urdf_name

    # Update actuator parameters — full asymmetric motor model.
    if tuned_cfg.get("actuators") and len(tuned_cfg["actuators"]) > 0:
        act = tuned_cfg["actuators"][0]

        # Asymmetric gear, damping, deadzone, frictionloss as [pos, neg].
        act["gear"] = [
            round(params.get("actuator_gear_pos", 0.424), 6),
            round(params.get("actuator_gear_neg", 0.425), 6),
        ]
        act["damping"] = [
            round(params.get("motor_damping_pos", 0.002), 6),
            round(params.get("motor_damping_neg", 0.015), 6),
        ]
        act["deadzone"] = [
            round(params.get("motor_deadzone_pos", 0.141), 6),
            round(params.get("motor_deadzone_neg", 0.078), 6),
        ]
        act["frictionloss"] = [
            round(params.get("motor_frictionloss_pos", 0.057), 6),
            round(params.get("motor_frictionloss_neg", 0.053), 6),
        ]
        if "actuator_filter_tau" in params:
            act["filter_tau"] = round(params["actuator_filter_tau"], 6)

        # Stribeck friction and action bias.
        if "stribeck_friction_boost" in params:
            act["stribeck_friction_boost"] = round(params["stribeck_friction_boost"], 6)
        if "stribeck_vel" in params:
            act["stribeck_vel"] = round(params["stribeck_vel"], 6)
        if "action_bias" in params:
            act["action_bias"] = round(params["action_bias"], 6)

        # ctrl_range from ctrl_limit parameter.
        if "ctrl_limit" in params:
            lim = round(params["ctrl_limit"], 6)
            act["ctrl_range"] = [-lim, lim]

    # Update joint overrides.
    if "joints" not in tuned_cfg:
        tuned_cfg["joints"] = {}

    if "motor_joint" not in tuned_cfg["joints"]:
        tuned_cfg["joints"]["motor_joint"] = {}
    mj = tuned_cfg["joints"]["motor_joint"]
    if "motor_armature" in params:
        mj["armature"] = round(params["motor_armature"], 6)
    # Frictionloss/damping = 0 in MuJoCo (motor model handles via qfrc_applied).
    mj["frictionloss"] = 0.0

    if "pendulum_joint" not in tuned_cfg["joints"]:
        tuned_cfg["joints"]["pendulum_joint"] = {}
    pj = tuned_cfg["joints"]["pendulum_joint"]
    if "pendulum_damping" in params:
        pj["damping"] = round(params["pendulum_damping"], 6)
    if "pendulum_frictionloss" in params:
        pj["frictionloss"] = round(params["pendulum_frictionloss"], 6)

    # Write tuned robot.yaml.
    tuned_yaml_path = robot_path / "robot_tuned.yaml"

    # Add a header comment.
    header = (
        "# Tuned robot config — generated by src.sysid.optimize\n"
        "# Original: robot.yaml\n"
        "# Run `python -m src.sysid.visualize` to compare real vs sim.\n\n"
    )
    tuned_yaml_path.write_text(
        header + yaml.dump(tuned_cfg, default_flow_style=False, sort_keys=False)
    )
    log.info("tuned_robot_yaml_written", path=str(tuned_yaml_path))

    return tuned_urdf_path, tuned_yaml_path


# ── CLI entry point ──────────────────────────────────────────────────


def main() -> None:
    import argparse
    import json

    parser = argparse.ArgumentParser(
        description="Export tuned URDF + robot.yaml from sysid results."
    )
    parser.add_argument(
        "--robot-path", type=str, default="assets/rotary_cartpole",
        help="Path to robot asset directory",
    )
    parser.add_argument(
        "--result", type=str, default=None,
        help="Path to sysid_result.json (auto-detected if omitted)",
    )
    args = parser.parse_args()

    robot_path = Path(args.robot_path).resolve()
    if args.result:
        result_path = Path(args.result)
    else:
        result_path = robot_path / "sysid_result.json"

    if not result_path.exists():
        raise FileNotFoundError(f"Result file not found: {result_path}")

    result = json.loads(result_path.read_text())

    export_tuned_files(
        robot_path=args.robot_path,
        params=result["best_params"],
    )


if __name__ == "__main__":
    main()

