emba-course-solver/app/src/solver/decisionTree.ts

import { ELECTIVE_SETS } from '../data/electiveSets';
import { coursesBySet } from '../data/lookups';
import type { Course, OptimizationMode } from '../data/types';
import { maximizeCount, priorityOrder } from './optimizer';
import { computeUpperBounds } from './feasibility';
import { makePriorityScorer, makePriorityRankWeight } from './priority';

export interface ChoiceOutcome {
  courseId: string;
  courseName: string;
  ceilingCount: number;
  ceilingSpecs: string[];
  evaluated: boolean;
}

export interface SetAnalysis {
  setId: string;
  setName: string;
  impact: number;
  choices: ChoiceOutcome[];
}

export interface PlanOutcome {
  courseAssignments: Record<string, string>; // setId -> courseId for open sets
  achievedSpecs: string[];
  priorityScore: number;
}

export interface SearchResult {
  topK: PlanOutcome[];
  setAnalyses: SetAnalysis[];
  partial: boolean;
  iterations: number;
  iterationsTotal: number;
}

export interface SearchCallbacks {
  onTopKUpdate?: (topK: PlanOutcome[], iterations: number) => void;
  onChoiceUpdate?: (setId: string, analysis: SetAnalysis) => void;
  onProgress?: (iterations: number, iterationsTotal: number) => void;
  onLeafEvaluated?: (leaf: PlanOutcome) => void;
}

const MAX_OPEN_SETS_FOR_ENUMERATION = 9;
const CREDIT_THRESHOLD = 9;
export const PROGRESS_THROTTLE_MS = 100;

function variance(values: number[]): number {
  if (values.length <= 1) return 0;
  const mean = values.reduce((a, b) => a + b, 0) / values.length;
  return values.reduce((sum, v) => sum + (v - mean) ** 2, 0) / values.length;
}

export function selectPriorityTarget(
  ranking: string[],
  upperBounds: Record<string, number>,
): string | null {
  for (const specId of ranking) {
    if ((upperBounds[specId] ?? 0) >= CREDIT_THRESHOLD) return specId;
  }
  return null;
}

export function reorderForTarget(
  setId: string,
  target: string | null,
  excludedCourseIds?: Set<string>,
): Course[] {
  const courses = coursesBySet[setId].filter(
    (c) => !excludedCourseIds?.has(c.id),
  );
  if (!target) return courses;
  const qualifying: Course[] = [];
  const others: Course[] = [];
  for (const c of courses) {
    if (c.qualifications.some((q) => q.specId === target)) qualifying.push(c);
    else others.push(c);
  }
  return [...qualifying, ...others];
}

/**
 * Reorder a set's courses so those qualifying for the most reachable specs
 * (upperBound >= 9) come first. Stable sort: ties keep declaration order.
 * Used by maximize-count mode to surface generalist courses early.
 */
export function reorderByReachableQualCount(
  setId: string,
  upperBounds: Record<string, number>,
  excludedCourseIds?: Set<string>,
): Course[] {
  const courses = coursesBySet[setId].filter(
    (c) => !excludedCourseIds?.has(c.id),
  );
  // Decorate-sort-undecorate for stability
  return courses
    .map((course, idx) => ({
      course,
      idx,
      score: course.qualifications.filter(
        (q) => (upperBounds[q.specId] ?? 0) >= CREDIT_THRESHOLD,
      ).length,
    }))
    .sort((a, b) => {
      if (b.score !== a.score) return b.score - a.score;
      return a.idx - b.idx;
    })
    .map((x) => x.course);
}

export function assignmentKey(assignments: Record<string, string>): string {
  return Object.keys(assignments)
    .sort()
    .map((k) => `${k}:${assignments[k]}`)
    .join('|');
}

export class BoundedRankedList<T> {
  private items: T[] = [];
  constructor(
    private capacity: number,
    private compare: (a: T, b: T) => number,
  ) {}

  tryInsert(item: T): boolean {
    let pos = 0;
    while (pos < this.items.length && this.compare(item, this.items[pos]) > 0) pos++;
    if (pos >= this.capacity) return false;
    this.items.splice(pos, 0, item);
    if (this.items.length > this.capacity) this.items.pop();
    return true;
  }

  toArray(): T[] {
    return [...this.items];
  }
}

