#!/usr/bin/env node
/**
 * WiFlow Supervised Pose Training Pipeline (ADR-079)
 *
 * Trains WiFlow pose estimation on paired CSI + camera keypoint data.
 * Extends the ruvllm training infrastructure with a simplified TCN architecture
 * and three-phase curriculum: contrastive pretraining, supervised keypoint
 * regression, and refinement with bone/temporal constraints.
 *
 * Input format (paired JSONL):
 *   {"csi": [[...128 or 8 floats...], ...20 frames], "keypoints": [[x,y],...17], "conf": [c0..c16], "timestamp": ...}
 *
 * Architecture:
 *   TCN (4 dilated causal conv blocks, k=7, dilation 1,2,4,8)
 *     input_dim -> 256 -> 192 -> 128
 *   Flatten [128*20] -> Linear 2560 -> 2048 -> Linear 2048 -> 34
 *   Reshape to [17, 2] keypoints in [0, 1]
 *
 * Phases:
 *   1. Contrastive (50 epochs) — representation learning on CSI windows
 *   2. Supervised (200 epochs) — confidence-weighted SmoothL1 on keypoints
 *      with curriculum: conf>0.9 -> conf>0.7 -> conf>0.5 -> all + augmentation
 *   3. Refinement (50 epochs) — combined loss with bone + temporal constraints
 *
 * Usage:
 *   node scripts/train-wiflow-supervised.js --data data/paired-csi-keypoints.jsonl
 *   node scripts/train-wiflow-supervised.js --data data/paired.jsonl --skip-contrastive --epochs 200
 *   node scripts/train-wiflow-supervised.js --data data/paired.jsonl --output models/wiflow-sup-v2
 *
 * ADR: docs/adr/ADR-079
 */

'use strict';

const fs = require('fs');
const path = require('path');
const { parseArgs } = require('util');

// ---------------------------------------------------------------------------
// Resolve ruvllm from vendor tree
// ---------------------------------------------------------------------------
const RUVLLM_PATH = path.resolve(__dirname, '..', 'vendor', 'ruvector', 'npm', 'packages', 'ruvllm', 'src');

const {
  ContrastiveTrainer,
  cosineSimilarity,
  infoNCELoss,
  computeGradient,
} = require(path.join(RUVLLM_PATH, 'contrastive.js'));

const {
  TrainingPipeline,
} = require(path.join(RUVLLM_PATH, 'training.js'));

const {
  EwcManager,
} = require(path.join(RUVLLM_PATH, 'sona.js'));

const {
  SafeTensorsWriter,
  ModelExporter,
} = require(path.join(RUVLLM_PATH, 'export.js'));

// ---------------------------------------------------------------------------
// CLI argument parsing
// ---------------------------------------------------------------------------
const { values: args } = parseArgs({
  options: {
    data:              { type: 'string',  short: 'd' },
    output:            { type: 'string',  short: 'o', default: 'models/wiflow-supervised' },
    epochs:            { type: 'string',  short: 'e', default: '300' },
    'batch-size':      { type: 'string',              default: '32' },
    lr:                { type: 'string',              default: '0.0001' },
    'skip-contrastive': { type: 'boolean',            default: false },
    'eval-split':      { type: 'string',              default: '0.2' },
    verbose:           { type: 'boolean', short: 'v', default: false },
  },
  strict: true,
});

if (!args.data) {
  console.error('Usage: node scripts/train-wiflow-supervised.js --data <paired-jsonl> [options]');
  console.error('');
  console.error('Options:');
  console.error('  --data <file>          Paired CSI+keypoint JSONL (required)');
  console.error('  --output <dir>         Output directory (default: models/wiflow-supervised)');
  console.error('  --epochs <n>           Total epochs across all phases (default: 300)');
  console.error('  --batch-size <n>       Batch size (default: 32)');
  console.error('  --lr <float>           Learning rate (default: 0.0001)');
  console.error('  --skip-contrastive     Skip phase 1 contrastive pretraining');
  console.error('  --eval-split <float>   Held-out eval fraction (default: 0.2)');
  console.error('  --verbose              Print detailed progress');
  process.exit(1);
}

const CONFIG = {
  dataPath:         args.data,
  outputDir:        args.output,
  totalEpochs:      parseInt(args.epochs, 10),
  batchSize:        parseInt(args['batch-size'], 10),
  lr:               parseFloat(args.lr),
  skipContrastive:  args['skip-contrastive'],
  evalSplit:        parseFloat(args['eval-split']),
  verbose:          args.verbose,

  // Phase epoch allocation (scaled to totalEpochs)
  contrastiveRatio: 50 / 300,
  supervisedRatio:  200 / 300,
  refinementRatio:  50 / 300,

  // Curriculum confidence thresholds (O1)
  curriculumStages: [0.9, 0.7, 0.5, 0.0],

  // Architecture
  timeSteps:        20,
  numKeypoints:     17,

  // SGD momentum
  momentum:         0.9,

  // Refinement loss weights
  boneWeight:       0.3,
  temporalWeight:   0.1,
};

// Compute phase epochs
const totalForPhases = CONFIG.skipContrastive
  ? CONFIG.totalEpochs
  : CONFIG.totalEpochs;
const contrastiveEpochs = CONFIG.skipContrastive ? 0 : Math.round(totalForPhases * CONFIG.contrastiveRatio);
const supervisedEpochs  = Math.round(totalForPhases * CONFIG.supervisedRatio);
const refinementEpochs  = totalForPhases - contrastiveEpochs - supervisedEpochs;

// ---------------------------------------------------------------------------
// Deterministic PRNG (xorshift32)
// ---------------------------------------------------------------------------

function createRng(seed) {
  let s = seed | 0 || 42;
  return () => {
    s ^= s << 13;
    s ^= s >> 17;
    s ^= s << 5;
    return (s >>> 0) / 4294967296;
  };
}

function gaussianRng(rng) {
  return () => {
    const u1 = rng() || 1e-10;
    const u2 = rng();
    return Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
  };
}

// ---------------------------------------------------------------------------
// Tensor utilities
// ---------------------------------------------------------------------------

