from __future__ import annotations

import argparse
from dataclasses import dataclass

import mujoco

from .controller import PDGaitConfig, PDGaitController
from .sim import load_model, reset_to_stand, snapshot


@dataclass
class EvalResult:
    seconds: float
    distance_m: float
    avg_speed_mps: float
    min_torso_z_m: float
    max_torso_z_m: float
    fallen: bool


def torso_hit_floor(model: mujoco.MjModel, data: mujoco.MjData) -> bool:
    torso_gid = model.geom("torso_geom").id
    floor_gid = model.geom("floor").id
    for i in range(data.ncon):
        c = data.contact[i]
        if (c.geom1 == torso_gid and c.geom2 == floor_gid) or (c.geom1 == floor_gid and c.geom2 == torso_gid):
            return True
    return False


def run_eval(seconds: float, gait_hz: float, kp: float, kd: float, model_path: str | None = None) -> EvalResult:
    model, data = load_model(model_path)
    reset_to_stand(model, data)
    controller = PDGaitController(model, data, PDGaitConfig(frequency_hz=gait_hz, kp=kp, kd=kd))

    start = snapshot(data)
    min_torso = start.torso_z
    max_torso = start.torso_z

    steps = int(seconds / model.opt.timestep)
    fallen = False

    for _ in range(steps):
        controller.step(data.time)
        mujoco.mj_step(model, data)
        s = snapshot(data)
        min_torso = min(min_torso, s.torso_z)
        max_torso = max(max_torso, s.torso_z)
        if torso_hit_floor(model, data) or s.torso_z < 0.08:
            fallen = True

    end = snapshot(data)
    elapsed = max(end.t - start.t, 1e-9)
    distance = end.torso_x - start.torso_x

    return EvalResult(
        seconds=elapsed,
        distance_m=distance,
        avg_speed_mps=distance / elapsed,
        min_torso_z_m=min_torso,
        max_torso_z_m=max_torso,
        fallen=fallen,
    )


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Evaluate baseline MuJoCo walker")
    p.add_argument("--model", type=str, default=None, help="Path to MJCF model")
    p.add_argument("--seconds", type=float, default=10.0, help="Evaluation duration")
    p.add_argument("--gait-hz", type=float, default=0.9, help="Gait frequency")
    p.add_argument("--kp", type=float, default=24.0, help="PD proportional gain")
    p.add_argument("--kd", type=float, default=2.5, help="PD derivative gain")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    result = run_eval(
        seconds=args.seconds,
        gait_hz=args.gait_hz,
        kp=args.kp,
        kd=args.kd,
        model_path=args.model,
    )

    print("Walker evaluation")
    print(f"- duration_s:      {result.seconds:.3f}")
    print(f"- distance_m:      {result.distance_m:.3f}")
    print(f"- avg_speed_mps:   {result.avg_speed_mps:.3f}")
    print(f"- min_torso_z_m:   {result.min_torso_z_m:.3f}")
    print(f"- max_torso_z_m:   {result.max_torso_z_m:.3f}")
    print(f"- fallen:          {result.fallen}")


if __name__ == "__main__":
    main()