#!/usr/bin/env npx tsx
/**
 * Graph Layout
 *
 * Computes a 2D layout for a general graph (not necessarily a tree).
 * 
 * - Primary nodes (from primary_edges.csv) are placed first in a radial layout
 * - Remaining nodes are placed near their connected primary neighbors
 * - Barnes-Hut force simulation relaxes the layout
 *
 * Reads: primary_edges.csv, secondary_edges.csv
 * Writes: node_positions.csv
 *
 * Usage: npx tsx scripts/compute_layout.ts
 */

import { writeFileSync, readFileSync, existsSync } from "fs";
import { join, dirname } from "path";
import { fileURLToPath } from "url";

const __dirname = dirname(fileURLToPath(import.meta.url));
const PUBLIC_DIR = join(__dirname, "..", "public");

// ══════════════════════════════════════════════════════════
// Configuration
// ══════════════════════════════════════════════════════════

const ITERATIONS = 200;         // Force iterations
const REPULSION_K = 200;        // Repulsion strength
const EDGE_LENGTH = 80;         // Desired edge rest length
const ATTRACTION_K = 0.005;     // Spring stiffness for edges
const INITIAL_MAX_DISP = 20;    // Starting max displacement
const COOLING = 0.995;          // Cooling per iteration
const MIN_DIST = 0.5;
const PRINT_EVERY = 20;         // Print progress every N iterations
const BH_THETA = 0.8;          // Barnes-Hut opening angle

// Primary node radial placement
const PRIMARY_RADIUS = 300;     // Radius for primary node ring

// ══════════════════════════════════════════════════════════
// Read edge data from CSVs
// ══════════════════════════════════════════════════════════

const primaryPath = join(PUBLIC_DIR, "primary_edges.csv");
const secondaryPath = join(PUBLIC_DIR, "secondary_edges.csv");

if (!existsSync(primaryPath) || !existsSync(secondaryPath)) {
    console.error(`Error: Missing input files!`);
    console.error(`  Expected: ${primaryPath}`);
    console.error(`  Expected: ${secondaryPath}`);
    console.error(`  Run 'npx tsx scripts/fetch_from_db.ts' first.`);
    process.exit(1);
}

function parseEdges(path: string): Array<[number, number]> {
    const content = readFileSync(path, "utf-8");
    const lines = content.trim().split("\n");
    const edges: Array<[number, number]> = [];
    for (let i = 1; i < lines.length; i++) {
        const line = lines[i].trim();
        if (!line) continue;
        const [src, tgt] = line.split(",").map(Number);
        if (!isNaN(src) && !isNaN(tgt)) {
            edges.push([src, tgt]);
        }
    }
    return edges;
}

const primaryEdges = parseEdges(primaryPath);
const secondaryEdges = parseEdges(secondaryPath);
const allEdges = [...primaryEdges, ...secondaryEdges];

// ══════════════════════════════════════════════════════════
// Build adjacency
// ══════════════════════════════════════════════════════════

const allNodes = new Set<number>();
const primaryNodes = new Set<number>();
const neighbors = new Map<number, Set<number>>();

function addNeighbor(a: number, b: number) {
    if (!neighbors.has(a)) neighbors.set(a, new Set());
    neighbors.get(a)!.add(b);
    if (!neighbors.has(b)) neighbors.set(b, new Set());
    neighbors.get(b)!.add(a);
}

for (const [src, dst] of primaryEdges) {
    allNodes.add(src);
    allNodes.add(dst);
    primaryNodes.add(src);
    primaryNodes.add(dst);
    addNeighbor(src, dst);
}

for (const [src, dst] of secondaryEdges) {
    allNodes.add(src);
    allNodes.add(dst);
    addNeighbor(src, dst);
}

const N = allNodes.size;
const nodeIds = Array.from(allNodes).sort((a, b) => a - b);
const idToIdx = new Map<number, number>();
nodeIds.forEach((id, idx) => idToIdx.set(id, idx));

console.log(
    `Read graph: ${N} nodes, ${allEdges.length} edges (P=${primaryEdges.length}, S=${secondaryEdges.length})`
);
console.log(`Primary nodes: ${primaryNodes.size}`);

// ══════════════════════════════════════════════════════════
// Initial placement
// ══════════════════════════════════════════════════════════

const x = new Float64Array(N);
const y = new Float64Array(N);

// Step 1: Place primary nodes in a radial layout
const primaryArr = Array.from(primaryNodes).sort((a, b) => a - b);
const angleStep = (2 * Math.PI) / Math.max(1, primaryArr.length);
const radius = PRIMARY_RADIUS * Math.max(1, Math.sqrt(primaryArr.length / 10));

