"""Small quadratic solvers and time integrators used in chapter 15.

Matrices are flat row-major arrays. These examples use ordinary scalar
coordinates; rotation updates are provided by rotation.py / quaternion.py.
"""
from math import cos, exp, sin, sqrt


def dot(a, b):
    return sum(x*y for x, y in zip(a, b))


def matvec(matrix, x):
    n = len(x)
    return [dot(matrix[n*i:n*(i+1)], x) for i in range(n)]


def objective(matrix, linear, x):
    return dot(x, matvec(matrix, x))/2 + dot(linear, x)


def solve2(matrix, right):
    a, b, c, d = matrix
    determinant = a*d-b*c
    if determinant <= 0 or a <= 0:
        raise ValueError("Expected a positive definite 2 by 2 matrix")
    return [(d*right[0]-b*right[1])/determinant,
            (a*right[1]-c*right[0])/determinant]


def complementarity_error(matrix, linear, force):
    gradient = [x+y for x, y in zip(matvec(matrix, force), linear)]
    return max(abs(x-max(0, x-g)) for x, g in zip(force, gradient))


def pgs(matrix, linear, start=None, sweeps=20):
    """Coordinate minimization on the nonnegative orthant; record each row."""
    n = len(linear)
    x = [0.0]*n if start is None else [max(0.0, a) for a in start]
    history = [{"x": x[:], "row": None, "cost": objective(matrix, linear, x)}]
    for _ in range(sweeps):
        for i in range(n):
            if matrix[n*i+i] <= 0:
                raise ValueError("Positive coordinate curvature required")
            residual = linear[i] + dot(matrix[n*i:n*(i+1)], x)
            x[i] = max(0, x[i]-residual/matrix[n*i+i])
            history.append({"x": x[:], "row": i, "cost": objective(matrix, linear, x)})
    return history


def orthant2(matrix, linear):
    """Enumerate the four active sets: independent reference for PGS."""
    a, b, c, d = matrix
    candidates = [[0.0, 0.0], [max(0, -linear[0]/a), 0],
                  [0, max(0, -linear[1]/d)]]
    unconstrained = solve2(matrix, [-v for v in linear])
    if min(unconstrained) >= 0:
        candidates.append(unconstrained)
    return min(candidates, key=lambda x: objective(matrix, linear, x))


def quadratic_descent(matrix, linear, start, method="cg", iterations=12):
    """Fixed SPD quadratic: exact line-search steepest descent, CG or Newton."""
    x = start[:]
    g = [a+b for a, b in zip(matvec(matrix, x), linear)]
    p = [-v for v in g]
    history = [{"x": x[:], "cost": objective(matrix, linear, x),
                "gradient": g[:], "direction": p[:], "alpha": 0}]
    for _ in range(iterations):
        if dot(g, g) < 1e-26:
            break
        if method == "newton":
            p = solve2(matrix, [-v for v in g])
            alpha = 1.0
        else:
            if method == "gradient":
                p = [-v for v in g]
            alpha = -dot(g, p)/dot(p, matvec(matrix, p))
        x = [a+alpha*b for a, b in zip(x, p)]
        next_g = [a+b for a, b in zip(matvec(matrix, x), linear)]
        history.append({"x": x[:], "cost": objective(matrix, linear, x),
                        "gradient": next_g[:], "direction": p[:], "alpha": alpha})
        beta = dot(next_g, next_g)/dot(g, g)
        p = [-a+beta*b for a, b in zip(next_g, p)]
        g = next_g
    return history


def primal_normal(mass, free, jacobian, reference, regularizer, acceleration):
    """Primal Gauss objective with independent soft normal contact rows."""
    n, count = len(free), len(reference)
    delta = [a-b for a, b in zip(acceleration, free)]
    gradient = matvec(mass, delta)
    hessian = mass[:]
    cost = dot(delta, gradient)/2
    force = []
    for i in range(count):
        row = jacobian[n*i:n*(i+1)]
        residual = dot(row, acceleration)-reference[i]
        active = residual < 0
        value = -residual/regularizer[i] if active else 0
        force.append(value)
        cost += min(0, residual)**2/(2*regularizer[i])
        for j in range(n):
            gradient[j] -= row[j]*value
            if active:
                for k in range(n):
                    hessian[n*j+k] += row[j]*row[k]/regularizer[i]
    return {"cost": cost, "gradient": gradient, "hessian": hessian, "force": force}


