"""MuJoCo 3.10 scalar soft-constraint parameters; Python standard library only.

error is pos-margin (metres for a translational contact), velocity is J v.
reference returns a target for J a, subtracting explicitly supplied curvature.
R is a base scalar regularizer; contact_regularization applies cone adjustments.
"""
import math

MINVAL = 1e-15
MINIMP, MAXIMP = 0.0001, 0.9999


def parameters(solimp):
    """Apply getsolparam's solimp bounds."""
    d0, dw, width, midpoint, power = map(float, solimp)
    clamp = lambda x: min(MAXIMP, max(MINIMP, x))
    return clamp(d0), clamp(dw), max(0, width), clamp(midpoint), max(1, power)


def impedance(error, solimp):
    """Return value and implementation derivative with respect to signed error."""
    d0, dw, width, midpoint, power = parameters(solimp)
    if d0 == dw or width <= MINVAL:
        return {"value": (d0+dw)/2, "derivative": 0.0}
    x = abs(error/width)
    if x >= 1 or x <= 0:
        return {"value": dw if x >= 1 else d0, "derivative": 0.0}
    if power == 1:
        y, yp = x, 1
    elif x <= midpoint:
        y = x**power/midpoint**(power-1)
        yp = power*(x/midpoint)**(power-1)
    else:
        y = 1-(1-x)**power/(1-midpoint)**(power-1)
        yp = power*((1-x)/(1-midpoint))**(power-1)
    return {"value": d0+(dw-d0)*y,
            "derivative": (dw-d0)*yp*math.copysign(1, error)/width}


def reference(error, velocity, solref, solimp, timestep, diag_approx=0.5,
              refsafe=True, curvature=0.0, friction=False):
    """Scalar KBIP/R/D/aref. friction=True sets K=0 (elliptic tangent/dry friction)."""
    ref = list(map(float, solref))
    if (ref[0] > 0) != (ref[1] > 0):
        ref = [0.02, 1.0]
    if refsafe and ref[0] > 0:
        ref[0] = max(ref[0], 2*timestep)
    dw = parameters(solimp)[1]
    k = (1/max(MINVAL, dw*dw*ref[0]**2*ref[1]**2)
         if ref[0] > 0 else -ref[0]/max(MINVAL, dw*dw))
    b = (2/max(MINVAL, dw*ref[0])
         if ref[1] > 0 else -ref[1]/max(MINVAL, dw))
    if friction:
        k = 0.0
    imp = impedance(error, solimp)
    d = imp["value"]
    r = max(MINVAL, (1-d)*diag_approx/d)
    return {"K": k, "B": b, "d": d, "dprime": imp["derivative"],
            "R": r, "D": 1/r, "aref": -b*velocity-k*d*error-curvature,
            "solref_used": ref}


def contact_regularization(normal_r, friction, impratio=1.0, pyramidal=False):
    """friction contains dim-1 positive sliding/spin/rolling coefficients."""
    if not friction:
        return [normal_r]
    if any(mu <= 0 for mu in friction):
        raise ValueError("Friction coefficients must be positive")
    tangent_r = normal_r/max(MINVAL, impratio)
    if pyramidal:
        # Caller supplies the first edge's base R, as mj_makeImpedance does.
        return [2*friction[0]**2*tangent_r]*(2*len(friction))
    return [normal_r]+[tangent_r*friction[0]**2/mu**2 for mu in friction]


def scalar_force(a_free, target, a_inverse, regularizer, unilateral=True):
    """Independent row: minimizer of .5(A+R)f²+(a_free-target)f."""
    force = (target-a_free)/(a_inverse+regularizer)
    if unilateral:
        force = max(0.0, force)
    return {"force": force, "acceleration": a_free+a_inverse*force}


if __name__ == "__main__":
    from efc import model_data
    c = model_data()
    error = c["position"][2]-c["radius"]-c["margin"]
    ref = reference(error, c["velocity"][2], c["solref"], c["solimp"],
                    c["timestep"], 1/c["mass"])
    print({"error": error, **ref})
    print(scalar_force(c["gravity"][2], ref["aref"], 1/c["mass"], ref["R"]))


def normalized_response(time, timeconst, ratio):
    """Exact r(t)/r(0), with zero initial speed and fixed d=dw, A=diagApprox."""
    decay = 1/timeconst
    if abs(ratio-1) < 1e-8:
        return math.exp(-decay*time)*(1+decay*time)
    if ratio < 1:
        frequency = decay*math.sqrt(1/ratio**2-1)
        return math.exp(-decay*time)*(math.cos(frequency*time)+decay/frequency*math.sin(frequency*time))
    h = math.sqrt(1-1/ratio**2)
    fast, slow = -(1+h)*decay, -(1-h)*decay
    return (fast*math.exp(slow*time)-slow*math.exp(fast*time))/(fast-slow)


def stationary_error(solref, solimp, timestep, gravity=9.81, diag_approx=0.5,
                     a_inverse=0.5, refsafe=True):
    """Bisection for a static active normal row with a_free=-gravity, K>0."""
    def acceleration(error):
        ref=reference(error,0,solref,solimp,timestep,diag_approx,refsafe)
        return scalar_force(-gravity,ref['aref'],a_inverse,ref['R'])['acceleration']
    low, high=-0.001,0.0
    while acceleration(low)<0:
        low*=2
        if abs(low)>1e6:
            raise ValueError('No finite restoring equilibrium')
    for _ in range(90):
        mid=(low+high)/2
        if acceleration(mid)>0:
            low=mid
        else:
            high=mid
    return (low+high)/2