function initKaiming(fanIn, fanOut, rng) {
  const std = Math.sqrt(2.0 / fanIn);
  const gauss = gaussianRng(rng);
  const arr = new Float32Array(fanIn * fanOut);
  for (let i = 0; i < arr.length; i++) arr[i] = gauss() * std;
  return arr;
}

function initXavier(fanIn, fanOut, rng) {
  const std = Math.sqrt(2.0 / (fanIn + fanOut));
  const gauss = gaussianRng(rng);
  const arr = new Float32Array(fanIn * fanOut);
  for (let i = 0; i < arr.length; i++) arr[i] = gauss() * std;
  return arr;
}

function relu(arr) {
  for (let i = 0; i < arr.length; i++) {
    if (arr[i] < 0) arr[i] = 0;
  }
  return arr;
}

function sigmoid(x) {
  return 1.0 / (1.0 + Math.exp(-x));
}

// ---------------------------------------------------------------------------
// SmoothL1 loss and gradient
// ---------------------------------------------------------------------------

function smoothL1(predicted, target, beta) {
  beta = beta || 0.05;
  let loss = 0;
  const n = Math.min(predicted.length, target.length);
  for (let i = 0; i < n; i++) {
    const diff = Math.abs(predicted[i] - target[i]);
    if (diff < beta) {
      loss += 0.5 * diff * diff / beta;
    } else {
      loss += diff - 0.5 * beta;
    }
  }
  return loss / n;
}

function smoothL1Grad(predicted, target, beta) {
  beta = beta || 0.05;
  const n = Math.min(predicted.length, target.length);
  const grad = new Float32Array(n);
  for (let i = 0; i < n; i++) {
    const diff = predicted[i] - target[i];
    const absDiff = Math.abs(diff);
    if (absDiff < beta) {
      grad[i] = diff / beta / n;
    } else {
      grad[i] = (diff > 0 ? 1 : -1) / n;
    }
  }
  return grad;
}

// ---------------------------------------------------------------------------
// COCO bone priors (ADR-079)
// ---------------------------------------------------------------------------

const BONE_CONNECTIONS = [
  [0, 1], [0, 2],         // nose -> eyes
  [1, 3], [2, 4],         // eyes -> ears
  [5, 7], [7, 9],         // left arm: shoulder-elbow, elbow-wrist
  [6, 8], [8, 10],        // right arm: shoulder-elbow, elbow-wrist
  [5, 11], [6, 12],       // torso: shoulder-hip
  [11, 13], [13, 15],     // left leg: hip-knee, knee-ankle
  [12, 14], [14, 16],     // right leg: hip-knee, knee-ankle
  [5, 6],                 // shoulder width
];

const BONE_LENGTH_PRIORS = [
  0.06, 0.06,   // nose-eye
  0.06, 0.06,   // eye-ear
  0.15, 0.13,   // left shoulder-elbow, elbow-wrist
  0.15, 0.13,   // right shoulder-elbow, elbow-wrist
  0.26, 0.26,   // shoulder-hip
  0.25, 0.25,   // left hip-knee, knee-ankle
  0.25, 0.25,   // right hip-knee, knee-ankle
  0.20,         // shoulder width
];

// ---------------------------------------------------------------------------
// Data loading — paired CSI + keypoint JSONL
// ---------------------------------------------------------------------------

/**
 * Load paired dataset from JSONL file.
 * Each line: { csi: [[...], ...], keypoints: [[x,y], ...17], conf: [...17], timestamp: ... }
 * csi shape: [subcarriers, timeSteps] or [features, timeSteps]
 */
function loadPairedData(filePath) {
  if (!fs.existsSync(filePath)) {
    console.error(`Data file not found: ${filePath}`);
    process.exit(1);
  }

  const content = fs.readFileSync(filePath, 'utf-8');
  const lines = content.split('\n').filter(l => l.trim());
  const samples = [];

  for (const line of lines) {
    try {
      const obj = JSON.parse(line);
      if (!obj.csi || !obj.keypoints) continue;

      const csi = obj.csi;           // 2D array [dim, T] or flat
      const kp  = obj.keypoints;     // [[x,y], ...] or flat [x,y,x,y,...]
      const conf = obj.conf || null;  // [c0, c1, ...c16] or null
      const ts  = obj.timestamp || 0;

      // Flatten keypoints to [34] = [x0, y0, x1, y1, ...]
      let kpFlat;
      if (Array.isArray(kp[0])) {
        kpFlat = new Float32Array(CONFIG.numKeypoints * 2);
        for (let i = 0; i < CONFIG.numKeypoints && i < kp.length; i++) {
          kpFlat[i * 2]     = kp[i][0];
          kpFlat[i * 2 + 1] = kp[i][1];
        }
      } else {
        kpFlat = new Float32Array(kp.slice(0, CONFIG.numKeypoints * 2));
      }

      // Confidence per keypoint
      let confArr;
      if (conf && conf.length >= CONFIG.numKeypoints) {
        confArr = new Float32Array(conf.slice(0, CONFIG.numKeypoints));
      } else {
        confArr = new Float32Array(CONFIG.numKeypoints).fill(1.0);
      }

      // Flatten CSI to Float32Array [dim * T]
      let csiFlat;
      let csiDim;
      if (Array.isArray(csi[0])) {
        csiDim = csi.length;
        const T = csi[0].length;
        csiFlat = new Float32Array(csiDim * T);
        for (let d = 0; d < csiDim; d++) {
          for (let t = 0; t < T; t++) {
            csiFlat[d * T + t] = csi[d][t] || 0;
          }
        }
      } else {
        // Assume flat 1D array, treat as [dim, 1] — shouldn't happen normally
        csiDim = csi.length;
        csiFlat = new Float32Array(csi);
      }

      samples.push({ csi: csiFlat, csiDim, keypoints: kpFlat, conf: confArr, timestamp: ts });
    } catch (_) {
      // Skip malformed lines
    }
  }

  return samples;
}

// ---------------------------------------------------------------------------
// Data augmentation (O2)
// ---------------------------------------------------------------------------

