"""Study A2 physics-setting overrides for the frozen insertion adapters.

Wraps `sweep_physx.setup` (PhysX) and `grasp_native.setup` (MuJoCo CPU) without
editing them. An empty override is a pure pass-through: same MJCF bytes, same
stage, same stepping. Every requested value is read back from the compiled
MuJoCo model or the live USD stage and checked before the first episode.

Physics substeps keep the 1 kHz controller: one effort command is held for k
physics steps of dt/k. Contacts are returned as the mean over the 1 ms control
interval (each substep's force divided by k), which is what the 1 ms run reports
(impulse / 1 ms). Raw 40 N / 4 Nm stops therefore see the same quantity.
"""
import math
import xml.etree.ElementTree as ET

import numpy as np

COMMON = {'substeps', 'friction'}
PHYSX = {'contact_offset_m', 'rest_offset_m', 'position_iterations', 'velocity_iterations',
         'solver', 'max_depenetration_velocity_m_s'}
MUJOCO = {'timeconst_s', 'dampratio', 'solimp', 'cone', 'impratio', 'noslip_iterations'}


def validate(physics, backend, control_dt=.001):
    physics = dict(physics or {})
    engine = PHYSX if backend.startswith('physx') else MUJOCO if backend == 'mujoco_cpu' else None
    if engine is None:
        raise ValueError('Ablation overrides support physx, physx_gpu and mujoco_cpu only')
    unknown = set(physics)-COMMON-engine
    if unknown:
        raise ValueError(f'Overrides not valid for {backend}: {sorted(unknown)}')
    k = physics.setdefault('substeps', 1)
    if type(k) is not int or not 1 <= k <= 16:
        raise ValueError('substeps must be an integer in [1, 16]')
    def positive(key, zero=False):
        if key in physics:
            v = physics[key]
            if type(v) not in (int, float) or not math.isfinite(v) or v < 0 or (v == 0 and not zero):
                raise ValueError(f'{key} must be finite and {"nonnegative" if zero else "positive"}')
            physics[key] = float(v)
    for key in ('contact_offset_m', 'max_depenetration_velocity_m_s', 'timeconst_s', 'dampratio', 'impratio'):
        positive(key)
    positive('friction', zero=True)
    positive('rest_offset_m', zero=True)
    for key, low in (('position_iterations', 1), ('velocity_iterations', 0), ('noslip_iterations', 0)):
        if key in physics and (type(physics[key]) is not int or not low <= physics[key] <= 255):
            raise ValueError(f'{key} must be an integer in [{low}, 255]')
    if physics.get('solver', 'TGS') not in ('TGS', 'PGS'):
        raise ValueError('solver must be TGS or PGS')
    if physics.get('cone', 'pyramidal') not in ('pyramidal', 'elliptic'):
        raise ValueError('cone must be pyramidal or elliptic')
    if physics.get('impratio', 1.) < 1:
        raise ValueError('impratio must be >= 1')
    if 'solimp' in physics:
        s = [float(x) for x in physics['solimp']]
        if len(s) != 5 or not (0 < s[0] <= s[1] < 1 and s[2] > 0 and 0 < s[3] < 1 and s[4] >= 1):
            raise ValueError('solimp must be (dmin <= dmax < 1, width > 0, midpoint in (0,1), power >= 1)')
        physics['solimp'] = s
    # MuJoCo silently clamps timeconst to 2 x timestep (refsafe). A clamped
    # level would be mislabelled, so it is rejected instead of run.
    if 'timeconst_s' in physics and physics['timeconst_s'] < 2*control_dt/k-1e-12:
        raise ValueError(f'timeconst {physics["timeconst_s"]} s is below 2 x substep and would be clamped')
    return physics


