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

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

/** Reserve a holdout only when collection already has four dated events.
 * Discrimination still waits until 3 training events remain.
 */
export const MIN_EVENTS_TO_RESERVE_HOLDOUT = 4;

export type HoldoutReservationStatus = "reserved" | "not_reserved_min_events";

export function holdoutReservationStatus(events: readonly DatedEventInput[]): HoldoutReservationStatus {
  const dated = events.filter((item) => item.year !== null && item.precision !== "unknown");
  if (dated.length < MIN_EVENTS_TO_RESERVE_HOLDOUT) return "not_reserved_min_events";
  return splitHoldoutEvents(events).some((item) => item.usage === "holdout")
    ? "reserved"
    : "not_reserved_min_events";
}

/**
 * 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 < MIN_EVENTS_TO_RESERVE_HOLDOUT) {
    return events.map((item) => ({
      ...item,
      usage: dated.some((row) => row.id === item.id) ? "training" as const : "unused" 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;
  }
  return (monthOrBetter.at(-1) ?? dated.at(-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));
}

export function holdoutDomainYears(events: readonly InferenceEvent[]): ReadonlySet<string> {
  return new Set(
    events
      .filter((item) => item.usage === "holdout" && item.year !== null)
      .map((item) => `${item.domain}:${item.year}`),
  );
}

/** Keep a previously reserved holdout sticky as more evidence arrives. */
export function stickyHoldoutEvents(
  events: readonly DatedEventInput[],
  previous?: readonly InferenceEvent[] | null,
): InferenceEvent[] {
  const previousHoldout = previous?.find((item) => item.usage === "holdout");
  if (previousHoldout && events.some((item) => item.id === previousHoldout.id)) {
    const dated = events.filter((item) => item.year !== null && item.precision !== "unknown");
    if (dated.length >= MIN_EVENTS_TO_RESERVE_HOLDOUT) {
      return events.map((item) => ({
        ...item,
        usage: item.id === previousHoldout.id
          ? "holdout" as const
          : dated.some((row) => row.id === item.id)
            ? "training" as const
            : "unused" as const,
      }));
    }
  }
  return splitHoldoutEvents(events);
}