function augmentSample(sample, rng, T) {
  const dim = sample.csiDim;
  const augCsi = new Float32Array(sample.csi);

  // Time shift: roll ±2 frames
  const shift = Math.floor(rng() * 5) - 2; // -2 to +2
  if (shift !== 0) {
    const temp = new Float32Array(dim * T);
    for (let d = 0; d < dim; d++) {
      for (let t = 0; t < T; t++) {
        let srcT = t - shift;
        if (srcT < 0) srcT = 0;
        if (srcT >= T) srcT = T - 1;
        temp[d * T + t] = augCsi[d * T + srcT];
      }
    }
    augCsi.set(temp);
  }

  // Amplitude noise: gaussian sigma=0.02
  const gauss = gaussianRng(rng);
  for (let i = 0; i < augCsi.length; i++) {
    augCsi[i] += gauss() * 0.02;
  }

  // Subcarrier dropout: zero 10% randomly
  for (let d = 0; d < dim; d++) {
    if (rng() < 0.10) {
      for (let t = 0; t < T; t++) {
        augCsi[d * T + t] = 0;
      }
    }
  }

  return {
    csi: augCsi,
    csiDim: dim,
    keypoints: sample.keypoints,
    conf: sample.conf,
    timestamp: sample.timestamp,
  };
}

// ---------------------------------------------------------------------------
// Deterministic shuffle
// ---------------------------------------------------------------------------

function shuffleArray(arr, seed) {
  const result = [...arr];
  let s = seed;
  for (let i = result.length - 1; i > 0; i--) {
    s ^= s << 13; s ^= s >> 17; s ^= s << 5;
    const j = (s >>> 0) % (i + 1);
    [result[i], result[j]] = [result[j], result[i]];
  }
  return result;
}

// ---------------------------------------------------------------------------
// WiFlow Supervised Model — simplified TCN + linear decoder
// ---------------------------------------------------------------------------

/**
 * 1D causal dilated convolution layer.
 * Weight shape: [outCh, inCh, kernel] stored as flat Float32Array.
 * Input/output layout: [channels, T].
 */
class CausalConv1d {
  constructor(inCh, outCh, kernel, dilation, rng) {
    this.inCh = inCh;
    this.outCh = outCh;
    this.kernel = kernel;
    this.dilation = dilation || 1;

    // Kaiming init
    this.weight = initKaiming(inCh * kernel, outCh, rng);
    this.bias = new Float32Array(outCh);

    // Momentum buffers for SGD
    this.weightMom = new Float32Array(this.weight.length);
    this.biasMom   = new Float32Array(outCh);
  }

  numParams() {
    return this.weight.length + this.bias.length;
  }

  /**
   * Forward: [inCh, T] -> [outCh, T] with causal (left) padding.
   */
  forward(input, T) {
    const effectiveK = this.kernel + (this.kernel - 1) * (this.dilation - 1);
    const padLeft = effectiveK - 1;
    const T_padded = T + padLeft;

    // Pad input
    const padded = new Float32Array(this.inCh * T_padded);
    for (let c = 0; c < this.inCh; c++) {
      for (let t = 0; t < T; t++) {
        padded[c * T_padded + (t + padLeft)] = input[c * T + t];
      }
    }

    // Convolve
    const output = new Float32Array(this.outCh * T);
    for (let oc = 0; oc < this.outCh; oc++) {
      for (let t = 0; t < T; t++) {
        let sum = this.bias[oc];
        for (let ic = 0; ic < this.inCh; ic++) {
          for (let k = 0; k < this.kernel; k++) {
            const tIdx = t + padLeft - k * this.dilation;
            if (tIdx >= 0 && tIdx < T_padded) {
              const wIdx = oc * (this.inCh * this.kernel) + ic * this.kernel + k;
              sum += this.weight[wIdx] * padded[ic * T_padded + tIdx];
            }
          }
        }
        output[oc * T + t] = sum;
      }
    }
    return output;
  }
}

/**
 * Batch normalization for 1D temporal data [channels, T].
 * Uses running mean/var for inference; batch stats for training.
 */
class BatchNorm1d {
  constructor(channels) {
    this.channels = channels;
    this.gamma = new Float32Array(channels).fill(1.0);
    this.beta  = new Float32Array(channels);
    this.runMean = new Float32Array(channels);
    this.runVar  = new Float32Array(channels).fill(1.0);
    this.momentum = 0.1;
    this.eps = 1e-5;

    // Momentum buffers
    this.gammaMom = new Float32Array(channels);
    this.betaMom  = new Float32Array(channels);
  }

  numParams() {
    return this.channels * 2;
  }

  /**
   * Forward: [channels, T] -> [channels, T], updates running stats.
   */
  forward(input, T) {
    const output = new Float32Array(input.length);
    for (let c = 0; c < this.channels; c++) {
      // Compute channel mean and var over T
      let mean = 0, varAcc = 0;
      for (let t = 0; t < T; t++) mean += input[c * T + t];
      mean /= T;
      for (let t = 0; t < T; t++) varAcc += (input[c * T + t] - mean) ** 2;
      varAcc /= T;

      // Update running stats
      this.runMean[c] = (1 - this.momentum) * this.runMean[c] + this.momentum * mean;
      this.runVar[c]  = (1 - this.momentum) * this.runVar[c]  + this.momentum * varAcc;

      // Normalize
      const invStd = 1.0 / Math.sqrt(varAcc + this.eps);
      for (let t = 0; t < T; t++) {
        output[c * T + t] = this.gamma[c] * (input[c * T + t] - mean) * invStd + this.beta[c];
      }
    }
    return output;
  }
}

/**
 * TCN block: Conv1d (causal, dilated) -> BN -> ReLU -> Conv1d -> BN + residual -> ReLU
 */
class TCNBlock {
  constructor(inCh, outCh, kernel, dilation, rng) {
    this.conv1 = new CausalConv1d(inCh, outCh, kernel, dilation, rng);
    this.bn1   = new BatchNorm1d(outCh);
    this.conv2 = new CausalConv1d(outCh, outCh, kernel, dilation, rng);
    this.bn2   = new BatchNorm1d(outCh);

    // Residual projection if dimensions differ
    this.hasResProj = (inCh !== outCh);
    if (this.hasResProj) {
      this.resConv = new CausalConv1d(inCh, outCh, 1, 1, rng);
    }
  }