def patch_mjcf(source, physics, dt):
    """Return MJCF with overrides applied; unchanged bytes when nothing applies."""
    k = physics['substeps']
    if k == 1 and not set(physics) & (MUJOCO | {'friction'}):
        return source
    tree = ET.fromstring(source)
    option = tree.find('option')
    if option is None:
        option = ET.SubElement(tree, 'option')
    option.set('timestep', repr(dt/k))
    for key in ('cone', 'impratio'):
        if key in physics:
            option.set(key, str(physics[key]))
    if 'noslip_iterations' in physics:
        option.set('noslip_iterations', str(physics['noslip_iterations']))
    for geom in tree.iter('geom'):
        if 'friction' in physics:
            values = (geom.get('friction') or '1 0.005 0.0001').split()
            values[0] = repr(physics['friction'])
            geom.set('friction', ' '.join(values))
        if 'timeconst_s' in physics or 'dampratio' in physics:
            current = [float(x) for x in (geom.get('solref') or '0.02 1').split()]
            geom.set('solref', f'{physics.get("timeconst_s", current[0])!r} {physics.get("dampratio", current[1])!r}')
        if 'solimp' in physics:
            geom.set('solimp', ' '.join(map(repr, physics['solimp'])))
    return ET.tostring(tree, encoding='unicode')


def unique(values, digits=12):
    """Collapse per-element readbacks to {value: count}; JSON-safe for inf."""
    counts = {}
    for v in values:
        if isinstance(v, (list, tuple, np.ndarray)):
            key = ' '.join(format(float(x), f'.{digits}g') for x in v)
        elif isinstance(v, (bool, np.bool_)) or v is None or isinstance(v, str):
            key = str(v)
        else:
            key = format(float(v), f'.{digits}g')
        counts[key] = counts.get(key, 0)+1
    return counts


def mujoco_readback(model):
    import mujoco
    collide = [g for g in range(model.ngeom) if model.geom_contype[g] or model.geom_conaffinity[g]]
    refsafe = not model.opt.disableflags & mujoco.mjtDisableBit.mjDSBL_REFSAFE
    timeconst = model.geom_solref[collide, 0]
    effective = np.where(timeconst > 0, np.maximum(timeconst, 2*model.opt.timestep), timeconst) if refsafe else timeconst
    return dict(source='compiled MjModel of the exact MJCF passed to the native adapter',
                timestep_s=float(model.opt.timestep), integrator=int(model.opt.integrator),
                solver=int(model.opt.solver), iterations=int(model.opt.iterations), tolerance=float(model.opt.tolerance),
                cone=['pyramidal', 'elliptic'][int(model.opt.cone)], impratio=float(model.opt.impratio),
                noslip_iterations=int(model.opt.noslip_iterations), refsafe=bool(refsafe),
                collision_geoms=len(collide), friction=unique(model.geom_friction[collide]),
                solref=unique(model.geom_solref[collide]), solimp=unique(model.geom_solimp[collide]),
                effective_timeconst_s=unique(effective), margin=unique(model.geom_margin[collide]))


def physx_apply(world, physics):
    from pxr import PhysxSchema, UsdPhysics
    if 'solver' in physics:
        world.get_physics_context().set_solver_type(physics['solver'])
    if 'friction' in physics:
        material = UsdPhysics.MaterialAPI(world.stage.GetPrimAtPath('/World/material'))
        material.GetStaticFrictionAttr().Set(physics['friction'])
        material.GetDynamicFrictionAttr().Set(physics['friction'])
    for prim in world.stage.Traverse():
        if prim.HasAPI(PhysxSchema.PhysxCollisionAPI):
            api = PhysxSchema.PhysxCollisionAPI(prim)
            if 'contact_offset_m' in physics:
                api.GetContactOffsetAttr().Set(physics['contact_offset_m'])
            if 'rest_offset_m' in physics:
                api.GetRestOffsetAttr().Set(physics['rest_offset_m'])
        if prim.HasAPI(PhysxSchema.PhysxArticulationAPI):
            api = PhysxSchema.PhysxArticulationAPI(prim)
            if 'position_iterations' in physics:
                api.GetSolverPositionIterationCountAttr().Set(physics['position_iterations'])
            if 'velocity_iterations' in physics:
                api.GetSolverVelocityIterationCountAttr().Set(physics['velocity_iterations'])
        if prim.HasAPI(PhysxSchema.PhysxRigidBodyAPI) and 'max_depenetration_velocity_m_s' in physics:
            PhysxSchema.PhysxRigidBodyAPI(prim).CreateMaxDepenetrationVelocityAttr(physics['max_depenetration_velocity_m_s'])