/**
 * Comparator for plan outcomes. Mode-dependent ordering uses lexicographic
 * rank-weight (top-ranked spec dominates any combination of lower-ranked
 * specs):
 *   - priority-order mode: (rankWeight desc, count desc, key asc)
 *   - maximize-count mode: (count desc, rankWeight desc, key asc)
 * Returns negative if a is better, positive if b is better.
 */
export function makeOutcomeComparator(
  mode: OptimizationMode,
  ranking: string[],
): (a: PlanOutcome, b: PlanOutcome) => number {
  const rankWeight = makePriorityRankWeight(ranking);
  return (a, b) => {
    const aw = rankWeight(a.achievedSpecs);
    const bw = rankWeight(b.achievedSpecs);
    if (mode === 'priority-order') {
      if (aw !== bw) return bw - aw;
      if (a.achievedSpecs.length !== b.achievedSpecs.length) return b.achievedSpecs.length - a.achievedSpecs.length;
    } else {
      if (a.achievedSpecs.length !== b.achievedSpecs.length) return b.achievedSpecs.length - a.achievedSpecs.length;
      if (aw !== bw) return bw - aw;
    }
    return assignmentKey(a.courseAssignments).localeCompare(
      assignmentKey(b.courseAssignments),
    );
  };
}

/** Default count-first comparator (uses default ranking), retained for backward compatibility with tests. */
export function compareOutcomes(a: PlanOutcome, b: PlanOutcome): number {
  // Default to alphabetical ranking; tests using this directly only exercise
  // simple count/key cases that are insensitive to ranking.
  return makeOutcomeComparator('maximize-count', [])(a, b);
}

interface CeilingComparable {
  count: number;
  specs: string[];
  key: string;
}

function makeCeilingComparator(
  mode: OptimizationMode,
  ranking: string[],
): (a: CeilingComparable, b: CeilingComparable) => number {
  const rankWeight = makePriorityRankWeight(ranking);
  return (a, b) => {
    const aw = rankWeight(a.specs);
    const bw = rankWeight(b.specs);
    if (mode === 'priority-order') {
      if (aw !== bw) return bw - aw;
      if (a.count !== b.count) return b.count - a.count;
    } else {
      if (a.count !== b.count) return b.count - a.count;
      if (aw !== bw) return bw - aw;
    }
    return a.key.localeCompare(b.key);
  };
}

