// Upscale a pose layer, snap it to its own k main colours, then smooth every colour boundary before tracing.
// The poses are only ~300-400 px across, so a plain upscale keeps 1-px stair-steps that vtracer traces as jagged edges.
// Smoothing = blur each colour's mask (and the transparent mask) and give every pixel the colour whose blurred mask is strongest.
// Thin anti-alias bands (a colour that sits between two others) are folded into the nearer neighbour first.
const sharp = require("module").createRequire("/home/ubuntu/kalki-cli/projects/mybazaar-web/package.json")("sharp");
// The outer silhouette gets its own, stronger blur: the cut-out's edge wobbles more than the drawn colour boundaries.
const [, , input, output, kArg = "9", scaleArg = "4", sigmaArg = "3.5", edgeSigmaArg = "7"] = process.argv;
const k = +kArg, scale = +scaleArg, sigma = +sigmaArg, edgeSigma = +edgeSigmaArg;
const dist2 = (a, b) => (a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2 + (a[2] - b[2]) ** 2;
(async () => {
	const meta = await sharp(input).metadata();
	const { data, info } = await sharp(input).ensureAlpha()
		.resize(meta.width * scale, meta.height * scale, { kernel: "lanczos3" })
		.raw().toBuffer({ resolveWithObject: true });
	const W = info.width, H = info.height, n = W * H;
	const samples = [];
	for (let p = 0; p < n; p += 5) if (data[p * 4 + 3] > 200) samples.push([data[p * 4], data[p * 4 + 1], data[p * 4 + 2]]);
	let centres = [samples[0]];
	while (centres.length < k) {
		let best = null, bestD = -1;
		for (let s = 0; s < samples.length; s += 11) {
			const d = Math.min(...centres.map((c) => dist2(c, samples[s])));
			if (d > bestD) { bestD = d; best = samples[s]; }
		}
		centres.push(best);
	}
	const nearestOf = (cs) => (px) => { let bi = 0, bd = Infinity; cs.forEach((c, i) => { const d = dist2(c, px); if (d < bd) { bd = d; bi = i; } }); return bi; };
	for (let it = 0; it < 20; it++) {
		const near = nearestOf(centres), acc = centres.map(() => [0, 0, 0, 0]);
		for (const s of samples) { const a = acc[near(s)]; a[0] += s[0]; a[1] += s[1]; a[2] += s[2]; a[3]++; }
		centres = centres.map((c, i) => (acc[i][3] ? acc[i].slice(0, 3).map((v) => v / acc[i][3]) : c));
	}
	const near = nearestOf(centres);
	const T = k; // the transparent label
	const label = new Uint8Array(n);
	for (let p = 0; p < n; p++) label[p] = data[p * 4 + 3] < 128 ? T : near([data[p * 4], data[p * 4 + 1], data[p * 4 + 2]]);

	// A blend colour lies on the line between two other colours. Each of its pixels that touches both end colours
	// is anti-aliasing and goes to the nearer end; free-standing ones (a gold sparkle between orange and white) stay.
	const pairsOf = centres.map((c, i) => {
		const pairs = [];
		centres.forEach((a, ai) => centres.forEach((b, bi) => {
			if (ai >= bi || ai === i || bi === i) return;
			const ab = [b[0] - a[0], b[1] - a[1], b[2] - a[2]], len2 = dist2(a, b);
			const t = ((c[0] - a[0]) * ab[0] + (c[1] - a[1]) * ab[1] + (c[2] - a[2]) * ab[2]) / len2;
			if (t < 0.15 || t > 0.85 || len2 < 60 ** 2) return;
			if (Math.sqrt(dist2(c, [a[0] + t * ab[0], a[1] + t * ab[1], a[2] + t * ab[2]])) < 14) pairs.push([ai, bi]);
		}));
		return pairs;
	});
	const count = new Array(k + 1).fill(0);
	for (let p = 0; p < n; p++) count[label[p]]++;
	const opaque = n - count[T];
	const reach = new Map();
	const nearMask = async (L) => {
		if (!reach.has(L)) {
			const m = Buffer.alloc(n);
			for (let p = 0; p < n; p++) if (label[p] === L) m[p] = 255;
			reach.set(L, await sharp(m, { raw: { width: W, height: H, channels: 1 } }).blur(2 * scale).extractChannel(0).raw().toBuffer());
		}
		return reach.get(L);
	};
	const blends = centres.map(() => false);
	const relabel = new Uint8Array(label);
	for (let i = 0; i < k; i++) {
		// Only small clusters can be anti-aliasing; a large one is a real shade in the art.
		if (!pairsOf[i].length || count[i] / opaque >= 0.06) continue;
		const masks = await Promise.all(pairsOf[i].map(async ([a, b]) => [a, b, await nearMask(a), await nearMask(b)]));
		let moved = 0;
		for (let p = 0; p < n; p++) {
			if (label[p] !== i) continue;
			const hit = masks.find(([, , ma, mb]) => ma[p] > 8 && mb[p] > 8);
			if (!hit) continue;
			const px = [data[p * 4], data[p * 4 + 1], data[p * 4 + 2]];
			relabel[p] = dist2(centres[hit[0]], px) < dist2(centres[hit[1]], px) ? hit[0] : hit[1];
			moved++;
		}
		blends[i] = moved / count[i] > 0.5;
	}
	label.set(relabel);

	// Two near-identical colours (e.g. a dark outline split in two by noise) become the larger one, unless the smaller
	// is a sizeable area of its own (a vampire cape's dark lining is close to the outline colour but is real art).
	const merged = centres.map((_, i) => i);
	centres.forEach((c, i) => centres.forEach((d, j) => {
		if (i !== j && !blends[i] && !blends[j] && dist2(c, d) < 40 ** 2 && count[i] < count[j] && count[i] / opaque < 0.04) merged[i] = merged[j] === i ? i : j;
	}));
	for (let p = 0; p < n; p++) if (label[p] !== T) label[p] = merged[label[p]];

	const opaqueMask = Buffer.alloc(n);
	for (let p = 0; p < n; p++) if (label[p] !== T) opaqueMask[p] = 255;
	const silhouette = await sharp(opaqueMask, { raw: { width: W, height: H, channels: 1 } }).blur(edgeSigma).extractChannel(0).raw().toBuffer();
	const live = [...new Set(label)].filter((L) => L !== T);
	const score = new Float32Array(n).fill(-1), winner = new Uint8Array(n).fill(T);
	for (const L of live) {
		const mask = Buffer.alloc(n);
		for (let p = 0; p < n; p++) if (label[p] === L) mask[p] = 255;
		const blurred = await sharp(mask, { raw: { width: W, height: H, channels: 1 } }).blur(sigma).extractChannel(0).raw().toBuffer();
		// A blend-like colour loses close calls, so only its solid areas (a sparkle) survive, not leftover fringes.
		const weight = pairsOf[L]?.length ? 0.6 : 1;
		for (let p = 0; p < n; p++) if (blurred[p] * weight > score[p]) { score[p] = blurred[p] * weight; winner[p] = L; }
	}
	// Inside the smoothed silhouette every pixel takes its strongest colour; outside it is transparent.
	for (let p = 0; p < n; p++) if (silhouette[p] < 128) winner[p] = T;
	// Tiny isolated patches (flecks of a lining colour along an outline, a speck inside a fang) join the colour around them.
	const minArea = 8 * scale * scale;
	const seen = new Uint8Array(n), stack = new Int32Array(n);
	for (let start = 0; start < n; start++) {
		if (seen[start] || winner[start] === T) continue;
		const L = winner[start], members = [];
		let top = 0;
		stack[top++] = start; seen[start] = 1;
		const around = new Map();
		while (top) {
			const p = stack[--top];
			members.push(p);
			const x = p % W;
			for (const q of [p - W, p + W, x > 0 ? p - 1 : -1, x < W - 1 ? p + 1 : -1]) {
				if (q < 0 || q >= n) continue;
				if (winner[q] === L) { if (!seen[q]) { seen[q] = 1; stack[top++] = q; } } else around.set(winner[q], (around.get(winner[q]) || 0) + 1);
			}
		}
		if (members.length >= minArea) continue;
		let into = T, most = 0;
		for (const [c, m] of around) if (c !== T && m > most) { most = m; into = c; }
		if (into !== T) for (const p of members) winner[p] = into;
	}
	const out = Buffer.alloc(n * 4);
	centres = centres.map((c) => c.map(Math.round));
	for (let p = 0; p < n; p++) {
		if (winner[p] === T) continue;
		const c = centres[winner[p]];
		out[p * 4] = c[0]; out[p * 4 + 1] = c[1]; out[p * 4 + 2] = c[2]; out[p * 4 + 3] = 255;
	}
	await sharp(out, { raw: { width: W, height: H, channels: 4 } }).png().toFile(output);
	const hex = (c) => "#" + c.map((v) => v.toString(16).padStart(2, "0")).join("");
	console.log(`${input.split("/").pop()}: ${centres.map((c, i) => `${hex(c)} ${(100 * count[i] / opaque).toFixed(1)}%${blends[i] ? " (blend)" : ""}`).join(", ")}`);
})();