for (let i = 0; i < primaryArr.length; i++) {
    const idx = idToIdx.get(primaryArr[i])!;
    const angle = i * angleStep;
    x[idx] = radius * Math.cos(angle);
    y[idx] = radius * Math.sin(angle);
}

console.log(`Placed ${primaryArr.length} primary nodes in radial layout (r=${radius.toFixed(0)})`);

// Step 2: Place remaining nodes near their connected neighbors
// BFS from already-placed nodes
const placed = new Set<number>(primaryNodes);
const queue: number[] = [...primaryArr];
let head = 0;

while (head < queue.length) {
    const nodeId = queue[head++];
    const nodeNeighbors = neighbors.get(nodeId);
    if (!nodeNeighbors) continue;

    for (const nbId of nodeNeighbors) {
        if (placed.has(nbId)) continue;
        placed.add(nbId);

        // Place near this neighbor with some jitter
        const parentIdx = idToIdx.get(nodeId)!;
        const childIdx = idToIdx.get(nbId)!;
        const jitterAngle = Math.random() * 2 * Math.PI;
        const jitterDist = EDGE_LENGTH * (0.5 + Math.random() * 0.5);
        x[childIdx] = x[parentIdx] + jitterDist * Math.cos(jitterAngle);
        y[childIdx] = y[parentIdx] + jitterDist * Math.sin(jitterAngle);

        queue.push(nbId);
    }
}

// Handle disconnected nodes (place randomly)
for (const id of nodeIds) {
    if (!placed.has(id)) {
        const idx = idToIdx.get(id)!;
        const angle = Math.random() * 2 * Math.PI;
        const r = radius * (1 + Math.random());
        x[idx] = r * Math.cos(angle);
        y[idx] = r * Math.sin(angle);
        placed.add(id);
    }
}

console.log(`Initial placement complete: ${placed.size} nodes`);

// ══════════════════════════════════════════════════════════
// Force-directed layout with Barnes-Hut
// ══════════════════════════════════════════════════════════

console.log(`Running force simulation (${ITERATIONS} iterations, ${N} nodes, ${allEdges.length} edges)...`);

const t0 = performance.now();
let maxDisp = INITIAL_MAX_DISP;

for (let iter = 0; iter < ITERATIONS; iter++) {
    const bhRoot = buildBHTree(x, y, N);
    const fx = new Float64Array(N);
    const fy = new Float64Array(N);

    // 1. Repulsion via Barnes-Hut
    for (let i = 0; i < N; i++) {
        calcBHForce(bhRoot, x[i], y[i], fx, fy, i, BH_THETA, x, y);
    }

    // 2. Edge attraction (spring force)
    for (const [aId, bId] of allEdges) {
        const a = idToIdx.get(aId)!;
        const b = idToIdx.get(bId)!;
        const dx = x[b] - x[a];
        const dy = y[b] - y[a];
        const d = Math.sqrt(dx * dx + dy * dy) || MIN_DIST;
        const displacement = d - EDGE_LENGTH;
        const f = ATTRACTION_K * displacement;
        const ux = dx / d;
        const uy = dy / d;
        fx[a] += ux * f;
        fy[a] += uy * f;
        fx[b] -= ux * f;
        fy[b] -= uy * f;
    }

    // 3. Apply forces with displacement capping
    let totalForce = 0;
    for (let i = 0; i < N; i++) {
        const mag = Math.sqrt(fx[i] * fx[i] + fy[i] * fy[i]);
        totalForce += mag;
        if (mag > 0) {
            const cap = Math.min(maxDisp, mag) / mag;
            x[i] += fx[i] * cap;
            y[i] += fy[i] * cap;
        }
    }

    maxDisp *= COOLING;

    if ((iter + 1) % PRINT_EVERY === 0 || iter === 0) {
        console.log(
            `  iter ${iter + 1}/${ITERATIONS}  max_disp=${maxDisp.toFixed(2)}  avg_force=${(totalForce / N).toFixed(2)}`
        );
    }
}

const elapsed = performance.now() - t0;
console.log(`Force simulation done in ${(elapsed / 1000).toFixed(1)}s`);

// ══════════════════════════════════════════════════════════
// Write output
// ══════════════════════════════════════════════════════════

const outLines: string[] = ["vertex,x,y"];
for (let i = 0; i < N; i++) {
    outLines.push(`${nodeIds[i]},${x[i]},${y[i]}`);
}