def primal_newton(mass, free, jacobian, reference, regularizer, start, iterations=16):
    """Two-dimensional piecewise quadratic primal; damped Newton with Armijo."""
    x, history = start[:], []
    for _ in range(iterations+1):
        state = primal_normal(mass, free, jacobian, reference, regularizer, x)
        history.append({"x": x[:], **state})
        if dot(state["gradient"], state["gradient"]) < 1e-24:
            break
        direction = solve2(state["hessian"], [-v for v in state["gradient"]])
        slope = dot(state["gradient"], direction)
        alpha = 1.0
        while alpha > 2**-30:
            trial = [a+alpha*b for a, b in zip(x, direction)]
            value = primal_normal(mass, free, jacobian, reference, regularizer, trial)["cost"]
            if value <= state["cost"] + 1e-4*alpha*slope:
                break
            alpha /= 2
        x = trial
        history[-1].update(alpha=alpha, direction=direction)
    return history


def oscillator_rhs(state, frequency=2.0, damping=0.0):
    q, v = state
    return [v, -frequency**2*q-damping*v]


def rk4_step(state, step, function):
    """Return the step and four evaluations, each with its actual trial state."""
    k1 = function(state)
    y2 = [x+step*k/2 for x, k in zip(state, k1)]
    k2 = function(y2)
    y3 = [x+step*k/2 for x, k in zip(state, k2)]
    k3 = function(y3)
    y4 = [x+step*k for x, k in zip(state, k3)]
    k4 = function(y4)
    result = [x+step*(a+2*b+2*c+d)/6 for x, a, b, c, d in zip(state, k1, k2, k3, k4)]
    return result, [{"fraction": t, "state": y, "slope": k} for t, y, k in
                    [(0, state[:], k1), (.5, y2, k2), (.5, y3, k3), (1, y4, k4)]]


def oscillator_step(state, step, frequency=2.0, damping=0.0, method="semi"):
    q, v = state
    a = -frequency**2*q-damping*v
    if method == "explicit":
        return [q+step*v, v+step*a]
    if method == "semi":
        new_v = v+step*a
        return [q+step*new_v, new_v]
    if method == "velocity":
        new_v = (v-step*frequency**2*q)/(1+step*damping)
        return [q+step*new_v, new_v]
    if method == "implicit":
        new_v = (v-step*frequency**2*q)/(1+step*damping+step**2*frequency**2)
        return [q+step*new_v, new_v]
    if method == "rk4":
        return rk4_step(state, step, lambda y: oscillator_rhs(y, frequency, damping))[0]
    raise ValueError("Unknown integrator")


def oscillator_exact(state, time, frequency=2.0, damping=0.0):
    q, v = state
    half = damping/2
    discriminant = frequency**2-half**2
    if discriminant > 1e-14:
        w = sqrt(discriminant)
        c, s = cos(w*time), sin(w*time)/w
    elif discriminant < -1e-14:
        w = sqrt(-discriminant)
        c, s = (exp(w*time)+exp(-w*time))/2, (exp(w*time)-exp(-w*time))/(2*w)
    else:
        c, s = 1, time
    decay = exp(-half*time)
    return [decay*(q*c+(v+half*q)*s), decay*(v*c-(frequency**2*q+half*v)*s)]


def oscillator_trace(step=.1, frequency=2., damping=0., duration=6., method="semi"):
    state, history, time = [1., 0.], [], 0.
    while time < duration-1e-12:
        history.append({"time": time, "state": state[:],
                        "energy": (state[1]**2+frequency**2*state[0]**2)/2})
        h = min(step, duration-time)
        state = oscillator_step(state, h, frequency, damping, method)
        time += h
    history.append({"time": time, "state": state[:],
                    "energy": (state[1]**2+frequency**2*state[0]**2)/2})
    return history