  numParams() {
    let p = this.conv1.numParams() + this.bn1.numParams() +
            this.conv2.numParams() + this.bn2.numParams();
    if (this.hasResProj) p += this.resConv.numParams();
    return p;
  }

  forward(input, T) {
    // Path 1: conv -> bn -> relu -> conv -> bn
    let x = this.conv1.forward(input, T);
    x = this.bn1.forward(x, T);
    relu(x);
    x = this.conv2.forward(x, T);
    x = this.bn2.forward(x, T);

    // Residual
    const res = this.hasResProj ? this.resConv.forward(input, T) : input;
    for (let i = 0; i < x.length; i++) x[i] += res[i];
    relu(x);
    return x;
  }
}

/**
 * Linear layer: [inDim] -> [outDim]
 */
class Linear {
  constructor(inDim, outDim, rng) {
    this.inDim  = inDim;
    this.outDim = outDim;
    this.weight = initXavier(inDim, outDim, rng);
    this.bias   = new Float32Array(outDim);

    // Momentum buffers
    this.weightMom = new Float32Array(this.weight.length);
    this.biasMom   = new Float32Array(outDim);
  }

  numParams() {
    return this.weight.length + this.bias.length;
  }

  forward(input) {
    const output = new Float32Array(this.outDim);
    for (let j = 0; j < this.outDim; j++) {
      let sum = this.bias[j];
      for (let i = 0; i < this.inDim; i++) {
        sum += input[i] * this.weight[i * this.outDim + j];
      }
      output[j] = sum;
    }
    return output;
  }
}

/**
 * WiFlow Supervised Model.
 *
 * TCN Stage: 4 dilated causal conv blocks (dilation 1,2,4,8), kernel 7
 *   input_dim -> 256 -> 192 -> 128
 * Flatten + Linear: [128 * 20] -> 2048 -> [17 * 2]
 * Sigmoid to [0, 1]
 */
class WiFlowSupervisedModel {
  constructor(inputDim, timeSteps, numKeypoints, seed) {
    this.inputDim    = inputDim;
    this.timeSteps   = timeSteps;
    this.numKeypoints = numKeypoints || 17;
    this.outDim      = this.numKeypoints * 2;

    const rng = createRng(seed || 42);

    // TCN blocks: inputDim -> 256 -> 256 -> 192 -> 128
    this.tcn1 = new TCNBlock(inputDim, 256, 7, 1, rng);
    this.tcn2 = new TCNBlock(256, 256, 7, 2, rng);
    this.tcn3 = new TCNBlock(256, 192, 7, 4, rng);
    this.tcn4 = new TCNBlock(192, 128, 7, 8, rng);

    // Flatten: 128 * timeSteps -> linear -> 34
    const flatDim = 128 * timeSteps;
    this.fc1 = new Linear(flatDim, 2048, rng);
    this.fc2 = new Linear(2048, this.outDim, rng);

    this._totalParams = null;
  }

  totalParams() {
    if (this._totalParams === null) {
      this._totalParams = this.tcn1.numParams() + this.tcn2.numParams() +
        this.tcn3.numParams() + this.tcn4.numParams() +
        this.fc1.numParams() + this.fc2.numParams();
    }
    return this._totalParams;
  }

  /**
   * Forward pass.
   * @param {Float32Array} csi - [inputDim * timeSteps] flat
   * @returns {Float32Array} keypoints [numKeypoints * 2] in [0, 1]
   */
  forward(csi) {
    const T = this.timeSteps;

    // TCN stages
    let x = this.tcn1.forward(csi, T);
    x = this.tcn2.forward(x, T);
    x = this.tcn3.forward(x, T);
    x = this.tcn4.forward(x, T);

    // Flatten: [128, T] -> [128*T]
    // x is already flat as [128 * T]

    // FC layers with ReLU
    let h = this.fc1.forward(x);
    relu(h);
    let out = this.fc2.forward(h);

    // Sigmoid to [0, 1]
    for (let i = 0; i < out.length; i++) {
      out[i] = sigmoid(out[i]);
    }

    return out;
  }

  /**
   * Encode CSI to embedding (for contrastive phase).
   * Returns the fc1 hidden layer (2048-dim).
   */
  encode(csi) {
    const T = this.timeSteps;
    let x = this.tcn1.forward(csi, T);
    x = this.tcn2.forward(x, T);
    x = this.tcn3.forward(x, T);
    x = this.tcn4.forward(x, T);

    let h = this.fc1.forward(x);
    relu(h);

    // L2 normalize for contrastive
    let norm = 0;
    for (let i = 0; i < h.length; i++) norm += h[i] * h[i];
    norm = Math.sqrt(norm) || 1;
    for (let i = 0; i < h.length; i++) h[i] /= norm;

    return h;
  }

  /**
   * Collect all weight arrays for gradient updates.
   * Returns array of { weight, mom, name } objects.
   */
  collectParams() {
    const params = [];
    const addConv = (conv, prefix) => {
      params.push({ weight: conv.weight, mom: conv.weightMom, name: `${prefix}.weight` });
      params.push({ weight: conv.bias,   mom: conv.biasMom,   name: `${prefix}.bias` });
    };
    const addBN = (bn, prefix) => {
      params.push({ weight: bn.gamma, mom: bn.gammaMom, name: `${prefix}.gamma` });
      params.push({ weight: bn.beta,  mom: bn.betaMom,  name: `${prefix}.beta` });
    };
    const addTCN = (tcn, prefix) => {
      addConv(tcn.conv1, `${prefix}.conv1`);
      addBN(tcn.bn1, `${prefix}.bn1`);
      addConv(tcn.conv2, `${prefix}.conv2`);
      addBN(tcn.bn2, `${prefix}.bn2`);
      if (tcn.hasResProj) addConv(tcn.resConv, `${prefix}.res`);
    };
    const addLinear = (linear, prefix) => {
      params.push({ weight: linear.weight, mom: linear.weightMom, name: `${prefix}.weight` });
      params.push({ weight: linear.bias,   mom: linear.biasMom,   name: `${prefix}.bias` });
    };

    addTCN(this.tcn1, 'tcn1');
    addTCN(this.tcn2, 'tcn2');
    addTCN(this.tcn3, 'tcn3');
    addTCN(this.tcn4, 'tcn4');
    addLinear(this.fc1, 'fc1');
    addLinear(this.fc2, 'fc2');

    return params;
  }

