import { DenseNetwork } from './network.js';

const DT = 1 / 60;
const clamp = (v, lo, hi) => Math.max(lo, Math.min(hi, v));
const wrap = a => Math.atan2(Math.sin(a), Math.cos(a));
const finiteVector = (v, n) => (Array.isArray(v) || ArrayBuffer.isView(v)) && v.length === n && Array.from(v).every(Number.isFinite);
const yaw = (v, a) => [Math.cos(a)*v[0]+Math.sin(a)*v[2], v[1], -Math.sin(a)*v[0]+Math.cos(a)*v[2]];

function quaternion(columns) {
  const [a,b,c] = columns, m00=a[0],m11=b[1],m22=c[2], trace=m00+m11+m22;
  let q;
  if (trace>0) { const s=2*Math.sqrt(trace+1);q=[(b[2]-c[1])/s,(c[0]-a[2])/s,(a[1]-b[0])/s,s/4]; }
  else if(m00>m11&&m00>m22) {const s=2*Math.sqrt(1+m00-m11-m22);q=[s/4,(b[0]+a[1])/s,(c[0]+a[2])/s,(b[2]-c[1])/s];}
  else if(m11>m22) {const s=2*Math.sqrt(1+m11-m00-m22);q=[(b[0]+a[1])/s,s/4,(c[1]+b[2])/s,(c[0]-a[2])/s];}
  else {const s=2*Math.sqrt(1+m22-m00-m11);q=[(c[0]+a[2])/s,(c[1]+b[2])/s,s/4,(a[1]-b[0])/s];}
  const length=Math.hypot(...q);return q.map(v=>v/length);
}

function rotation(value, offset) {
  const x=Array.from(value.subarray(offset,offset+3)), b=Array.from(value.subarray(offset+3,offset+6));
  const nx=Math.hypot(...x);if(nx<1e-6)throw Error('Degenerate predicted rotation.');
  for(let i=0;i<3;i++)x[i]/=nx;
  const dot=x.reduce((s,v,i)=>s+v*b[i],0),y=b.map((v,i)=>v-dot*x[i]),ny=Math.hypot(...y);
  if(ny<1e-6)throw Error('Degenerate predicted rotation.');
  for(let i=0;i<3;i++)y[i]/=ny;
  return [x,y,[x[1]*y[2]-x[2]*y[1],x[2]*y[0]-x[0]*y[2],x[0]*y[1]-x[1]*y[0]]];
}

/** Decode an absolute 143-value pose using the canonical rig. */
export function decodeMotionPose(output,rig,position=[0,0,0],facing=0) {
  if(!finiteVector(output,143)||!finiteVector(position,3)||!Number.isFinite(facing))throw Error('Invalid predicted motion pose.');
  const localQuaternions=[];let articulated=0;
  for(const end of rig.end_sites) {
    if(end){localQuaternions.push([0,0,0,1]);continue;}
    let columns=rotation(output,articulated*6);
    if(articulated===0)columns=columns.map(v=>yaw(v,facing));
    localQuaternions.push(quaternion(columns));articulated++;
  }
  const offset=yaw(Array.from(output.subarray(138,141)),facing);
  return {rootPosition:position.map((v,i)=>v+offset[i]),localQuaternions,
    contacts:Array.from(output.subarray(141),v=>1/(1+Math.exp(-v)))};
}

/** Fixed-step, engine-independent locomotion. Units: metres, seconds, radians. */
export class MotionController {
  constructor(manifest, weights, config) {
    if(config.format!=='superhq-motion-controller-v1'||config.dt!==DT||
       manifest.input_contract!=='trajectory22_causal_phase3_v1'||
       manifest.layers[0]?.inputs!==25||manifest.layers.at(-1)?.outputs!==143||
       manifest.original_feature_indices?.some((v,i)=>v!==843+i)||manifest.original_feature_indices?.length!==25||
       !finiteVector(config.cadence,2)||!Array.isArray(config.rig?.end_sites)||
       config.rig.end_sites.filter(v=>!v).length!==23||config.rig.names?.length!==config.rig.end_sites.length||
       ![config.max_speed_mps,config.acceleration_mps2,config.yaw_rate_radps].every(v=>Number.isFinite(v)&&v>0)) {
      throw Error('Incompatible Motion controller contract.');
    }
    this.network=new DenseNetwork(manifest,weights);this.config=config;this.rig=config.rig;
    this.reset();
  }

  static async load(baseURL, {fetch: fetcher=globalThis.fetch}={}) {
    const base=baseURL instanceof URL?baseURL:new URL(baseURL,globalThis.location?.href);
    async function request(path, binary=false) {
      const r=await fetcher(new URL(path,base));if(!r.ok)throw Error(`Cannot load ${path}: HTTP ${r.status}`);
      return binary?r.arrayBuffer():r.json();
    }
    const [manifest,config]=await Promise.all([request('network.json'),request('controller.json')]);
    if(manifest.weights!=='weights.bin')throw Error('Unexpected weight filename.');
    const weights=await request(manifest.weights,true);
    if(!globalThis.crypto?.subtle)throw Error('Weight verification requires Web Crypto (HTTPS or localhost).');
    const digest=Array.from(new Uint8Array(await crypto.subtle.digest('SHA-256',weights)),v=>v.toString(16).padStart(2,'0')).join('');
    if(digest!==manifest.sha256)throw Error('Motion weight checksum mismatch.');
    return new MotionController(manifest,weights,config);
  }

  reset({position=[0,0,0],facing=0,velocity=[0,0],phase=0}={}) {
    if(!finiteVector(position,3)||Math.abs(position[1])>1e-8||!finiteVector(velocity,2)||
       !Number.isFinite(facing)||!Number.isFinite(phase)||Math.hypot(...velocity)>this.config.max_speed_mps+1e-6) {
      throw Error('Expected a finite flat-ground reset transform and bounded velocity.');
    }
    this.position=Array.from(position);this.velocity=[velocity[0],0,velocity[1]];
    this.facing=facing;this.phase=phase;this.frame=0;
  }