def amplification(z, method):
    """Scalar test equation multiplier, with complex z = h*lambda."""
    if method == "explicit":
        return 1+z
    if method == "implicit":
        return 1/(1-z)
    if method == "rk4":
        return 1+z+z*z/2+z**3/6+z**4/24
    raise ValueError("Unknown scalar amplification function")


def cone2(diagonal, linear, friction):
    """Exact minimization over fn >= 0, |ft| <= mu*fn; 2D elliptic-cone slice."""
    if min(diagonal) <= 0 or friction < 0:
        raise ValueError("Positive curvature and nonnegative friction required")
    n, t = [-b/h for b, h in zip(linear, diagonal)]
    candidates = [[0., 0.]]
    if n >= 0 and abs(t) <= friction*n:
        candidates.append([n, t])
    for sign in [-1, 1]:
        normal = max(0, -(linear[0]+sign*friction*linear[1]) /
                     (diagonal[0]+friction**2*diagonal[1]))
        candidates.append([normal, sign*friction*normal])
    def cost(x):
        return sum(h*a*a/2+b*a for h, b, a in zip(diagonal, linear, x))
    return min(candidates, key=cost)


def contact_model():
    import json
    from pathlib import Path
    text = (Path(__file__).resolve().parents[1]/"assets/contact-model.js").read_text()
    return json.loads(text.split("window.ContactModelData =", 1)[1].strip().rstrip(";"))


def contact_response(state, model=None, step=None):
    """One normal and one tangent of the shared translational sphere.

    The two-dimensional convex cone is minimized jointly. Rn = Rt uses
    isotropic friction and impratio = 1. The plane is fixed, orientation locked.
    """
    from soft_contact import reference
    model = contact_model() if model is None else model
    step = model["timestep"] if step is None else step
    x, z, vx, vz = state
    mass, g = model["mass"], -model["gravity"][2]
    distance = z-model["radius"]
    # Local MuJoCo 3.10.0: query distance margin+gap; active threshold margin.
    inclusion = model["margin"]
    error = distance-inclusion
    if error >= 0:
        return {"distance": distance, "error": error, "force": [0., 0.],
                "acceleration": [0., -g], "active": False}
    ref = reference(error, vz, model["solref"], model["solimp"], step, 1/mass)
    tangent_reference = -ref["B"]*vx
    force = cone2([1/mass+ref["R"]]*2,
                  [-g-ref["aref"], -tangent_reference], model["friction"])
    return {"distance": distance, "error": error, "force": force,
            "acceleration": [force[1]/mass, -g+force[0]/mass],
            "active": True, "reference": ref, "tangent_reference": tangent_reference}


def contact_trace(height=.6, speed=.8, step=None, duration=2., friction=None):
    model = contact_model()
    if friction is not None:
        model["friction"] = friction
    step = model["timestep"] if step is None else step
    if step <= 0 or duration < 0:
        raise ValueError("Positive step and nonnegative duration required")
    state, history, time = [0., height, speed, model["velocity"][2]], [], 0.
    while True:
        response = contact_response(state, model, step)
        history.append({"time": time, "state": state[:], **response})
        if time >= duration-1e-12:
            break
        h = min(step, duration-time)
        state[2] += h*response["acceleration"][0]
        state[3] += h*response["acceleration"][1]
        state[0] += h*state[2]
        state[1] += h*state[3]
        time += h
    return history


if __name__ == "__main__":
    H, b = [2, 1, 1, 2], [-1, -2]
    print("PGS:", pgs(H, b, sweeps=30)[-1]["x"], "reference:", orthant2(H, b))
    for name in ["explicit", "semi", "velocity", "implicit", "rk4"]:
        final = oscillator_trace(method=name)[-1]
        print(name, final)