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

import { SPECIALIZATIONS } from '../data/specializations';
import { COURSES } from '../data/courses';
import { coursesBySpec, setIdByCourse } from '../data/lookups';
import type { AllocationResult, SpecStatus } from '../data/types';
import {
  checkFeasibility,
  enumerateS2Choices,
  preFilterCandidates,
  computeUpperBounds,
} from './feasibility';

const CREDIT_THRESHOLD = 9;
const CREDIT_PER_COURSE = 2.5;

/**
 * Generate all combinations of k items from arr.
 */
function combinations<T>(arr: T[], k: number): T[][] {
  if (k === 0) return [[]];
  if (k > arr.length) return [];
  const result: T[][] = [];
  for (let i = 0; i <= arr.length - k; i++) {
    for (const rest of combinations(arr.slice(i + 1), k - 1)) {
      result.push([arr[i], ...rest]);
    }
  }
  return result;
}

/**
 * Try to check feasibility for a target set, handling S2 enumeration for Strategy.
 */
function checkWithS2(
  selectedCourseIds: string[],
  targetSpecIds: string[],
): { feasible: boolean; allocations: Record<string, Record<string, number>> } {
  const hasStrategy = targetSpecIds.includes('STR');
  if (!hasStrategy) {
    return checkFeasibility(selectedCourseIds, targetSpecIds);
  }

  // Enumerate S2 choices
  const s2Choices = enumerateS2Choices(selectedCourseIds);
  for (const s2Choice of s2Choices) {
    const result = checkFeasibility(selectedCourseIds, targetSpecIds, s2Choice);
    if (result.feasible) return result;
  }
  return { feasible: false, allocations: {} };
}

/**
 * Maximize Count mode: find the largest set of achievable specializations.
 * Among equal-size feasible subsets, prefer the one with the highest priority score.
 */
export function maximizeCount(
  selectedCourseIds: string[],
  ranking: string[],
  openSetIds: string[],
  excludedCourseIds?: Set<string>,
): { achieved: string[]; allocations: Record<string, Record<string, number>> } {
  const candidates = preFilterCandidates(selectedCourseIds, openSetIds, excludedCourseIds);

  // Only check specs that can be achieved from selected courses alone (not open sets)
  // Filter to candidates that have qualifying selected courses
  const achievable = candidates.filter((specId) => {
    const entries = coursesBySpec[specId] || [];
    return entries.some((e) => selectedCourseIds.includes(e.courseId));
  });

  // Priority score: sum of (15 - rank position) for each spec in subset
  const rankIndex = new Map(ranking.map((id, i) => [id, i]));
  function priorityScore(specs: string[]): number {
    return specs.reduce((sum, id) => sum + (15 - (rankIndex.get(id) ?? 14)), 0);
  }

  // Try from size 3 down to 0
  const maxSize = Math.min(3, achievable.length);
  for (let size = maxSize; size >= 1; size--) {
    const subsets = combinations(achievable, size);

    // Sort subsets by priority score descending — check best-scoring first
    subsets.sort((a, b) => priorityScore(b) - priorityScore(a));

    let bestResult: { achieved: string[]; allocations: Record<string, Record<string, number>> } | null = null;
    let bestScore = -1;

    for (const subset of subsets) {
      const result = checkWithS2(selectedCourseIds, subset);
      if (result.feasible) {
        const score = priorityScore(subset);
        if (score > bestScore) {
          bestScore = score;
          bestResult = { achieved: subset, allocations: result.allocations };
        }
      }
    }

    if (bestResult) return bestResult;
  }

  return { achieved: [], allocations: {} };
}

/**
 * Priority Order mode: process specializations in rank order,
 * adding each if feasible with the current achieved set.
 */
export function priorityOrder(
  selectedCourseIds: string[],
  ranking: string[],
  openSetIds: string[],
  excludedCourseIds?: Set<string>,
): { achieved: string[]; allocations: Record<string, Record<string, number>> } {
  const candidates = new Set(preFilterCandidates(selectedCourseIds, openSetIds, excludedCourseIds));

  // Only consider specs that have qualifying selected courses
  const withSelectedCourses = new Set(
    SPECIALIZATIONS.filter((spec) => {
      const entries = coursesBySpec[spec.id] || [];
      return entries.some((e) => selectedCourseIds.includes(e.courseId));
    }).map((s) => s.id),
  );

  const achieved: string[] = [];
  let lastAllocations: Record<string, Record<string, number>> = {};

  for (const specId of ranking) {
    if (!candidates.has(specId)) continue;
    if (!withSelectedCourses.has(specId)) continue;
    if (achieved.length >= 3) break;

    const trySet = [...achieved, specId];
    const result = checkWithS2(selectedCourseIds, trySet);
    if (result.feasible) {
      achieved.push(specId);
      lastAllocations = result.allocations;
    }
  }

  return { achieved, allocations: lastAllocations };
}

/**
 * Determine the status of each specialization after optimization.
 */
export function determineStatuses(
  selectedCourseIds: string[],
  openSetIds: string[],
  achieved: string[],
  excludedCourseIds?: Set<string>,
): Record<string, SpecStatus> {
  const achievedSet = new Set(achieved);
  const selectedSet = new Set(selectedCourseIds);
  const openSetSet = new Set(openSetIds);
  const upperBounds = computeUpperBounds(selectedCourseIds, openSetIds, excludedCourseIds);
  const statuses: Record<string, SpecStatus> = {};

  for (const spec of SPECIALIZATIONS) {
    if (achievedSet.has(spec.id)) {
      statuses[spec.id] = 'achieved';
      continue;
    }

    // Check required course gate
    if (spec.requiredCourseId) {
      if (!selectedSet.has(spec.requiredCourseId)) {
        const requiredCourse = COURSES.find((c) => c.id === spec.requiredCourseId)!;
        if (!openSetSet.has(requiredCourse.setId)) {
          statuses[spec.id] = 'missing_required';
          continue;
        }
      }
    }

    // Check upper bound
    if (upperBounds[spec.id] < CREDIT_THRESHOLD) {
      statuses[spec.id] = 'unreachable';
      continue;
    }

    statuses[spec.id] = 'achievable';
  }

  return statuses;
}

/**
 * Run the full optimization pipeline.
 */
export function optimize(
  selectedCourseIds: string[],
  ranking: string[],
  openSetIds: string[],
  mode: 'maximize-count' | 'priority-order',
  excludedCourseIds?: Set<string>,
): AllocationResult {
  const fn = mode === 'maximize-count' ? maximizeCount : priorityOrder;
  const { achieved, allocations } = fn(selectedCourseIds, ranking, openSetIds, excludedCourseIds);
  const statuses = determineStatuses(selectedCourseIds, openSetIds, achieved, excludedCourseIds);
  const upperBounds = computeUpperBounds(selectedCourseIds, openSetIds, excludedCourseIds);

  return { achieved, allocations, statuses, upperBounds };
}