def physx_readback(world):
    from pxr import PhysxSchema, UsdPhysics
    stage = world.stage
    out = dict(source='live USD stage after world.reset(); unauthored attributes report the schema fallback',
               physics_dt_s=float(world.get_physics_dt()))
    scenes = [p for p in stage.Traverse() if p.IsA(UsdPhysics.Scene)]
    if len(scenes) != 1:
        raise ValueError(f'Expected one physics scene, found {len(scenes)}')
    scene = PhysxSchema.PhysxSceneAPI(scenes[0])
    out.update(solver=str(scene.GetSolverTypeAttr().Get()), gpu_dynamics=bool(scene.GetEnableGPUDynamicsAttr().Get()),
               time_steps_per_second=scene.GetTimeStepsPerSecondAttr().Get(),
               broadphase=str(scene.GetBroadphaseTypeAttr().Get()))
    material = stage.GetPrimAtPath('/World/material')
    out['material'] = dict(static_friction=float(UsdPhysics.MaterialAPI(material).GetStaticFrictionAttr().Get()),
                           dynamic_friction=float(UsdPhysics.MaterialAPI(material).GetDynamicFrictionAttr().Get()),
                           friction_combine=str(PhysxSchema.PhysxMaterialAPI(material).GetFrictionCombineModeAttr().Get())
                           if material.HasAPI(PhysxSchema.PhysxMaterialAPI) else 'unauthored (PhysX default: average)')
    contact, rest, position, velocity, depenetration, authored = [], [], [], [], [], []
    for prim in stage.Traverse():
        if prim.HasAPI(PhysxSchema.PhysxCollisionAPI):
            api = PhysxSchema.PhysxCollisionAPI(prim)
            contact.append(api.GetContactOffsetAttr().Get()); rest.append(api.GetRestOffsetAttr().Get())
        if prim.HasAPI(PhysxSchema.PhysxArticulationAPI):
            api = PhysxSchema.PhysxArticulationAPI(prim)
            position.append(api.GetSolverPositionIterationCountAttr().Get())
            velocity.append(api.GetSolverVelocityIterationCountAttr().Get())
        if prim.HasAPI(PhysxSchema.PhysxRigidBodyAPI):
            attr = PhysxSchema.PhysxRigidBodyAPI(prim).GetMaxDepenetrationVelocityAttr()
            depenetration.append(attr.Get()); authored.append(attr.HasAuthoredValue())
    out.update(contact_offset_m=unique(contact), rest_offset_m=unique(rest),
               position_iterations=unique(position), velocity_iterations=unique(velocity),
               max_depenetration_velocity_m_s=unique(depenetration),
               max_depenetration_velocity_authored=unique(authored))
    return out


def check(physics, readback, physics_dt):
    """Fail closed unless every requested value is what the engine received."""
    def only(counts, expected, name):
        values = [float(v) if v not in ('True', 'False') else v for v in counts]
        if len(values) != 1 or (isinstance(expected, str) and values[0] != expected) or \
                (not isinstance(expected, str) and not math.isclose(values[0], expected, rel_tol=1e-6, abs_tol=1e-12)):
            raise ValueError(f'Readback mismatch for {name}: requested {expected}, got {counts}')
    if 'physics_dt_s' in readback:
        if not math.isclose(readback['physics_dt_s'], physics_dt, rel_tol=1e-9):
            raise ValueError(f'PhysX step {readback["physics_dt_s"]} != {physics_dt}')
        for key in ('contact_offset_m', 'rest_offset_m', 'position_iterations', 'velocity_iterations',
                    'max_depenetration_velocity_m_s'):
            if key in physics:
                only(readback[key], physics[key], key)
        if 'solver' in physics and readback['solver'] != physics['solver']:
            raise ValueError(f'Readback mismatch for solver: {readback["solver"]}')
        if 'friction' in physics:
            for key in ('static_friction', 'dynamic_friction'):
                only({readback['material'][key]: 1}, physics['friction'], key)
        return
    if not math.isclose(readback['timestep_s'], physics_dt, rel_tol=1e-9):
        raise ValueError(f'MuJoCo step {readback["timestep_s"]} != {physics_dt}')
    for key in ('cone', 'impratio', 'noslip_iterations'):
        if key in physics and readback[key] != physics[key]:
            raise ValueError(f'Readback mismatch for {key}: {readback[key]}')
    if 'friction' in physics:
        only({k.split()[0]: 1 for k in readback['friction']}, physics['friction'], 'friction')
    if 'timeconst_s' in physics:
        only({k.split()[0]: 1 for k in readback['solref']}, physics['timeconst_s'], 'timeconst_s')
        only(readback['effective_timeconst_s'], physics['timeconst_s'], 'effective_timeconst_s')
    if 'solimp' in physics:
        if list(readback['solimp']) != [' '.join(format(x, '.12g') for x in physics['solimp'])]:
            raise ValueError(f'Readback mismatch for solimp: {readback["solimp"]}')