  /**
   * Get all weights as a flat Float32Array (for export).
   */
  getAllWeights() {
    const params = this.collectParams();
    let totalLen = 0;
    for (const p of params) totalLen += p.weight.length;
    const flat = new Float32Array(totalLen);
    let offset = 0;
    for (const p of params) {
      flat.set(p.weight, offset);
      offset += p.weight.length;
    }
    return flat;
  }
}

// ---------------------------------------------------------------------------
// SGD with momentum + cosine LR decay
// ---------------------------------------------------------------------------

/**
 * Numerical gradient estimation using finite differences.
 * Computes gradient of lossFn w.r.t. each parameter in paramObj.weight.
 */
function computeNumericalGrad(model, sample, lossFn, paramObj, eps) {
  eps = eps || 1e-4;
  const w = paramObj.weight;
  const grad = new Float32Array(w.length);

  for (let i = 0; i < w.length; i++) {
    const orig = w[i];

    w[i] = orig + eps;
    const lossPlus = lossFn(model, sample);

    w[i] = orig - eps;
    const lossMinus = lossFn(model, sample);

    w[i] = orig;
    grad[i] = (lossPlus - lossMinus) / (2 * eps);
  }

  return grad;
}

/**
 * Apply SGD with momentum to a single parameter.
 */
function sgdStep(paramObj, grad, lr, momentum) {
  const w   = paramObj.weight;
  const mom = paramObj.mom;
  for (let i = 0; i < w.length; i++) {
    mom[i] = momentum * mom[i] + grad[i];
    w[i] -= lr * mom[i];
  }
}

/**
 * Cosine annealing learning rate.
 */
function cosineDecayLR(baseLR, epoch, totalEpochs) {
  return baseLR * 0.5 * (1 + Math.cos(Math.PI * epoch / totalEpochs));
}

// ---------------------------------------------------------------------------
// Loss functions
// ---------------------------------------------------------------------------

/**
 * Confidence-weighted SmoothL1 loss for keypoints.
 * L = (1/N) * sum(conf_i * smoothL1(pred_i, gt_i, beta=0.05))
 */
function supervisedLoss(predicted, target, conf, beta) {
  beta = beta || 0.05;
  const nKp = conf.length;
  let loss = 0;
  let weightSum = 0;

  for (let k = 0; k < nKp; k++) {
    const px = predicted[k * 2], py = predicted[k * 2 + 1];
    const tx = target[k * 2],    ty = target[k * 2 + 1];

    const diffX = Math.abs(px - tx);
    const diffY = Math.abs(py - ty);

    let lx = diffX < beta ? 0.5 * diffX * diffX / beta : diffX - 0.5 * beta;
    let ly = diffY < beta ? 0.5 * diffY * diffY / beta : diffY - 0.5 * beta;

    loss += conf[k] * (lx + ly);
    weightSum += conf[k];
  }

  return weightSum > 0 ? loss / weightSum : 0;
}

/**
 * Bone length constraint loss.
 */
function boneLoss(predicted) {
  let loss = 0;
  for (let b = 0; b < BONE_CONNECTIONS.length; b++) {
    const [i, j] = BONE_CONNECTIONS[b];
    const prior = BONE_LENGTH_PRIORS[b];
    const dx = predicted[i * 2] - predicted[j * 2];
    const dy = predicted[i * 2 + 1] - predicted[j * 2 + 1];
    const boneLen = Math.sqrt(dx * dx + dy * dy);
    const deviation = boneLen - prior;
    loss += deviation * deviation;
  }
  return loss / BONE_CONNECTIONS.length;
}

/**
 * Temporal consistency loss between consecutive predictions.
 */
function temporalLoss(predCurrent, predPrev) {
  if (!predPrev) return 0;
  return smoothL1(predCurrent, predPrev, 0.05);
}

// ---------------------------------------------------------------------------
// Evaluation: PCK@threshold
// ---------------------------------------------------------------------------

function pck(predicted, target, threshold) {
  threshold = threshold || 0.2;
  let correct = 0;
  const nKp = Math.min(predicted.length, target.length) / 2;
  for (let k = 0; k < nKp; k++) {
    const dx = predicted[k * 2] - target[k * 2];
    const dy = predicted[k * 2 + 1] - target[k * 2 + 1];
    if (Math.sqrt(dx * dx + dy * dy) < threshold) correct++;
  }
  return correct / nKp;
}

/**
 * Evaluate model on held-out set, return average loss and PCK@20.
 */
function evaluate(model, evalSet) {
  let totalLoss = 0;
  let totalPck = 0;

  for (const sample of evalSet) {
    const pred = model.forward(sample.csi);
    totalLoss += supervisedLoss(pred, sample.keypoints, sample.conf);
    totalPck  += pck(pred, sample.keypoints, 0.2);
  }

  return {
    loss: evalSet.length > 0 ? totalLoss / evalSet.length : 0,
    pck20: evalSet.length > 0 ? totalPck / evalSet.length : 0,
  };
}

// ---------------------------------------------------------------------------
// Stochastic gradient estimation for a mini-batch
// ---------------------------------------------------------------------------

/**
 * Estimate gradient via forward-mode perturbation for a mini-batch.
 * This uses simultaneous perturbation (SPSA-like) which scales O(1) per
 * parameter rather than O(n) for naive numerical differentiation.
 */