  _intent(intent) {
    if(!finiteVector(intent?.velocity,2)||!Number.isFinite(intent.facing??this.facing))throw Error('Expected velocity [x,z] and finite facing radians.');
    const speed=Math.hypot(...intent.velocity),scale=Math.min(1,this.config.max_speed_mps/Math.max(speed,1e-8));
    return {desired:[intent.velocity[0]*scale,0,intent.velocity[1]*scale],facing:intent.facing??this.facing,clamped:speed>this.config.max_speed_mps};
  }

  /** Call once per 1/60 s. Optional resolver supplies a collision-resolved root. */
  step(intent, {resolveMovement}={}) {
    const {desired,facing,clamped}=this._intent(intent),oldPosition=this.position.slice(),oldFacing=this.facing;
    const change=desired.map((v,i)=>v-this.velocity[i]),distance=Math.hypot(...change),increment=this.config.acceleration_mps2*DT;
    const direction=change.map(v=>v/Math.max(distance,1e-8)),turn=wrap(facing-oldFacing);
    const points=[],velocities=[],angles=[];
    for(const n of [12,24,36]) {
      const accelerating=Math.min(n,Math.trunc(distance/increment));
      const sum=increment*accelerating*(accelerating+1)/2+(n-accelerating)*distance;
      points.push(oldPosition.map((v,i)=>v+DT*(n*this.velocity[i]+direction[i]*sum)));
      velocities.push(this.velocity.map((v,i)=>v+direction[i]*Math.min(n*increment,distance)));
      angles.push(oldFacing+clamp(turn,-n*this.config.yaw_rate_radps*DT,n*this.config.yaw_rate_radps*DT));
    }
    let nextVelocity=this.velocity.map((v,i)=>v+change[i]*Math.min(1,increment/Math.max(distance,1e-8)));
    let nextPosition=oldPosition.map((v,i)=>v+nextVelocity[i]*DT);
    let nextFacing=oldFacing+clamp(turn,-this.config.yaw_rate_radps*DT,this.config.yaw_rate_radps*DT);
    if(resolveMovement!==undefined) {
      if(typeof resolveMovement!=='function')throw Error('resolveMovement must be a function.');
      const resolved=resolveMovement({position:oldPosition.slice(),facing:oldFacing},{position:nextPosition.slice(),facing:nextFacing});
      if(!finiteVector(resolved?.position,3)||Math.abs(resolved.position[1])>1e-8||!Number.isFinite(resolved.facing??nextFacing))throw Error('Invalid resolved flat-ground transform.');
      nextPosition=Array.from(resolved.position);nextFacing=resolved.facing??nextFacing;
      nextVelocity=nextPosition.map((v,i)=>(v-oldPosition[i])/DT);
      if(Math.hypot(...nextVelocity)>this.config.max_speed_mps+1e-5)throw Error('Collision resolution exceeds bounded movement; use reset for teleports.');
      // Forecast actual movement when a host resolves collisions. No contact or
      // physics prediction is implied by this constant-velocity preview.
      for(let j=0;j<3;j++){const n=[12,24,36][j];points[j]=oldPosition.map((v,i)=>v+nextVelocity[i]*n*DT);velocities[j]=nextVelocity.slice();angles[j]=nextFacing;}
    }
    const speed=Math.hypot(...(resolveMovement?nextVelocity:this.velocity));
    const activity=clamp((speed-.06)/.19,0,1),frequency=clamp(this.config.cadence[0]+this.config.cadence[1]*speed,.4,1.8);
    const nextPhase=(this.phase+2*Math.PI*frequency*activity*DT)%(2*Math.PI);
    const input=new Float32Array(25),firstVelocity=yaw(velocities[0],-oldFacing),firstAngle=angles[0]-oldFacing;
    input.set([firstVelocity[0],firstVelocity[2],Math.sin(firstAngle),Math.cos(firstAngle)]);
    for(let j=0;j<3;j++) {
      const p=yaw(points[j].map((v,i)=>v-oldPosition[i]),-oldFacing),v=yaw(velocities[j],-oldFacing),a=angles[j]-oldFacing;
      input.set([p[0],p[2],v[0],v[2],Math.sin(a),Math.cos(a)],4+j*6);
    }
    input.set([Math.sin(nextPhase)*activity,Math.cos(nextPhase)*activity,frequency*activity],22);
    const output=this.network.forward(input),localQuaternions=[];let articulated=0;
    for(const end of this.rig.end_sites) {
      if(end){localQuaternions.push([0,0,0,1]);continue;}
      let columns=rotation(output,articulated*6);
      if(articulated===0)columns=columns.map(v=>yaw(v,nextFacing));
      localQuaternions.push(quaternion(columns));articulated++;
    }
    const offset=yaw(Array.from(output.subarray(138,141)),nextFacing),rootPosition=nextPosition.map((v,i)=>v+offset[i]);
    if(!rootPosition.every(Number.isFinite)||rootPosition[1]<.25||rootPosition[1]>1.8)throw Error('Predicted pelvis outside valid range.');
    this.position=nextPosition;this.velocity=nextVelocity;this.facing=nextFacing;this.phase=nextPhase;this.frame++;
    return {rootPosition,localQuaternions,contacts:Array.from(output.subarray(141),v=>1/(1+Math.exp(-v))),
      simulationPosition:nextPosition.slice(),facing:nextFacing,phase:nextPhase,frame:this.frame,
      speed:Math.hypot(...nextVelocity),inputClamped:clamped};
  }
}