Add `pose_estimator` module to complete the natural-feature AR pipeline

import { matrix_t } from "../matrix_t/matrix_t";

/** One correspondence, equivalent to OpenCV's `cv::DMatch`. */
export class match_t {
    public queryIdx: number;
    public trainIdx: number;
    public distance: number;

    constructor(queryIdx = 0, trainIdx = 0, distance = 0) {
        this.queryIdx = queryIdx;
        this.trainIdx = trainIdx;
        this.distance = distance;
    }
}

export class bfmatcher {
    /** Mirrors OpenCV's NORM_HAMMING constant value. Only norm supported here. */
    static readonly NORM_HAMMING = 6;

    public norm_type: number;
    public cross_check: boolean;

    constructor(norm_type: number = bfmatcher.NORM_HAMMING, cross_check = false) {
        this.norm_type = norm_type;
        this.cross_check = cross_check;
    }

    // SWAR population count — identical to the sample's popcnt32, so results
    // match bit-for-bit and this can serve as a reference oracle for the Rust
    // port (u32::count_ones / i32.popcnt).
    private static popcnt32(n: number): number {
        n -= (n >> 1) & 0x55555555;
        n = (n & 0x33333333) + ((n >> 2) & 0x33333333);
        return (((n + (n >> 4)) & 0x0f0f0f0f) * 0x01010101) >> 24;
    }

    /** Read a U8 descriptor matrix as an Int32Array view (mirrors `buffer.i32`). */
    private static asInt32(m: matrix_t): Int32Array {
        const d = m.data as Uint8Array;
        return new Int32Array(d.buffer, d.byteOffset, d.byteLength >> 2);
    }

    private static hamming(
        qb: Int32Array, qoff: number,
        tb: Int32Array, toff: number,
        int_len: number
    ): number {
        let d = 0;
        for (let k = 0; k < int_len; ++k) {
            d += bfmatcher.popcnt32(qb[qoff + k] ^ tb[toff + k]);
        }
        return d;
    }

    /**
     * Nearest-neighbour match. With `cross_check`, keeps only mutually-best
     * pairs (q,t): q must also be t's best match back.
     */
    match(query: matrix_t, train: matrix_t, max_distance = 256): match_t[] {
        const q_cnt = query.rows;
        const t_cnt = train.rows;
        const int_len = query.cols >> 2;
        const qb = bfmatcher.asInt32(query);
        const tb = bfmatcher.asInt32(train);

        const out: match_t[] = [];
        const fwd = new Int32Array(q_cnt);

        for (let qi = 0; qi < q_cnt; ++qi) {
            const qoff = qi * int_len;
            let best = 0x7fffffff;
            let best_idx = -1;
            for (let ti = 0; ti < t_cnt; ++ti) {
                const d = bfmatcher.hamming(qb, qoff, tb, ti * int_len, int_len);
                if (d < best) { best = d; best_idx = ti; }
            }
            fwd[qi] = best_idx;
            if (!this.cross_check && best_idx >= 0 && best <= max_distance) {
                out.push(new match_t(qi, best_idx, best));
            }
        }

        if (!this.cross_check) return out;

        for (let qi = 0; qi < q_cnt; ++qi) {
            const ti = fwd[qi];
            if (ti < 0) continue;
            const toff = ti * int_len;
            let best = 0x7fffffff;
            let back_idx = -1;
            for (let qj = 0; qj < q_cnt; ++qj) {
                const d = bfmatcher.hamming(qb, qj * int_len, tb, toff, int_len);
                if (d < best) { best = d; back_idx = qj; }
            }
            if (back_idx === qi && best <= max_distance) {
                out.push(new match_t(qi, ti, best));
            }
        }
        return out;
    }

    /** k-nearest matches per query, sorted ascending by distance. */
    knnMatch(query: matrix_t, train: matrix_t, k = 2): match_t[][] {
        const q_cnt = query.rows;
        const t_cnt = train.rows;
        const int_len = query.cols >> 2;
        const qb = bfmatcher.asInt32(query);
        const tb = bfmatcher.asInt32(train);

        const result: match_t[][] = [];
        for (let qi = 0; qi < q_cnt; ++qi) {
            const qoff = qi * int_len;
            const top: match_t[] = [];
            for (let ti = 0; ti < t_cnt; ++ti) {
                const d = bfmatcher.hamming(qb, qoff, tb, ti * int_len, int_len);
                if (top.length < k) {
                    top.push(new match_t(qi, ti, d));
                    top.sort((a, b) => a.distance - b.distance);
                } else if (d < top[top.length - 1].distance) {
                    const slot = top[top.length - 1];
                    slot.queryIdx = qi; slot.trainIdx = ti; slot.distance = d;
                    top.sort((a, b) => a.distance - b.distance);
                }
            }
            result.push(top);
        }
        return result;
    }

    /** Lowe's ratio test over `knnMatch(query, train, 2)`. */
    static ratio_test(knn: match_t[][], ratio = 0.75): match_t[] {
        const good: match_t[] = [];
        for (let i = 0; i < knn.length; ++i) {
            const m = knn[i];
            if (m.length >= 2) {
                if (m[0].distance < ratio * m[1].distance) good.push(m[0]);
            } else if (m.length === 1) {
                good.push(m[0]);
            }
        }
        return good;
    }
}

import { matrix_t } from "../matrix_t/matrix_t";
import { JSFEAT_CONSTANTS } from "../constants/constants";