function estimateBatchGrad(model, batch, lossFn, paramObj, rng) {
  const eps = 1e-4;
  const w = paramObj.weight;
  const n = w.length;
  const grad = new Float32Array(n);

  // Use SPSA: perturb all weights simultaneously with random direction
  const delta = new Float32Array(n);
  for (let i = 0; i < n; i++) {
    delta[i] = rng() < 0.5 ? 1 : -1;
  }

  // Compute loss at w + eps*delta
  for (let i = 0; i < n; i++) w[i] += eps * delta[i];
  let lossPlus = 0;
  for (const sample of batch) lossPlus += lossFn(model, sample);
  lossPlus /= batch.length;

  // Compute loss at w - eps*delta
  for (let i = 0; i < n; i++) w[i] -= 2 * eps * delta[i];
  let lossMinus = 0;
  for (const sample of batch) lossMinus += lossFn(model, sample);
  lossMinus /= batch.length;

  // Restore weights
  for (let i = 0; i < n; i++) w[i] += eps * delta[i];

  // SPSA gradient estimate
  const scale = (lossPlus - lossMinus) / (2 * eps);
  for (let i = 0; i < n; i++) {
    grad[i] = scale / delta[i];
  }

  return grad;
}

// ---------------------------------------------------------------------------
// Main training pipeline
// ---------------------------------------------------------------------------