export function searchDecisionTree(
  pinnedCourseIds: string[],
  openSetIds: string[],
  ranking: string[],
  mode: OptimizationMode,
  K: number,
  callbacks?: SearchCallbacks,
  excludedCourseIds?: Set<string>,
  skipKeys?: Set<string>,
  pinnedAssignments?: Record<string, string>,
  externalCredits?: Record<string, number>,
): SearchResult {
  const fn = mode === 'maximize-count' ? maximizeCount : priorityOrder;
  const scorer = makePriorityScorer(ranking);
  const upperBounds = computeUpperBounds(
    pinnedCourseIds,
    openSetIds,
    excludedCourseIds,
    externalCredits,
  );
  const priorityTarget = selectPriorityTarget(ranking, upperBounds);
  // Pinned assignments (setId -> courseId) for any pinned sets — included in
  // the leaf's full courseAssignments so cache keys are stable across pin/unpin.
  const pinnedMap = pinnedAssignments ?? {};

  // Initialize per-set analyses with unevaluated cells, ordered by mode
  const setAnalyses: Record<string, SetAnalysis> = {};
  const orderedCoursesPerSet: Record<string, Course[]> = {};
  let iterationsTotal = 1;
  for (const setId of openSetIds) {
    const set = ELECTIVE_SETS.find((s) => s.id === setId)!;
    const ordered =
      mode === 'maximize-count'
        ? reorderByReachableQualCount(setId, upperBounds, excludedCourseIds)
        : reorderForTarget(setId, priorityTarget, excludedCourseIds);
    orderedCoursesPerSet[setId] = ordered;
    iterationsTotal *= ordered.length || 1;
    setAnalyses[setId] = {
      setId,
      setName: set.name,
      impact: 0,
      choices: ordered.map((c) => ({
        courseId: c.id,
        courseName: c.name,
        ceilingCount: 0,
        ceilingSpecs: [],
        evaluated: false,
      })),
    };
  }
  const choiceKey: Record<string, string> = {};

  const outcomeComparator = makeOutcomeComparator(mode, ranking);
  const ceilingComparator = makeCeilingComparator(mode, ranking);
  const topK = new BoundedRankedList<PlanOutcome>(K, outcomeComparator);
  let iterations = 0;
  const partial = false;
  let lastProgressEmit = 0;

  function emitProgress() {
    if (!callbacks?.onProgress) return;
    const now = Date.now();
    if (now - lastProgressEmit >= PROGRESS_THROTTLE_MS) {
      lastProgressEmit = now;
      callbacks.onProgress(iterations, iterationsTotal);
    }
  }

  function evaluateLeaf(accumulated: Record<string, string>): void {
    iterations++;

    // Build the full 12-set assignment so cache keys remain stable across
    // pin/unpin operations.
    const fullAssignment: Record<string, string> = { ...pinnedMap, ...accumulated };
    const aKey = assignmentKey(fullAssignment);
    if (skipKeys?.has(aKey)) {
      emitProgress();
      return;
    }

    const courses: string[] = [];
    for (const setId of openSetIds) courses.push(accumulated[setId]);
    const selected = [...pinnedCourseIds, ...courses];
    const result = fn(selected, ranking, [], excludedCourseIds, externalCredits);
    const score = scorer(result.achieved);

    const outcome: PlanOutcome = {
      courseAssignments: fullAssignment,
      achievedSpecs: result.achieved,
      priorityScore: score,
    };

    callbacks?.onLeafEvaluated?.(outcome);

    if (topK.tryInsert(outcome)) {
      callbacks?.onTopKUpdate?.(topK.toArray(), iterations);
    }

    // Per-set ceiling + evaluated-flag updates
    for (const setId of openSetIds) {
      const courseId = accumulated[setId];
      const analysis = setAnalyses[setId];
      const choice = analysis.choices.find((c) => c.courseId === courseId)!;
      const wasEvaluated = choice.evaluated;
      const currentKey = `${setId}:${courseId}`;
      const existing: CeilingComparable = {
        count: choice.ceilingCount,
        specs: choice.ceilingSpecs,
        key: choiceKey[currentKey] ?? '',
      };
      const candidate: CeilingComparable = {
        count: result.achieved.length,
        specs: result.achieved,
        key: aKey,
      };
      const ceilingImproved = ceilingComparator(candidate, existing) < 0;
      if (ceilingImproved) {
        choice.ceilingCount = candidate.count;
        choice.ceilingSpecs = result.achieved;
        choiceKey[currentKey] = aKey;
      }
      // Mark evaluated regardless of improvement
      choice.evaluated = true;

      if (!wasEvaluated || ceilingImproved) {
        const impact = variance(analysis.choices.map((c) => c.ceilingCount));
        analysis.impact = impact;
        const updated: SetAnalysis = {
          ...analysis,
          impact,
          choices: analysis.choices.map((c) => ({ ...c })),
        };
        callbacks?.onChoiceUpdate?.(setId, updated);
      }
    }

    emitProgress();
  }

  function dfs(setIdx: number, accumulated: Record<string, string>) {
    if (setIdx >= openSetIds.length) {
      evaluateLeaf(accumulated);
      return;
    }
    const setId = openSetIds[setIdx];
    const courses = orderedCoursesPerSet[setId];
    for (const course of courses) {
      accumulated[setId] = course.id;
      dfs(setIdx + 1, accumulated);
    }
    delete accumulated[setId];
  }

  if (openSetIds.length > 0 && openSetIds.every((s) => orderedCoursesPerSet[s].length > 0)) {
    dfs(0, {});
  }

  // Final progress emit so consumers see the completion count
  if (callbacks?.onProgress) callbacks.onProgress(iterations, iterationsTotal);

  // Final impact recomputation + sort
  for (const a of Object.values(setAnalyses)) {
    a.impact = variance(a.choices.map((c) => c.ceilingCount));
  }
  const setOrder = new Map(ELECTIVE_SETS.map((s, i) => [s.id, i]));
  const sortedAnalyses = Object.values(setAnalyses).sort((a, b) => {
    if (b.impact !== a.impact) return b.impact - a.impact;
    return (setOrder.get(a.setId) ?? 0) - (setOrder.get(b.setId) ?? 0);
  });

  return {
    topK: topK.toArray(),
    setAnalyses: sortedAnalyses,
    partial,
    iterations,
    iterationsTotal,
  };
}

/**
 * Pure derivation of {topK, setAnalyses} from a collection of leaf outcomes.
 * Used by the main thread when filtering the leaf cache, and reusable
 * elsewhere as needed. Does NOT run any optimizer calls — leaves carry
 * their own pre-computed achievedSpecs/priorityScore.
 */
