import type { EventUsage, InferenceEvent } from "./types.ts";

export type DatedEventInput = Readonly<{
  id: string;
  domain: string;
  year: number | null;
  precision: InferenceEvent["precision"];
}>;

/**
 * Domain-stratified holdout: keep at least one dated event out of training
 * so the winner is not certified by the same fact that selected it.
 */
export function splitHoldoutEvents(events: readonly DatedEventInput[]): InferenceEvent[] {
  const dated = events.filter((item) => item.year !== null && item.precision !== "unknown");
  if (dated.length < 4) {
    return events.map((item) => ({ ...item, usage: "training" as const }));
  }
  const byDomain = new Map<string, DatedEventInput[]>();
  for (const item of dated) {
    const rows = byDomain.get(item.domain) ?? [];
    rows.push(item);
    byDomain.set(item.domain, rows);
  }
  const holdoutId = pickHoldoutId(byDomain, dated);
  return events.map((item) => ({
    ...item,
    usage: usageFor(item, holdoutId, dated),
  }));
}

function usageFor(item: DatedEventInput, holdoutId: string | null, dated: readonly DatedEventInput[]): EventUsage {
  if (item.year === null || item.precision === "unknown") return "unused";
  if (item.id === holdoutId) return "holdout";
  return dated.some((row) => row.id === item.id) ? "training" : "unused";
}

function pickHoldoutId(
  byDomain: Map<string, DatedEventInput[]>,
  dated: readonly DatedEventInput[],
): string | null {
  const monthOrBetter = dated.filter((item) => item.precision === "day" || item.precision === "month");
  const singletonDomain = [...byDomain.entries()].find(([, rows]) => rows.length === 1);
  if (singletonDomain && dated.length - 1 >= 3) {
    const preferred = monthOrBetter.find((item) => item.domain === singletonDomain[0]);
    return (preferred ?? singletonDomain[1][0])?.id ?? null;
  }
  const ranked = [...monthOrBetter, ...dated];
  return ranked[ranked.length - 1]?.id ?? null;
}

export function trainingEventIds(events: readonly InferenceEvent[]): ReadonlySet<string> {
  return new Set(events.filter((item) => item.usage === "training").map((item) => item.id));
}

export function holdoutEventIds(events: readonly InferenceEvent[]): ReadonlySet<string> {
  return new Set(events.filter((item) => item.usage === "holdout").map((item) => item.id));
}