import { maskAnd, maskOr, morphDilate, dropSmall, suppressThinLines } from "@/utils/maskOps";
import { coverageOfMask, leakRatio } from "./lattice";

type AutoTuneOptions = {
  targetLow: number;
  targetHigh: number;
  minArea: number;
  thinWidth: number;
  thinLength: number;
  leakLimit: number;
  maxDilateSteps: number;
  u2BoostMask?: Uint8Array | null;
};

export type AutoTuneDiagnostics = {
  dilateSteps: number;
  shrinkApplied: boolean;
  u2BoostApplied: boolean;
};

export type AutoTuneResult = {
  mask: Uint8Array;
  coverage: number;
  leak: number;
  adjustments: AutoTuneDiagnostics;
};

type LeakExtras = {
  weights?: Float32Array;
  axisMap?: Int8Array;
  gradX?: Float32Array;
  gradY?: Float32Array;
  gradientThreshold?: number;
  angleThresholdRad?: number;
};

export function autoTuneMask(
  baseMask: Uint8Array,
  interiorMask: Uint8Array,
  lineMask: Uint8Array,
  width: number,
  height: number,
  options: AutoTuneOptions,
  leakExtras: LeakExtras = {},
): AutoTuneResult {
  const clamp01 = (v: number) => Math.max(0, Math.min(1, v));
  let tuned = maskAnd(baseMask, interiorMask);

  const thinWidth = Math.max(0, Math.round(options.thinWidth));
  const thinLength = clamp01(options.thinLength);
  if (thinWidth > 0) {
    tuned = suppressThinLines(tuned, width, height, thinWidth, thinLength) as Uint8Array;
  }
  tuned = dropSmall(tuned, width, height, Math.max(1, Math.round(options.minArea))) as Uint8Array;

  let coverage = coverageOfMask(tuned);
  let leak = leakRatio(tuned, lineMask, leakExtras);
  const adjustments: AutoTuneDiagnostics = {
    dilateSteps: 0,
    shrinkApplied: false,
    u2BoostApplied: false,
  };

  const leakLimit = Math.max(0, options.leakLimit);
  while (coverage < options.targetLow && adjustments.dilateSteps < options.maxDilateSteps) {
    const grown = maskAnd(morphDilate(tuned, width, height, 1), interiorMask) as Uint8Array;
    const leakCandidate = leakRatio(grown, lineMask, leakExtras);
    if (leakCandidate > leakLimit + 0.01) break;
    tuned = grown;
    adjustments.dilateSteps += 1;
    if (thinWidth > 0) {
      tuned = suppressThinLines(tuned, width, height, thinWidth, thinLength) as Uint8Array;
    }
    tuned = dropSmall(tuned, width, height, Math.max(1, Math.round(options.minArea))) as Uint8Array;
    coverage = coverageOfMask(tuned);
    leak = leakRatio(tuned, lineMask, leakExtras);
    if (coverage >= options.targetLow) break;
  }

  if (coverage > options.targetHigh + 0.02) {
    const tighter = dropSmall(tuned, width, height, Math.round(options.minArea * 1.2)) as Uint8Array;
    const tightCoverage = coverageOfMask(tighter);
    if (tightCoverage >= options.targetLow - 0.02) {
      tuned = tighter;
      coverage = tightCoverage;
      leak = leakRatio(tuned, lineMask, leakExtras);
      adjustments.shrinkApplied = true;
    }
  }

  if (coverage < options.targetLow && options.u2BoostMask) {
    const boosted = maskOr(tuned, maskAnd(options.u2BoostMask, interiorMask) as Uint8Array) as Uint8Array;
    const leakBoost = leakRatio(boosted, lineMask, leakExtras);
    const boostedCoverage = coverageOfMask(boosted);
    if (boostedCoverage > coverage && leakBoost <= leakLimit + 0.01) {
      tuned = dropSmall(boosted, width, height, Math.max(1, Math.round(options.minArea))) as Uint8Array;
      coverage = coverageOfMask(tuned);
      leak = leakRatio(tuned, lineMask, leakExtras);
      adjustments.u2BoostApplied = true;
    }
  }

  return { mask: tuned, coverage, leak, adjustments };
}