# tire_force_reader.py
#
# Listener to read tire/ground contact forces in AGX Dynamics.
# Usage example (in your main script):
#
#   from tire_force_reader import TireForceReader
#   sim = application().getSimulation()
#   reader = TireForceReader(
#       sim,
#       tire_body_paths=[
#           "Scene.dump_truck.wheel_L1",
#           "Scene.dump_truck.wheel_R1",
#           "Scene.dump_truck.wheel_L2",
#           "Scene.dump_truck.wheel_R2",
#       ],
#       terrain_path="Scene.terrain",
#       use_geometry_if_no_body=True,
#       print_every_n_steps=1
#   )
#   sim.add(reader)

from __future__ import annotations

from typing import Dict, List, Optional, Sequence, Union

import agx
import agxSDK


TerrainTarget = Union["agx.RigidBody", "agxCollide.Geometry"]


class TireForceReader(agxSDK.StepEventListener):
    """
    Reads summed contact forces between one or more wheel rigid bodies and the terrain.

    - Uses agx.ContactForceReader
    - Reads forces in post-step
    - Works with either:
        (A) terrain as a RigidBody, or
        (B) terrain as a Geometry (fallback if no rigid body exists)

    Forces are returned in WORLD coordinates and include normal + friction components.
    """

    def __init__(
        self,
        sim: "agxSDK.Simulation",
        tire_body_paths: Sequence[str],
        terrain_path: str,
        *,
        use_geometry_if_no_body: bool = True,
        print_every_n_steps: int = 1,
    ):
        super().__init__()
        self._sim = sim
        self._reader = agx.ContactForceReader(sim)

        if not tire_body_paths:
            raise ValueError("tire_body_paths must contain at least one rigid body path.")

        self._wheel_paths: List[str] = list(tire_body_paths)        
        self._wheels: Dict[str, agx.RigidBody] = {}
        self._terrain_path = terrain_path
        self._terrain: Optional[TerrainTarget] = None

        self._use_geometry_if_no_body = use_geometry_if_no_body
        self._print_every_n_steps = max(1, int(print_every_n_steps))
        self._step_counter = 0

        self._resolve_handles()

    # -------- public API --------

    def get_wheel_force_world(self, wheel_path: str) -> agx.Vec3:
        """
        Return summed wheel↔terrain force vector in WORLD coordinates for the given wheel path.
        """
        wheel = self._wheels.get(wheel_path)
        if wheel is None:
            raise KeyError(f"Unknown wheel_path '{wheel_path}'. Known: {list(self._wheels.keys())}")

        if self._terrain is None:
            raise RuntimeError("Terrain target is not resolved.")

        return self._reader.getContactForces(wheel, self._terrain)

    def get_all_forces_world(self) -> Dict[str, agx.Vec3]:
        """
        Return dict: wheel_path -> summed force vector (WORLD coordinates).
        """
        return {wp: self.get_wheel_force_world(wp) for wp in self._wheel_paths}

    # -------- agxSDK.StepEventListener --------

    def post(self, t: float) -> None:
        """
        Called after a simulation step has been solved. This is where contact forces are valid.
        """
        self._step_counter += 1
        if (self._step_counter % self._print_every_n_steps) != 0:
            return

        # Re-resolve if something was missing at init (e.g., scene not fully loaded yet)
        if self._terrain is None or not self._wheels:
            self._resolve_handles()
            if self._terrain is None:
                print("No terrain")
                return
            if not self._wheels:
                print("No wheel")
                return

        forces = self.get_all_forces_world()

        # Extract vertical forces (Z)
        fz = {wp: F.z() for wp, F in forces.items()}

        # Wheel order assumption:
        # L1, R1, L2, R2, L3, R3
        L1 = fz.get(self._wheel_paths[0], 0.0) / 9.80665
        R1 = fz.get(self._wheel_paths[1], 0.0) / 9.80665
        L2 = fz.get(self._wheel_paths[2], 0.0) / 9.80665
        R2 = fz.get(self._wheel_paths[3], 0.0) / 9.80665
        L3 = fz.get(self._wheel_paths[4], 0.0) / 9.80665
        R3 = fz.get(self._wheel_paths[5], 0.0) / 9.80665

        axle1 = L1 + R1
        axle2 = L2 + R2
        axle3 = L3 + R3
        total = axle1 + axle2 + axle3
        
        axle1_pct = axle1/total*100
        axle2_pct = axle2/total*100
        axle3_pct = axle3/total*100
        
        
        print(f"\nTime {t:8.3f} s")
        print("-----------------------------------------------------------")
        print(" Axle |  Left [kg] |  Right [kg] |  Sum L+R [kg]")
        print("-----------------------------------------------------------")
        print(f"   1  | {L1:10.0f} |  {R1:10.0f} |   {axle1:6.0f}  ({axle1_pct:5.1f}%)")
        print(f"   2  | {L2:10.0f} |  {R2:10.0f} |   {axle2:6.0f}  ({axle2_pct:5.1f}%)")
        print(f"   3  | {L3:10.0f} |  {R3:10.0f} |   {axle3:6.0f}  ({axle3_pct:5.1f}%)")
        print("-----------------------------------------------------------")
        print(f" Total vehicle vertical load [kg]: {total:8.0f}  ({(axle1_pct+axle2_pct+axle3_pct):5.1f}%)\n\n")

    # -------- internal helpers --------

    def _resolve_handles(self) -> None:
        # Resolve wheels
        self._wheels.clear()
        for wp in self._wheel_paths:
            rb = self._sim.getRigidBody(wp)
            if rb is None:
                # Keep going; scene may not be fully created yet.
                continue
            self._wheels[wp] = rb

        # Resolve terrain: prefer rigid body, optionally fall back to geometry
        terrain_body = self._sim.getRigidBody(self._terrain_path)
        if terrain_body is not None:
            self._terrain = terrain_body
            return

        if not self._use_geometry_if_no_body:
            self._terrain = None
            return

        # Fallback: geometry in the simulation space
        space = self._sim.getSpace()
        try:
            terrain_geom = space.getGeometry(self._terrain_path)
        except Exception:
            terrain_geom = None

        self._terrain = terrain_geom
