shiftcraft/server/utils/solver/modelBuilder.ts

import type { SolveInput } from '~/shared/types/schedule'
import { getDaysInRange, getWeekday } from '~/app/utils/dateHelpers'
import type { ConstraintJSON } from '~/shared/types/constraint'

export function buildModel(input: SolveInput) {
  const { employees, framework, period_start, period_end, constraints } = input
  const days = getDaysInRange(period_start, period_end)
  const E = employees.length
  const S = framework.shifts.length
  const D = days.length

  // Variable naming: x_e_s_d = 1 if employee e works shift s on day d
  const varNames: string[] = []
  const varIndex: Record<string, number> = {}

  for (let e = 0; e < E; e++) {
    for (let s = 0; s < S; s++) {
      for (let d = 0; d < D; d++) {
        const name = `x_${e}_${s}_${d}`
        varIndex[name] = varNames.length
        varNames.push(name)
      }
    }
  }

  // Penalty variables for soft constraints
  const penaltyVarStart = varNames.length
  let penaltyVarCount = 0
  const softConstraintPenalties: Array<{ varIdx: number; weight: number }> = []

  function addPenaltyVar(weight: number): number {
    const idx = penaltyVarStart + penaltyVarCount++
    varNames.push(`penalty_${idx}`)
    softConstraintPenalties.push({ varIdx: idx, weight })
    return idx
  }

  // Objective: minimize sum of soft penalties
  const objectiveCoeffs: Record<string, number> = {}
  for (const p of softConstraintPenalties) {
    objectiveCoeffs[varNames[p.varIdx]] = p.weight
  }

  const bounds: Record<string, [number, number]> = {}
  const binaries = new Set<string>()

  for (const name of varNames) {
    bounds[name] = [0, 1]
    if (name.startsWith('x_')) binaries.add(name)
  }

  const rows: Array<{
    name: string
    vars: Record<string, number>
    lb: number
    ub: number
  }> = []

  let rowIdx = 0
  function addRow(vars: Record<string, number>, lb: number, ub: number) {
    rows.push({ name: `r${rowIdx++}`, vars, lb, ub })
  }

  // Structural constraint 1: at most 1 shift per employee per day
  for (let e = 0; e < E; e++) {
    for (let d = 0; d < D; d++) {
      const vars: Record<string, number> = {}
      for (let s = 0; s < S; s++) {
        vars[`x_${e}_${s}_${d}`] = 1
      }
      addRow(vars, 0, 1)
    }
  }

  // Structural constraint 2: coverage requirements
  for (let s = 0; s < S; s++) {
    const shift = framework.shifts[s]
    for (let d = 0; d < D; d++) {
      const weekday = getWeekday(days[d])
      const applicable = shift.days_applicable.length === 0 || shift.days_applicable.includes(weekday)
      if (!applicable) {
        // Force all x_e_s_d = 0 for this day
        for (let e = 0; e < E; e++) {
          addRow({ [`x_${e}_${s}_${d}`]: 1 }, 0, 0)
        }
        continue
      }

      const coverageVars: Record<string, number> = {}
      for (let e = 0; e < E; e++) {
        coverageVars[`x_${e}_${s}_${d}`] = 1
      }
      addRow(coverageVars, shift.workers_required, E)
    }
  }

  // Structural constraint 3: availability
  for (let e = 0; e < E; e++) {
    const emp = employees[e]

    for (let d = 0; d < D; d++) {
      const date = days[d]

      // Block unavailable dates
      if (emp.unavailable_dates.includes(date)) {
        for (let s = 0; s < S; s++) {
          addRow({ [`x_${e}_${s}_${d}`]: 1 }, 0, 0)
        }
        continue
      }

      // Block periods not available
      for (let s = 0; s < S; s++) {
        const shift = framework.shifts[s]
        const period = framework.periods.find(p => p.id === shift.period_id)
        if (period && emp.available_periods.length > 0 && !emp.available_periods.includes(period.id)) {
          addRow({ [`x_${e}_${s}_${d}`]: 1 }, 0, 0)
        }
      }
    }
  }

  // Process constraints from ConstraintJSON
  for (const c of constraints) {
    const cj = c.constraint_json as ConstraintJSON
    const params = cj.params as Record<string, unknown>

    // Determine affected employees
    let affectedEmployees: number[] = Array.from({ length: E }, (_, i) => i)
    if (cj.scope.type === 'employee') {
      const empIdx = employees.findIndex(e => e.id === (cj.scope as { type: 'employee'; employee_id: string }).employee_id)
      if (empIdx >= 0) affectedEmployees = [empIdx]
      else continue
    } else if (cj.scope.type === 'role') {
      const role = (cj.scope as { type: 'role'; role: string }).role
      affectedEmployees = employees
        .map((e, i) => ({ e, i }))
        .filter(({ e }) => e.roles.includes(role))
        .map(({ i }) => i)
    }

    switch (cj.type) {
      case 'max_consecutive_shifts': {
        const max = params.max_count as number
        for (const e of affectedEmployees) {
          for (let d = 0; d <= D - (max + 1); d++) {
            const vars: Record<string, number> = {}
            for (let dd = d; dd <= d + max; dd++) {
              for (let s = 0; s < S; s++) {
                vars[`x_${e}_${s}_${dd}`] = 1
              }
            }
            addRow(vars, 0, max)
          }
        }
        break
      }

      case 'max_consecutive_shift_type': {
        const max = params.max_count as number
        const periodId = params.period_id as string
        const shiftIndices = framework.shifts
          .map((sh, i) => ({ sh, i }))
          .filter(({ sh }) => sh.period_id === periodId)
          .map(({ i }) => i)

        if (shiftIndices.length === 0) break

        for (const e of affectedEmployees) {
          for (let d = 0; d <= D - (max + 1); d++) {
            const vars: Record<string, number> = {}
            for (let dd = d; dd <= d + max; dd++) {
              for (const s of shiftIndices) {
                vars[`x_${e}_${s}_${dd}`] = 1
              }
            }
            addRow(vars, 0, max)
          }
        }
        break
      }

      case 'forbidden_shift_sequence': {
        const firstPeriod = params.first_period_id as string
        const secondPeriod = params.second_period_id as string
        const firstShifts = framework.shifts
          .map((sh, i) => ({ sh, i }))
          .filter(({ sh }) => sh.period_id === firstPeriod)
          .map(({ i }) => i)
        const secondShifts = framework.shifts
          .map((sh, i) => ({ sh, i }))
          .filter(({ sh }) => sh.period_id === secondPeriod)
          .map(({ i }) => i)

        for (const e of affectedEmployees) {
          for (let d = 0; d < D - 1; d++) {
            for (const s1 of firstShifts) {
              for (const s2 of secondShifts) {
                addRow({
                  [`x_${e}_${s1}_${d}`]: 1,
                  [`x_${e}_${s2}_${d + 1}`]: 1
                }, 0, 1)
              }
            }
          }
        }
        break
      }

      case 'employee_avoids_period': {
        const periodId = params.period_id as string
        const shiftIndices = framework.shifts
          .map((sh, i) => ({ sh, i }))
          .filter(({ sh }) => sh.period_id === periodId)
          .map(({ i }) => i)

        if (cj.hard) {
          for (const e of affectedEmployees) {
            for (let d = 0; d < D; d++) {
              for (const s of shiftIndices) {
                addRow({ [`x_${e}_${s}_${d}`]: 1 }, 0, 0)
              }
            }
          }
        } else {
          // Soft: penalty per assignment
          const weight = (c.weight || 70) * 10
          for (const e of affectedEmployees) {
            for (let d = 0; d < D; d++) {
              for (const s of shiftIndices) {
                const pv = addPenaltyVar(weight)
                objectiveCoeffs[varNames[pv]] = weight
                bounds[varNames[pv]] = [0, 1]
                addRow({
                  [varNames[pv]]: 1,
                  [`x_${e}_${s}_${d}`]: -1
                }, 0, 1)
              }
            }
          }
        }
        break
      }

      case 'max_weekend_shifts_per_month': {
        const max = params.max_count as number
        const weekendDays = days
          .map((day, i) => ({ day, i }))
          .filter(({ day }) => {
            const wd = getWeekday(day)
            return wd === 5 || wd === 6
          })
          .map(({ i }) => i)

        for (const e of affectedEmployees) {
          const vars: Record<string, number> = {}
          for (const d of weekendDays) {
            for (let s = 0; s < S; s++) {
              vars[`x_${e}_${s}_${d}`] = 1
            }
          }
          if (Object.keys(vars).length > 0) {
            addRow(vars, 0, max)
          }
        }
        break
      }

      case 'min_rest_between_shifts':
        // Already handled implicitly by "max 1 shift per day"
        break

      default:
        break
    }
  }

  // Rebuild objective with finalized penalty vars
  for (const p of softConstraintPenalties) {
    objectiveCoeffs[varNames[p.varIdx]] = p.weight
  }

  // Initialize bounds for penalty vars
  for (let i = penaltyVarStart; i < varNames.length; i++) {
    if (!bounds[varNames[i]]) {
      bounds[varNames[i]] = [0, 1]
    }
  }

  return {
    name: 'ShiftCraft',
    sense: 'minimize' as const,
    offset: 0,
    col_lower: varNames.map(n => bounds[n]?.[0] ?? 0),
    col_upper: varNames.map(n => bounds[n]?.[1] ?? 1),
    col_cost: varNames.map(n => objectiveCoeffs[n] ?? 0),
    col_names: varNames,
    row_lower: rows.map(r => r.lb),
    row_upper: rows.map(r => r.ub),
    row_names: rows.map(r => r.name),
    a_matrix: {
      format: 'colwise' as const,
      start: buildColwiseStart(varNames, rows),
      index: buildColwiseIndex(varNames, rows),
      value: buildColwiseValues(varNames, rows),
    },
    integrality: varNames.map(n => binaries.has(n) ? 1 : 0),
  }
}

function buildColwiseStart(varNames: string[], rows: Array<{ vars: Record<string, number> }>): number[] {
  const starts: number[] = [0]
  let count = 0
  for (const name of varNames) {
    for (const row of rows) {
      if (name in row.vars) count++
    }
    starts.push(count)
  }
  return starts
}

function buildColwiseIndex(varNames: string[], rows: Array<{ vars: Record<string, number> }>): number[] {
  const indices: number[] = []
  for (const name of varNames) {
    rows.forEach((row, i) => {
      if (name in row.vars) indices.push(i)
    })
  }
  return indices
}

function buildColwiseValues(varNames: string[], rows: Array<{ vars: Record<string, number> }>): number[] {
  const values: number[] = []
  for (const name of varNames) {
    for (const row of rows) {
      if (name in row.vars) values.push(row.vars[name])
    }
  }
  return values
}