export class pose_t {
    public R: matrix_t;     // 3x3, row-major, OpenCV camera frame
    public t: Float64Array; // length 3
    public good: boolean;   // false if H/K were degenerate

    constructor() {
        this.R = new matrix_t(3, 3, JSFEAT_CONSTANTS.F64_t | JSFEAT_CONSTANTS.C1_t);
        this.t = new Float64Array(3);
        this.good = false;
    }
}

export class pose_estimator {
    private Kinv: Float64Array; // inverse intrinsics, row-major, zero skew

    constructor(K: matrix_t) {
        this.Kinv = pose_estimator.invertIntrinsics(K);
    }

    setIntrinsics(K: matrix_t): void {
        this.Kinv = pose_estimator.invertIntrinsics(K);
    }

    private static invertIntrinsics(K: matrix_t): Float64Array {
        const k = K.data;
        const fx = k[0], cx = k[2], fy = k[4], cy = k[5];
        return new Float64Array([
            1 / fx, 0, -cx / fx,
            0, 1 / fy, -cy / fy,
            0, 0, 1,
        ]);
    }

    /** Recover pose from a homography mapping model-plane points to pixels. */
    estimate(H: matrix_t, out?: pose_t): pose_t {
        const pose = out || new pose_t();
        const h = H.data;
        const ki = this.Kinv;

        // B = Kinv * H, row-major
        const B = new Float64Array(9);
        for (let r = 0; r < 3; ++r) {
            for (let c = 0; c < 3; ++c) {
                B[r * 3 + c] =
                    ki[r * 3] * h[c] +
                    ki[r * 3 + 1] * h[3 + c] +
                    ki[r * 3 + 2] * h[6 + c];
            }
        }

        const b1 = [B[0], B[3], B[6]];
        const b2 = [B[1], B[4], B[7]];
        const b3 = [B[2], B[5], B[8]];

        const n1 = Math.hypot(b1[0], b1[1], b1[2]);
        const n2 = Math.hypot(b2[0], b2[1], b2[2]);
        if (n1 < 1e-12 || n2 < 1e-12) { pose.good = false; return pose; }

        const lambda = 2.0 / (n1 + n2);
        const s = b3[2] >= 0 ? 1 : -1; // target in front of camera (t.z > 0)

        const r1 = [s * b1[0] / n1, s * b1[1] / n1, s * b1[2] / n1];
        const r2 = [s * b2[0] / n2, s * b2[1] / n2, s * b2[2] / n2];
        const t  = [s * b3[0] * lambda, s * b3[1] * lambda, s * b3[2] * lambda];

        // Closed-form re-orthonormalization of (r1, r2):
        const c = pose_estimator.normalize([r1[0] + r2[0], r1[1] + r2[1], r1[2] + r2[2]]);
        const p = pose_estimator.cross(r1, r2);
        const d = pose_estimator.normalize(pose_estimator.cross(c, p));
        const q = Math.SQRT1_2;

        const c1 = [(c[0] + d[0]) * q, (c[1] + d[1]) * q, (c[2] + d[2]) * q];
        const c2 = [(c[0] - d[0]) * q, (c[1] - d[1]) * q, (c[2] - d[2]) * q];
        const c3 = pose_estimator.cross(c1, c2);

        const R = pose.R.data; // store columns [c1 c2 c3] into row-major R
        R[0] = c1[0]; R[1] = c2[0]; R[2] = c3[0];
        R[3] = c1[1]; R[4] = c2[1]; R[5] = c3[1];
        R[6] = c1[2]; R[7] = c2[2]; R[8] = c3[2];

        pose.t[0] = t[0]; pose.t[1] = t[1]; pose.t[2] = t[2];
        pose.good = true;
        return pose;
    }

    /** Column-major 4x4 modelview for three.js / WebGL, CV -> GL via diag(1,-1,-1). */
    modelViewGL(pose: pose_t, out?: Float32Array): Float32Array {
        const m = out || new Float32Array(16);
        const R = pose.R.data;
        const t = pose.t;
        m[0] = R[0];  m[1] = -R[3];  m[2] = -R[6];  m[3] = 0;
        m[4] = R[1];  m[5] = -R[4];  m[6] = -R[7];  m[7] = 0;
        m[8] = R[2];  m[9] = -R[5];  m[10] = -R[8]; m[11] = 0;
        m[12] = t[0]; m[13] = -t[1]; m[14] = -t[2]; m[15] = 1;
        return m;
    }

    /** Rough pinhole intrinsics from image size + horizontal FOV (uncalibrated bootstrap). */
    static intrinsics(width: number, height: number, fovXdeg = 60): matrix_t {
        const K = new matrix_t(3, 3, JSFEAT_CONSTANTS.F64_t | JSFEAT_CONSTANTS.C1_t);
        const f = 0.5 * width / Math.tan(0.5 * fovXdeg * Math.PI / 180);
        const k = K.data;
        k[0] = f; k[1] = 0; k[2] = 0.5 * width;
        k[3] = 0; k[4] = f; k[5] = 0.5 * height;
        k[6] = 0; k[7] = 0; k[8] = 1;
        return K;
    }

    private static cross(a: number[], b: number[]): number[] {
        return [
            a[1] * b[2] - a[2] * b[1],
            a[2] * b[0] - a[0] * b[2],
            a[0] * b[1] - a[1] * b[0],
        ];
    }

    private static normalize(v: number[]): number[] {
        const n = Math.hypot(v[0], v[1], v[2]) || 1;
        return [v[0] / n, v[1] / n, v[2] / n];
    }
}