def merge_checks(checks):
    """Combine per-substep native checks over one control interval."""
    merged = dict(checks[-1]) if checks else {}
    for key in ('penetration_m', 'count', 'normal_matrix_error_N'):
        values = [c[key] for c in checks if c.get(key) is not None]
        if values:
            merged[key] = max(values)
    arms = [c['constraint_arm_Nm'] for c in checks if c.get('constraint_arm_Nm') is not None]
    if arms:
        merged['constraint_arm_Nm'] = np.mean(arms, axis=0).tolist()
    return merged


def substepped(state, step, k):
    """Hold one effort for k physics steps; return interval-mean contact forces."""
    if k == 1:
        return step
    native_check = getattr(state, 'native_contact_check', None)
    merged = {}

    def advance(u):
        contacts, checks = [], []
        for _ in range(k):
            contacts += [(a, b, p, np.asarray(f, dtype=float)/k) for a, b, p, f in step(u)]
            if native_check is not None:
                checks.append(native_check())
        merged.clear(); merged.update(merge_checks(checks))
        return contacts
    if native_check is not None:
        state.native_contact_check = lambda: dict(merged)
    return advance


def install(physics, backend):
    """Replace the module-level insertion setups used by rl_env for this process."""
    physics = validate(physics, backend)
    k = physics['substeps']
    if backend.startswith('physx'):
        import sweep_physx
        original = sweep_physx.setup

        def setup(engine, source, spec, dt, seed, extra_setup=None, **kwargs):
            if 'bore_origin_m' not in spec:
                raise ValueError('Ablation overrides are declared for the insertion scene only')
            if 'rest_offset_m' in physics and physics['rest_offset_m'] >= physics.get('contact_offset_m', spec['contact_offset_m']):
                raise ValueError('rest offset must be below contact offset')
            captured = {}

            def configure(world):
                captured['world'] = world
                physx_apply(world, physics)
                if extra_setup is not None:
                    extra_setup(world)
            # A wider contact envelope reaches more bore facets; the adapter's
            # 256-slot readback would fail closed (a buffer limit, not physics).
            capacity = 4096 if 'contact_offset_m' in physics else spec.get('contact_readback_capacity', 256)
            if capacity != spec.get('contact_readback_capacity', 256):
                spec = dict(spec, contact_readback_capacity=capacity)
            state, step, effective, app = original(engine, source, spec, dt/k, seed, extra_setup=configure, **kwargs)
            readback = physx_readback(captured['world'])
            check(physics, readback, dt/k)
            effective['ablation'] = dict(requested=physics, control_dt_s=dt, readback=readback,
                                         contact_readback_capacity=capacity)
            return state, substepped(state, step, k), effective, app
        sweep_physx.setup = setup
    else:
        import mujoco
        import grasp_native
        original = grasp_native.setup

        def setup(engine, source, spec, dt, seed):
            if 'bore_origin_m' not in spec:
                raise ValueError('Ablation overrides are declared for the insertion scene only')
            patched = patch_mjcf(source, physics, dt)
            readback = mujoco_readback(mujoco.MjModel.from_xml_string(patched))
            check(physics, readback, dt/k)
            state, step, effective, app = original(engine, patched, spec, dt, seed)
            # Live cross-check on the adapter's own compiled meshes.
            for shape in state.native_geometry()['shapes']:
                if 'friction' in physics and not math.isclose(shape['friction'][0], physics['friction']):
                    raise ValueError('Live MuJoCo friction differs from request')
                if 'timeconst_s' in physics and not math.isclose(shape['solref'][0], physics['timeconst_s']):
                    raise ValueError('Live MuJoCo solref differs from request')
            effective['ablation'] = dict(requested=physics, control_dt_s=dt, readback=readback,
                                         mjcf_changed=patched != source)
            return state, substepped(state, step, k), effective, app
        grasp_native.setup = setup
    return physics