export function deriveFromLeaves(
  leaves: Iterable<PlanOutcome>,
  K: number,
  mode: OptimizationMode,
  ranking: string[],
  openSetIds: string[],
  excludedCourseIds?: Set<string>,
  externalCredits?: Record<string, number>,
): { topK: PlanOutcome[]; setAnalyses: SetAnalysis[] } {
  const upperBounds = computeUpperBounds([], openSetIds, excludedCourseIds, externalCredits);
  const priorityTarget = selectPriorityTarget(ranking, upperBounds);

  const setAnalyses: Record<string, SetAnalysis> = {};
  for (const setId of openSetIds) {
    const set = ELECTIVE_SETS.find((s) => s.id === setId)!;
    const ordered =
      mode === 'maximize-count'
        ? reorderByReachableQualCount(setId, upperBounds, excludedCourseIds)
        : reorderForTarget(setId, priorityTarget, excludedCourseIds);
    setAnalyses[setId] = {
      setId,
      setName: set.name,
      impact: 0,
      choices: ordered.map((c) => ({
        courseId: c.id,
        courseName: c.name,
        ceilingCount: 0,
        ceilingSpecs: [],
        evaluated: false,
      })),
    };
  }
  const choiceKey: Record<string, string> = {};
  const ceilingComparator = makeCeilingComparator(mode, ranking);
  const outcomeComparator = makeOutcomeComparator(mode, ranking);
  const topK = new BoundedRankedList<PlanOutcome>(K, outcomeComparator);

  for (const leaf of leaves) {
    topK.tryInsert(leaf);
    const aKey = assignmentKey(leaf.courseAssignments);
    for (const setId of openSetIds) {
      const courseId = leaf.courseAssignments[setId];
      if (!courseId) continue;
      const analysis = setAnalyses[setId];
      const choice = analysis.choices.find((c) => c.courseId === courseId);
      if (!choice) continue;
      const currentKey = `${setId}:${courseId}`;
      const existing: CeilingComparable = {
        count: choice.ceilingCount,
        specs: choice.ceilingSpecs,
        key: choiceKey[currentKey] ?? '',
      };
      const candidate: CeilingComparable = {
        count: leaf.achievedSpecs.length,
        specs: leaf.achievedSpecs,
        key: aKey,
      };
      if (ceilingComparator(candidate, existing) < 0) {
        choice.ceilingCount = candidate.count;
        choice.ceilingSpecs = leaf.achievedSpecs;
        choiceKey[currentKey] = aKey;
      }
      choice.evaluated = true;
    }
  }

  for (const a of Object.values(setAnalyses)) {
    a.impact = variance(a.choices.map((c) => c.ceilingCount));
  }
  const setOrder = new Map(ELECTIVE_SETS.map((s, i) => [s.id, i]));
  const sortedAnalyses = Object.values(setAnalyses).sort((a, b) => {
    if (b.impact !== a.impact) return b.impact - a.impact;
    return (setOrder.get(a.setId) ?? 0) - (setOrder.get(b.setId) ?? 0);
  });

  return { topK: topK.toArray(), setAnalyses: sortedAnalyses };
}

/**
 * Backward-compatible wrapper: produces only the per-set ceiling table.
 * Internally runs searchDecisionTree with K=10 and emits each set's analysis
 * once per choice update via the legacy onSetComplete callback.
 */
export function analyzeDecisionTree(
  pinnedCourseIds: string[],
  openSetIds: string[],
  ranking: string[],
  mode: OptimizationMode,
  onSetComplete?: (analysis: SetAnalysis) => void,
  excludedCourseIds?: Set<string>,
): SetAnalysis[] {
  if (openSetIds.length > MAX_OPEN_SETS_FOR_ENUMERATION) {
    return openSetIds.map((setId) => {
      const set = ELECTIVE_SETS.find((s) => s.id === setId)!;
      return { setId, setName: set.name, impact: 0, choices: [] };
    });
  }
  const result = searchDecisionTree(
    pinnedCourseIds,
    openSetIds,
    ranking,
    mode,
    10,
    onSetComplete
      ? { onChoiceUpdate: (_setId, analysis) => onSetComplete(analysis) }
      : undefined,
    excludedCourseIds,
  );
  return result.setAnalyses;
}