"""Original source-authored feedback controller; no model imports or learned weights.

Experimental simulation artifact, unscored. Raw Arena obs -> residual joint targets.
Only ONNX and NumPy are required. This script never runs a simulator or network call.
"""
from pathlib import Path
import hashlib,json
import numpy as np
import onnx
from onnx import TensorProto,helper,numpy_helper

# Interface joint order: left hip yaw/roll/pitch/knee/ankle; neck/head four;
# right hip yaw/roll/pitch/knee/ankle. Parameters below are authored here.
# These are conservative initial hypotheses, not identified control gains.
PARAMS=dict(knee_flex=0.08,back_lean=0.02,pitch_kp=0.12,pitch_kd=0.02,
            roll_kp=0.10,roll_kd=0.02,joint_damping=0.003)

def coefficients(params=PARAMS):
    w=np.zeros((61,14),np.float32); b=np.zeros(14,np.float32)
    c=params['knee_flex']
    # Mirrored sagittal axes; +hip -knee +ankle is the left-chain rotation.
    b[[2,3,4]]=[c/2,c,c/2+params['back_lean']]
    b[[11,12,13]]=[-c/2,-c,-c/2-params['back_lean']]
    for j,sign in [(4,1),(13,-1)]:
        w[3,j]=sign*params['pitch_kp']  # projected gravity x ~ trunk pitch
        w[1,j]=sign*params['pitch_kd'] # gyro y, corrective ankle response
    for j in (1,10):
        w[4,j]=-params['roll_kp']      # projected gravity y ~ -trunk roll
        w[0,j]=params['roll_kd']
    active=[1,2,3,4,10,11,12,13]
    for j in active: w[20+j,j]=-params['joint_damping']
    limits=np.zeros(14,np.float32)
    limits[active]=[.12,.15,.30,.20,.12,.15,.30,.20]
    return w,b,-limits,limits

def act(obs,params=PARAMS):
    w,b,lo,hi=coefficients(params)
    return np.clip(np.asarray(obs,np.float32)@w+b,lo,hi)

def build(path,params=PARAMS):
    w,b,lo,hi=coefficients(params)
    constants=[numpy_helper.from_array(a,n) for a,n in [(w,'feedback'),(b,'pose_bias'),(lo,'lower'),(hi,'upper')]]
    graph=helper.make_graph([
        helper.make_node('MatMul',['obs','feedback'],['response']),
        helper.make_node('Add',['response','pose_bias'],['raw_action']),
        helper.make_node('Max',['raw_action','lower'],['floored_action']),
        helper.make_node('Min',['floored_action','upper'],['action'])],
        'original_observation_feedback_brace',
        [helper.make_tensor_value_info('obs',TensorProto.FLOAT,[1,61])],
        [helper.make_tensor_value_info('action',TensorProto.FLOAT,[1,14])],constants)
    model=helper.make_model(graph,producer_name='source-authored CPU brace prototype',opset_imports=[helper.make_opsetid('',13)])
    model.ir_version=8
    model.doc_string='Unscored original affine feedback brace. No reference policy weights, clock, seeds, recurrence, or training.'
    onnx.checker.check_model(model)
    onnx.save(model,str(path))
    return {'parameters':params,'sha256':hashlib.sha256(path.read_bytes()).hexdigest(),'bytes':path.stat().st_size,
            'learned_weights':False,'reference_weights_used':False,'simulated':False}

if __name__=='__main__':
    root=Path(__file__).resolve().parent
    info=build(root/'policy.onnx')
    (root/'build.json').write_text(json.dumps(info,indent=2)+'\n')
    print(json.dumps(info,indent=2))