const outPath = join(PUBLIC_DIR, "node_positions.csv");
writeFileSync(outPath, outLines.join("\n") + "\n");
console.log(`Wrote ${N} positions to ${outPath}`);
console.log(`Layout complete.`);

// ══════════════════════════════════════════════════════════
// Barnes-Hut Helpers
// ══════════════════════════════════════════════════════════

interface BHNode {
    mass: number;
    cx: number;
    cy: number;
    minX: number;
    maxX: number;
    minY: number;
    maxY: number;
    children?: BHNode[];
    pointIdx?: number;
}

function buildBHTree(x: Float64Array, y: Float64Array, n: number): BHNode {
    let minX = Infinity, maxX = -Infinity, minY = Infinity, maxY = -Infinity;
    for (let i = 0; i < n; i++) {
        if (x[i] < minX) minX = x[i];
        if (x[i] > maxX) maxX = x[i];
        if (y[i] < minY) minY = y[i];
        if (y[i] > maxY) maxY = y[i];
    }
    const cx = (minX + maxX) / 2;
    const cy = (minY + maxY) / 2;
    const halfDim = Math.max(maxX - minX, maxY - minY) / 2 + 0.01;

    const root: BHNode = {
        mass: 0, cx: 0, cy: 0,
        minX: cx - halfDim, maxX: cx + halfDim,
        minY: cy - halfDim, maxY: cy + halfDim,
    };

    for (let i = 0; i < n; i++) {
        insertBH(root, i, x[i], y[i], x, y);
    }
    calcBHMass(root, x, y);
    return root;
}

function insertBH(node: BHNode, idx: number, px: number, py: number, x: Float64Array, y: Float64Array) {
    if (!node.children && node.pointIdx === undefined) {
        node.pointIdx = idx;
        return;
    }

    if (!node.children && node.pointIdx !== undefined) {
        const oldIdx = node.pointIdx;
        node.pointIdx = undefined;
        subdivideBH(node);
        insertBH(node, oldIdx, x[oldIdx], y[oldIdx], x, y);
    }

    if (node.children) {
        const mx = (node.minX + node.maxX) / 2;
        const my = (node.minY + node.maxY) / 2;
        let q = 0;
        if (px > mx) q += 1;
        if (py > my) q += 2;
        insertBH(node.children[q], idx, px, py, x, y);
    }
}

function subdivideBH(node: BHNode) {
    const mx = (node.minX + node.maxX) / 2;
    const my = (node.minY + node.maxY) / 2;
    node.children = [
        { mass: 0, cx: 0, cy: 0, minX: node.minX, maxX: mx, minY: node.minY, maxY: my },
        { mass: 0, cx: 0, cy: 0, minX: mx, maxX: node.maxX, minY: node.minY, maxY: my },
        { mass: 0, cx: 0, cy: 0, minX: node.minX, maxX: mx, minY: my, maxY: node.maxY },
        { mass: 0, cx: 0, cy: 0, minX: mx, maxX: node.maxX, minY: my, maxY: node.maxY },
    ];
}

function calcBHMass(node: BHNode, x: Float64Array, y: Float64Array) {
    if (node.pointIdx !== undefined) {
        node.mass = 1;
        node.cx = x[node.pointIdx];
        node.cy = y[node.pointIdx];
        return;
    }
    if (node.children) {
        let m = 0, sx = 0, sy = 0;
        for (const c of node.children) {
            calcBHMass(c, x, y);
            m += c.mass;
            sx += c.cx * c.mass;
            sy += c.cy * c.mass;
        }
        node.mass = m;
        if (m > 0) {
            node.cx = sx / m;
            node.cy = sy / m;
        } else {
            node.cx = (node.minX + node.maxX) / 2;
            node.cy = (node.minY + node.maxY) / 2;
        }
    }
}

function calcBHForce(
    node: BHNode,
    px: number, py: number,
    fx: Float64Array, fy: Float64Array,
    idx: number, theta: number,
    x: Float64Array, y: Float64Array,
) {
    const dx = px - node.cx;
    const dy = py - node.cy;
    const d2 = dx * dx + dy * dy;
    const dist = Math.sqrt(d2);
    const width = node.maxX - node.minX;

    if (width / dist < theta || !node.children) {
        if (node.mass > 0 && node.pointIdx !== idx) {
            const dEff = Math.max(dist, MIN_DIST);
            const f = (REPULSION_K * node.mass) / (dEff * dEff);
            fx[idx] += (dx / dEff) * f;
            fy[idx] += (dy / dEff) * f;
        }
    } else {
        for (const c of node.children) {
            calcBHForce(c, px, py, fx, fy, idx, theta, x, y);
        }
    }
}