"""Export tuned motor parameters to MJCF and robot.yaml files.

Reads the original motor.xml and robot.yaml, patches with optimised
parameter values, and writes motor_tuned.xml + robot_tuned.yaml.

Usage:
    python -m src.sysid.motor.export \
        --asset-path assets/motor \
        --result assets/motor/motor_sysid_result.json
"""

from __future__ import annotations

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

import structlog
import yaml

log = structlog.get_logger()

_DEFAULT_ASSET = "assets/motor"


def export_tuned_files(
    asset_path: str | Path,
    params: dict[str, float],
) -> tuple[Path, Path]:
    """Write tuned MJCF and robot.yaml files.

    Returns (tuned_mjcf_path, tuned_robot_yaml_path).
    """
    asset_path = Path(asset_path).resolve()

    robot_yaml_path = asset_path / "robot.yaml"
    robot_cfg = yaml.safe_load(robot_yaml_path.read_text())
    mjcf_path = asset_path / robot_cfg["mjcf"]

    # ── Tune MJCF ────────────────────────────────────────────────
    tree = ET.parse(str(mjcf_path))
    root = tree.getroot()

    # Actuator — use average gear for the MJCF model.
    gear_pos = params.get("actuator_gear_pos", params.get("actuator_gear"))
    gear_neg = params.get("actuator_gear_neg", params.get("actuator_gear"))
    gear_avg = None
    if gear_pos is not None and gear_neg is not None:
        gear_avg = (gear_pos + gear_neg) / 2.0
    elif gear_pos is not None:
        gear_avg = gear_pos
    filter_tau = params.get("actuator_filter_tau")
    for act_el in root.iter("general"):
        if act_el.get("name") == "motor":
            if gear_avg is not None:
                act_el.set("gear", str(gear_avg))
            if filter_tau is not None:
                if filter_tau > 0:
                    act_el.set("dyntype", "filter")
                    act_el.set("dynprm", str(filter_tau))
                else:
                    act_el.set("dyntype", "none")

    # Joint — average damping & friction for MJCF (asymmetry in runtime).
    fl_pos = params.get("motor_frictionloss_pos", params.get("motor_frictionloss"))
    fl_neg = params.get("motor_frictionloss_neg", params.get("motor_frictionloss"))
    fl_avg = None
    if fl_pos is not None and fl_neg is not None:
        fl_avg = (fl_pos + fl_neg) / 2.0
    elif fl_pos is not None:
        fl_avg = fl_pos
    damp_pos = params.get("motor_damping_pos", params.get("motor_damping"))
    damp_neg = params.get("motor_damping_neg", params.get("motor_damping"))
    damp_avg = None
    if damp_pos is not None and damp_neg is not None:
        damp_avg = (damp_pos + damp_neg) / 2.0
    elif damp_pos is not None:
        damp_avg = damp_pos
    for jnt in root.iter("joint"):
        if jnt.get("name") == "motor_joint":
            if damp_avg is not None:
                jnt.set("damping", str(damp_avg))
            if "motor_armature" in params:
                jnt.set("armature", str(params["motor_armature"]))
            if fl_avg is not None:
                jnt.set("frictionloss", str(fl_avg))

    # Rotor mass.
    if "rotor_mass" in params:
        for geom in root.iter("geom"):
            if geom.get("name") == "rotor_disk":
                geom.set("mass", str(params["rotor_mass"]))

    # Write tuned MJCF.
    tuned_mjcf_name = mjcf_path.stem + "_tuned" + mjcf_path.suffix
    tuned_mjcf_path = asset_path / tuned_mjcf_name
    ET.indent(tree, space="  ")
    tree.write(str(tuned_mjcf_path), xml_declaration=True, encoding="unicode")
    log.info("tuned_mjcf_written", path=str(tuned_mjcf_path))

    # ── Tune robot.yaml ──────────────────────────────────────────
    tuned_cfg = copy.deepcopy(robot_cfg)
    tuned_cfg["mjcf"] = tuned_mjcf_name

    if tuned_cfg.get("actuators") and len(tuned_cfg["actuators"]) > 0:
        act = tuned_cfg["actuators"][0]
        if gear_avg is not None:
            act["gear"] = round(gear_avg, 6)
        if "actuator_filter_tau" in params:
            act["filter_tau"] = round(params["actuator_filter_tau"], 6)
        if "motor_damping" in params:
            act["damping"] = round(params["motor_damping"], 6)

    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)
    if fl_avg is not None:
        mj["frictionloss"] = round(fl_avg, 6)

    # Asymmetric / hardware-realism / nonlinear parameters.
    realism = {}
    for key in [
        "actuator_gear_pos", "actuator_gear_neg",
        "motor_damping_pos", "motor_damping_neg",
        "motor_frictionloss_pos", "motor_frictionloss_neg",
        "motor_deadzone_pos", "motor_deadzone_neg",
        "action_bias",
        "viscous_quadratic", "back_emf_gain",
        "stribeck_friction_boost", "stribeck_vel",
        "gearbox_backlash",
    ]:
        if key in params:
            realism[key] = round(params[key], 6)
    if realism:
        tuned_cfg["hardware_realism"] = realism

    tuned_yaml_path = asset_path / "robot_tuned.yaml"
    header = (
        "# Tuned motor config — generated by src.sysid.motor.optimize\n"
        "# Original: robot.yaml\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_mjcf_path, tuned_yaml_path


# ── CLI ──────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Export tuned motor parameters to MJCF + robot.yaml."
    )
    parser.add_argument("--asset-path", type=str, default=_DEFAULT_ASSET)
    parser.add_argument(
        "--result", type=str, default=None,
        help="Path to motor_sysid_result.json (auto-detected if omitted)",
    )
    args = parser.parse_args()

    asset_path = Path(args.asset_path).resolve()
    if args.result:
        result_path = Path(args.result)
    else:
        result_path = asset_path / "motor_sysid_result.json"

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

    result = json.loads(result_path.read_text())
    params = result["best_params"]

    export_tuned_files(asset_path=args.asset_path, params=params)


if __name__ == "__main__":
    main()