"""Friction teaching module; standard-library solvers with explicit degeneracy handling.
Run: python3 examples/friction_contact.py
"""
import math,json
from convex_contact import solve,objective


def add_diagonal(A,value):
    return [[x+(value if i==j else 0) for j,x in enumerate(row)] for i,row in enumerate(A)]


def fixed_normal_qcqp(H,b,mu,normal):
    if normal<0 or any(x<0 for x in mu):raise ValueError('nonnegative scales required')
    active=[i for i,x in enumerate(mu) if x>0]
    if normal==0 or not active:
        return dict(force=[0.0]*len(mu),multiplier=None if normal==0 else 0.0,active=normal==0)
    B=[[mu[i]*H[i][j]*mu[j] for j in active] for i in active]
    c=[mu[i]*b[i] for i in active]
    def evaluate(multiplier):
        y=solve(add_diagonal(B,multiplier),[-v for v in c])
        return y,sum(v*v for v in y)-normal*normal
    y,value=evaluate(0)
    multiplier=0.0
    if value>0:
        low,high=0.0,1.0
        while evaluate(high)[1]>0:high*=2
        # Independent from the browser's Newton derivative and Cholesky implementation.
        for _ in range(110):
            multiplier=(low+high)/2
            y,value=evaluate(multiplier)
            if value>0:low=multiplier
            else:high=multiplier
        multiplier=(low+high)/2
        y,value=evaluate(multiplier)
    result=[0.0]*len(mu)
    for i,v in zip(active,y):result[i]=mu[i]*v
    return dict(force=result,multiplier=multiplier,active=multiplier>0)


def sliding(normal,mu,velocity):
    length=math.sqrt(sum((m*v)**2 for m,v in zip(mu,velocity)))
    force=[-normal*m*m*v/length for m,v in zip(mu,velocity)] if length else [0.0]*len(mu)
    return force,sum(f*v for f,v in zip(force,velocity))


def decode_pyramid(alpha,mu):
    return [sum(alpha)]+[m*(alpha[2*i]-alpha[2*i+1]) for i,m in enumerate(mu)]


def encode_pyramid(force,mu):
    normal=force[0]
    if not mu:return [normal]
    beta=[force[i+1]/m if m else 0 for i,m in enumerate(mu)]
    if any(m==0 and force[i+1]!=0 for i,m in enumerate(mu)):return None
    slack=normal-sum(abs(x) for x in beta)
    if slack < -1e-12:return None
    share=max(0,slack)/(2*len(mu))
    return [v for x in beta for v in (max(x,0)+share,max(-x,0)+share)]


def soft_cone(jar,Rn,mu,ratio=1):
    positive=[m for m in mu if m>0]
    if not positive:
        fn=max(0,-jar[0]/Rn)
        return dict(force=[fn]+[0.0]*len(mu),cost=.5*Rn*fn*fn,region='top' if jar[0]>=0 else 'bottom')
    kappa=Rn*positive[0]**2/ratio
    u=[m*j for m,j in zip(mu,jar[1:])]
    T=math.sqrt(sum(x*x for x in u));jn=jar[0]
    if jn>=T:return dict(force=[0.0]*len(jar),cost=0.0,region='top')
    if kappa*jn+Rn*T<=0:
        f=[-jn/Rn]+[-m*x/kappa for m,x in zip(mu,u)]
        return dict(force=f,cost=.5*jn*jn/Rn+.5*T*T/kappa,region='bottom')
    fn=(T-jn)/(Rn+kappa)
    return dict(force=[fn]+[-fn*m*x/T for m,x in zip(mu,u)],cost=.5*(T-jn)**2/(Rn+kappa),region='middle')


if __name__=='__main__':
    H,b,mu=[[2,.7],[.7,1]],[-3,-1],[1,.5]
    result=fixed_normal_qcqp(H,b,mu,1)
    free=solve(H,[-v for v in b]);length=math.sqrt(sum((x/m)**2 for x,m in zip(free,mu)))
    radial=[x/max(1,length) for x in free]
    print(json.dumps({'example':result,'objective':objective(H,b,result['force']),
        'radial':radial,'radial_objective':objective(H,b,radial),
        'pyramid':encode_pyramid([10,2,-.5],[.5,.25]),
        'soft_regions':[soft_cone([n,2,0],.1,[.5,.5]) for n in [2,-1,-5]]},indent=2))