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

import { ELECTIVE_SETS } from '../data/electiveSets';
import { coursesBySet } from '../data/lookups';
import type { OptimizationMode } from '../data/types';
import { maximizeCount, priorityOrder } from './optimizer';

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

export interface SetAnalysis {
  setId: string;
  setName: string;
  impact: number; // variance in ceiling outcomes
  choices: ChoiceOutcome[];
}

const MAX_OPEN_SETS_FOR_ENUMERATION = 9;

/**
 * Compute the ceiling outcome for a single course choice:
 * the best achievable result assuming that course is pinned
 * and all other open sets are chosen optimally.
 */
function computeCeiling(
  basePinnedCourses: string[],
  chosenCourseId: string,
  otherOpenSetIds: string[],
  ranking: string[],
  mode: OptimizationMode,
  excludedCourseIds?: Set<string>,
): { count: number; specs: string[] } {
  const fn = mode === 'maximize-count' ? maximizeCount : priorityOrder;

  if (otherOpenSetIds.length === 0) {
    // No other open sets — just solve with this choice added
    const selected = [...basePinnedCourses, chosenCourseId];
    const result = fn(selected, ranking, [], excludedCourseIds);
    return { count: result.achieved.length, specs: result.achieved };
  }

  // Enumerate all combinations of remaining open sets
  let bestCount = 0;
  let bestSpecs: string[] = [];

  function enumerate(setIndex: number, accumulated: string[]) {
    // Early termination: already found max (3)
    if (bestCount >= 3) return;

    if (setIndex >= otherOpenSetIds.length) {
      const selected = [...basePinnedCourses, chosenCourseId, ...accumulated];
      const result = fn(selected, ranking, [], excludedCourseIds);
      if (result.achieved.length > bestCount) {
        bestCount = result.achieved.length;
        bestSpecs = result.achieved;
      }
      return;
    }

    const setId = otherOpenSetIds[setIndex];
    const courses = coursesBySet[setId];
    for (const course of courses) {
      if (excludedCourseIds?.has(course.id)) continue;
      enumerate(setIndex + 1, [...accumulated, course.id]);
      if (bestCount >= 3) return;
    }
  }

  enumerate(0, []);
  return { count: bestCount, specs: bestSpecs };
}

/**
 * Compute variance of an array of numbers.
 */
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;
}

/**
 * Analyze all open sets and compute per-choice ceiling outcomes.
 * Returns sets ordered by decision impact (highest first).
 *
 * onSetComplete is called progressively as each set's analysis finishes.
 */
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) {
    // Fallback: return empty analyses (caller uses upper bounds instead)
    return openSetIds.map((setId) => {
      const set = ELECTIVE_SETS.find((s) => s.id === setId)!;
      return { setId, setName: set.name, impact: 0, choices: [] };
    });
  }

  const analyses: SetAnalysis[] = [];

  for (const setId of openSetIds) {
    const set = ELECTIVE_SETS.find((s) => s.id === setId)!;
    const otherOpenSets = openSetIds.filter((id) => id !== setId);
    const courses = coursesBySet[setId];

    const choices: ChoiceOutcome[] = courses
      .filter((course) => !excludedCourseIds?.has(course.id))
      .map((course) => {
        const ceiling = computeCeiling(
          pinnedCourseIds,
          course.id,
          otherOpenSets,
          ranking,
          mode,
          excludedCourseIds,
        );
        return {
          courseId: course.id,
          courseName: course.name,
          ceilingCount: ceiling.count,
          ceilingSpecs: ceiling.specs,
        };
      });

    const impact = variance(choices.map((c) => c.ceilingCount));
    const analysis: SetAnalysis = { setId, setName: set.name, impact, choices };
    analyses.push(analysis);

    onSetComplete?.(analysis);
  }

  // Sort by impact descending, then by set order (chronological) for ties
  const setOrder = new Map(ELECTIVE_SETS.map((s, i) => [s.id, i]));
  analyses.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 analyses;
}