async function main() {
  const startTime = Date.now();
  console.log('=== WiFlow Supervised Pose Training Pipeline (ADR-079) ===');
  console.log(`Config: totalEpochs=${CONFIG.totalEpochs} batch=${CONFIG.batchSize} lr=${CONFIG.lr}`);
  console.log(`        phases: contrastive=${contrastiveEpochs} supervised=${supervisedEpochs} refinement=${refinementEpochs}`);
  console.log(`        momentum=${CONFIG.momentum} evalSplit=${CONFIG.evalSplit}`);
  console.log('');

  // -----------------------------------------------------------------------
  // Step 1: Load paired data
  // -----------------------------------------------------------------------
  console.log('[1/6] Loading paired CSI+keypoint data...');
  const allSamples = loadPairedData(CONFIG.dataPath);
  if (allSamples.length === 0) {
    console.error('No valid paired samples found in data file.');
    process.exit(1);
  }

  // Auto-detect input dimension
  const inputDim = allSamples[0].csiDim;
  const T = CONFIG.timeSteps;
  console.log(`  Loaded ${allSamples.length} paired samples`);
  console.log(`  Auto-detected input dim: ${inputDim} (${inputDim === 128 ? 'full CSI subcarriers' : inputDim + '-dim feature vectors'})`);
  console.log(`  Time steps: ${T}`);

  // Train/eval split
  const shuffled = shuffleArray(allSamples, 42);
  const splitIdx = Math.floor(shuffled.length * (1 - CONFIG.evalSplit));
  const trainSet = shuffled.slice(0, splitIdx);
  const evalSet  = shuffled.slice(splitIdx);
  console.log(`  Train: ${trainSet.length}  Eval: ${evalSet.length}`);
  console.log('');

  // -----------------------------------------------------------------------
  // Step 2: Initialize model
  // -----------------------------------------------------------------------
  console.log('[2/6] Initializing WiFlow supervised model...');
  const model = new WiFlowSupervisedModel(inputDim, T, CONFIG.numKeypoints, 42);
  console.log(`  Parameters: ${model.totalParams().toLocaleString()}`);
  console.log(`  Architecture: TCN(${inputDim}->256->256->192->128, k=7, d=[1,2,4,8]) -> FC(${128 * T}->2048->34)`);
  console.log('');

  const trainingLog = {
    config: { ...CONFIG, inputDim, contrastiveEpochs, supervisedEpochs, refinementEpochs },
    phases: [],
  };

  const allParams = model.collectParams();
  const rng = createRng(123);
  let globalEpoch = 0;

  // -----------------------------------------------------------------------
  // Phase 1: Contrastive pretraining
  // -----------------------------------------------------------------------
  if (!CONFIG.skipContrastive && contrastiveEpochs > 0) {
    console.log(`[3/6] Phase 1: Contrastive pretraining (${contrastiveEpochs} epochs)...`);

    const contrastiveLog = { phase: 'contrastive', epochs: [] };
    const trainer = new ContrastiveTrainer({
      margin: 0.3,
      temperature: 0.07,
    });

    for (let epoch = 0; epoch < contrastiveEpochs; epoch++) {
      const lr = cosineDecayLR(CONFIG.lr * 10, epoch, contrastiveEpochs); // Higher LR for contrastive
      const shuffledTrain = shuffleArray(trainSet, epoch * 7 + 1);

      let epochLoss = 0;
      let nBatches = 0;

      for (let b = 0; b < shuffledTrain.length - 2; b += CONFIG.batchSize) {
        const batchEnd = Math.min(b + CONFIG.batchSize, shuffledTrain.length - 2);
        let batchLoss = 0;
        let nTriplets = 0;

        // Create temporal triplets: anchor=frame[i], positive=frame[i+1], negative=frame[j] (far)
        for (let i = b; i < batchEnd; i++) {
          const anchorEmb   = Array.from(model.encode(shuffledTrain[i].csi));
          const positiveEmb = Array.from(model.encode(shuffledTrain[i + 1].csi));
          // Negative: pick a distant sample
          const negIdx = (i + Math.floor(shuffledTrain.length / 2)) % shuffledTrain.length;
          const negativeEmb = Array.from(model.encode(shuffledTrain[negIdx].csi));

          trainer.addTriplet(
            `anchor-${i}`, anchorEmb,
            `pos-${i}`, positiveEmb,
            `neg-${i}`, negativeEmb,
          );

          const sim_pos = cosineSimilarity(anchorEmb, positiveEmb);
          const sim_neg = cosineSimilarity(anchorEmb, negativeEmb);
          batchLoss += Math.max(0, 0.3 - sim_pos + sim_neg);
          nTriplets++;
        }

        if (nTriplets > 0) batchLoss /= nTriplets;

        // SPSA gradient update on all params
        for (const p of allParams) {
          const lossFn = (m, s) => {
            const emb = m.encode(s.csi);
            // Simple self-consistency loss
            let norm = 0;
            for (let i = 0; i < emb.length; i++) norm += emb[i] * emb[i];
            return 1.0 - norm; // push toward unit norm
          };

          const batch = shuffledTrain.slice(b, batchEnd);
          const grad = estimateBatchGrad(model, batch, lossFn, p, rng);
          sgdStep(p, grad, lr, CONFIG.momentum);
        }

        epochLoss += batchLoss;
        nBatches++;
      }

      epochLoss = nBatches > 0 ? epochLoss / nBatches : 0;
      const evalResult = evaluate(model, evalSet);

      contrastiveLog.epochs.push({
        epoch: globalEpoch,
        loss: epochLoss,
        evalLoss: evalResult.loss,
        pck20: evalResult.pck20,
        lr,
      });

      if ((epoch + 1) % 10 === 0 || epoch === 0) {
        console.log(`  [contrastive] epoch ${epoch + 1}/${contrastiveEpochs}  loss=${epochLoss.toFixed(6)}  eval_loss=${evalResult.loss.toFixed(6)}  PCK@20=${(evalResult.pck20 * 100).toFixed(1)}%  lr=${lr.toExponential(2)}`);
      }
      globalEpoch++;
    }

    trainingLog.phases.push(contrastiveLog);
    console.log('');
  } else {
    console.log('[3/6] Phase 1: Contrastive pretraining SKIPPED');
    console.log('');
  }

  // -----------------------------------------------------------------------
  // Phase 2: Supervised training with curriculum (O1)
  // -----------------------------------------------------------------------
  console.log(`[4/6] Phase 2: Supervised keypoint regression (${supervisedEpochs} epochs, 4-stage curriculum)...`);

  const supervisedLog = { phase: 'supervised', epochs: [] };
  const epochsPerStage = Math.floor(supervisedEpochs / CONFIG.curriculumStages.length);

  for (let epoch = 0; epoch < supervisedEpochs; epoch++) {
    // Determine curriculum stage
    const stageIdx = Math.min(
      Math.floor(epoch / epochsPerStage),
      CONFIG.curriculumStages.length - 1
    );
    const confThreshold = CONFIG.curriculumStages[stageIdx];
    const useAugmentation = (stageIdx === CONFIG.curriculumStages.length - 1);

    const lr = cosineDecayLR(CONFIG.lr, epoch, supervisedEpochs);

    // Filter training samples by confidence threshold
    let trainSubset;
    if (confThreshold > 0) {
      trainSubset = trainSet.filter(s => {
        let meanConf = 0;
        for (let i = 0; i < s.conf.length; i++) meanConf += s.conf[i];
        meanConf /= s.conf.length;
        return meanConf >= confThreshold;
      });
    } else {
      trainSubset = trainSet;
    }

    // Apply augmentation in final stage
    if (useAugmentation) {
      const augmented = [];
      for (const s of trainSubset) {
        augmented.push(s);
        augmented.push(augmentSample(s, createRng(epoch * 1000 + augmented.length), T));
      }
      trainSubset = augmented;
    }

    if (trainSubset.length === 0) {
      // Skip if no samples pass threshold
      globalEpoch++;
      continue;
    }

    const shuffledTrain = shuffleArray(trainSubset, epoch * 13 + 3);

    let epochLoss = 0;
    let nBatches = 0;

    for (let b = 0; b < shuffledTrain.length; b += CONFIG.batchSize) {
      const batchEnd = Math.min(b + CONFIG.batchSize, shuffledTrain.length);
      const batch = shuffledTrain.slice(b, batchEnd);

      // Compute batch loss
      const lossFn = (m, s) => {
        const pred = m.forward(s.csi);
        return supervisedLoss(pred, s.keypoints, s.conf);
      };

      let batchLoss = 0;
      for (const s of batch) batchLoss += lossFn(model, s);
      batchLoss /= batch.length;

      // SPSA gradient update
      for (const p of allParams) {
        const grad = estimateBatchGrad(model, batch, lossFn, p, rng);
        sgdStep(p, grad, lr, CONFIG.momentum);
      }

      epochLoss += batchLoss;
      nBatches++;
    }

    epochLoss = nBatches > 0 ? epochLoss / nBatches : 0;
    const evalResult = evaluate(model, evalSet);

    supervisedLog.epochs.push({
      epoch: globalEpoch,
      stage: stageIdx + 1,
      confThreshold,
      loss: epochLoss,
      evalLoss: evalResult.loss,
      pck20: evalResult.pck20,
      lr,
      trainSamples: trainSubset.length,
    });

    if ((epoch + 1) % 10 === 0 || epoch === 0) {
      console.log(`  [supervised] epoch ${epoch + 1}/${supervisedEpochs}  stage=${stageIdx + 1}/4 (conf>${confThreshold.toFixed(1)})  loss=${epochLoss.toFixed(6)}  eval_loss=${evalResult.loss.toFixed(6)}  PCK@20=${(evalResult.pck20 * 100).toFixed(1)}%  lr=${lr.toExponential(2)}  samples=${trainSubset.length}`);
    }
    globalEpoch++;
  }

  trainingLog.phases.push(supervisedLog);
  console.log('');

  // -----------------------------------------------------------------------
  // Phase 3: Refinement with bone + temporal constraints
  // -----------------------------------------------------------------------
  console.log(`[5/6] Phase 3: Refinement with bone + temporal constraints (${refinementEpochs} epochs)...`);

  const refinementLog = { phase: 'refinement', epochs: [] };

  for (let epoch = 0; epoch < refinementEpochs; epoch++) {
    const lr = cosineDecayLR(CONFIG.lr * 0.5, epoch, refinementEpochs); // Lower LR
    const shuffledTrain = shuffleArray(trainSet, epoch * 17 + 7);

    // Apply augmentation
    const augmented = [];
    for (const s of shuffledTrain) {
      augmented.push(s);
      augmented.push(augmentSample(s, createRng(epoch * 2000 + augmented.length), T));
    }

    let epochLoss = 0;
    let epochBone = 0;
    let epochTemporal = 0;
    let nBatches = 0;

    for (let b = 0; b < augmented.length; b += CONFIG.batchSize) {
      const batchEnd = Math.min(b + CONFIG.batchSize, augmented.length);
      const batch = augmented.slice(b, batchEnd);

      // Combined loss function
      const lossFn = (m, s, prevPred) => {
        const pred = m.forward(s.csi);
        const lSup  = supervisedLoss(pred, s.keypoints, s.conf);
        const lBone = boneLoss(pred);
        const lTemp = prevPred ? temporalLoss(pred, prevPred) : 0;
        return lSup + CONFIG.boneWeight * lBone + CONFIG.temporalWeight * lTemp;
      };

      // Compute batch loss with temporal tracking
      let batchLoss = 0;
      let batchBone = 0;
      let batchTemporal = 0;
      let prevPred = null;
      for (const s of batch) {
        const pred = model.forward(s.csi);
        const lSup  = supervisedLoss(pred, s.keypoints, s.conf);
        const lBone = boneLoss(pred);
        const lTemp = prevPred ? temporalLoss(pred, prevPred) : 0;
        batchLoss += lSup + CONFIG.boneWeight * lBone + CONFIG.temporalWeight * lTemp;
        batchBone += lBone;
        batchTemporal += lTemp;
        prevPred = pred;
      }
      batchLoss /= batch.length;
      batchBone /= batch.length;
      batchTemporal /= batch.length;

      // SPSA gradient update with combined loss
      const combinedLossFn = (m, s) => {
        const pred = m.forward(s.csi);
        return supervisedLoss(pred, s.keypoints, s.conf) +
               CONFIG.boneWeight * boneLoss(pred);
      };

      for (const p of allParams) {
        const grad = estimateBatchGrad(model, batch, combinedLossFn, p, rng);
        sgdStep(p, grad, lr, CONFIG.momentum);
      }

      epochLoss += batchLoss;
      epochBone += batchBone;
      epochTemporal += batchTemporal;
      nBatches++;
    }

    epochLoss = nBatches > 0 ? epochLoss / nBatches : 0;
    epochBone = nBatches > 0 ? epochBone / nBatches : 0;
    epochTemporal = nBatches > 0 ? epochTemporal / nBatches : 0;
    const evalResult = evaluate(model, evalSet);

    refinementLog.epochs.push({
      epoch: globalEpoch,
      loss: epochLoss,
      boneLoss: epochBone,
      temporalLoss: epochTemporal,
      evalLoss: evalResult.loss,
      pck20: evalResult.pck20,
      lr,
    });

    if ((epoch + 1) % 10 === 0 || epoch === 0) {
      console.log(`  [refinement] epoch ${epoch + 1}/${refinementEpochs}  loss=${epochLoss.toFixed(6)}  bone=${epochBone.toFixed(6)}  temporal=${epochTemporal.toFixed(6)}  eval_loss=${evalResult.loss.toFixed(6)}  PCK@20=${(evalResult.pck20 * 100).toFixed(1)}%  lr=${lr.toExponential(2)}`);
    }
    globalEpoch++;
  }

  trainingLog.phases.push(refinementLog);
  console.log('');

  // -----------------------------------------------------------------------
  // Step 6: Export
  // -----------------------------------------------------------------------
  console.log('[6/6] Exporting model and results...');

  fs.mkdirSync(CONFIG.outputDir, { recursive: true });

  // Export model weights as JSON
  const weights = model.getAllWeights();
  const modelExport = {
    format: 'wiflow-supervised-v1',
    adr: 'ADR-079',
    architecture: {
      inputDim,
      timeSteps: T,
      numKeypoints: CONFIG.numKeypoints,
      tcnChannels: [inputDim, 256, 256, 192, 128],
      tcnKernel: 7,
      tcnDilations: [1, 2, 4, 8],
      fcDims: [128 * T, 2048, CONFIG.numKeypoints * 2],
    },
    totalParams: model.totalParams(),
    weightsBase64: Buffer.from(weights.buffer).toString('base64'),
    trainingSamples: trainSet.length,
    evalSamples: evalSet.length,
    createdAt: new Date().toISOString(),
  };

  const modelPath = path.join(CONFIG.outputDir, 'wiflow-v1.json');
  fs.writeFileSync(modelPath, JSON.stringify(modelExport, null, 2));
  console.log(`  Model weights: ${modelPath} (${(fs.statSync(modelPath).size / 1024).toFixed(0)} KB)`);

  // Export training log
  const logPath = path.join(CONFIG.outputDir, 'training-log.json');
  fs.writeFileSync(logPath, JSON.stringify(trainingLog, null, 2));
  console.log(`  Training log:  ${logPath}`);

  // Export held-out predictions
  const evalPath = path.join(CONFIG.outputDir, 'eval-holdout.jsonl');
  const evalLines = [];
  for (const sample of evalSet) {
    const pred = model.forward(sample.csi);
    const pckScore = pck(pred, sample.keypoints, 0.2);
    evalLines.push(JSON.stringify({
      timestamp: sample.timestamp,
      predicted: Array.from(pred),
      groundTruth: Array.from(sample.keypoints),
      conf: Array.from(sample.conf),
      pck20: pckScore,
    }));
  }
  fs.writeFileSync(evalPath, evalLines.join('\n') + '\n');
  console.log(`  Eval holdout:  ${evalPath} (${evalSet.length} samples)`);

  // Final evaluation summary
  const finalEval = evaluate(model, evalSet);
  const elapsed = ((Date.now() - startTime) / 1000).toFixed(1);

  console.log('');
  console.log('=== Training Complete ===');
  console.log(`  Total epochs:     ${globalEpoch}`);
  console.log(`  Final eval loss:  ${finalEval.loss.toFixed(6)}`);
  console.log(`  Final PCK@20:     ${(finalEval.pck20 * 100).toFixed(1)}%`);
  console.log(`  Total parameters: ${model.totalParams().toLocaleString()}`);
  console.log(`  Elapsed:          ${elapsed}s`);
}

main().catch(err => {
  console.error('Training failed:', err);
  process.exit(1);
});