export interface HendrixNote {
  onset: number;
  offset: number;
  pitch: number;
  program: number;
  isDrum: boolean;
  velocity: number;
}

interface Event {
  type: "PAD" | "EOS" | "UNK" | "shift" | "pitch" | "velocity" | "tie" | "program" | "drum";
  value: number;
}

interface NoteEvent {
  isDrum: boolean;
  program: number;
  time: number | null;
  velocity: number;
  pitch: number;
}

interface DecodedWindow {
  notes: NoteEvent[];
  ties: NoteEvent[];
  lastActivity: Array<[number, number]>;
  start: number;
}

const PROGRAM_GROUPS: Array<[string, number[]]> = [
  ["acoustic_piano", [0, 1, 3, 6, 7]],
  ["electric_piano", [2, 4, 5]],
  ["chromatic_percussion", range(8, 16)],
  ["organ", range(16, 24)],
  ["acoustic_guitar", [24, 25]],
  ["clean_electric_guitar", [26, 27, 28]],
  ["distorted_electric_guitar", [29, 30, 31]],
  ["acoustic_bass", [32, 35]],
  ["electric_bass", [33, 34, 36, 37, 38, 39]],
  ["violin", [40]], ["viola", [41]], ["cello", [42]], ["contrabass", [43]],
  ["orchestral_harp", [46]], ["timpani", [47]],
  ["string_ensemble", [48, 49, 44, 45]], ["synth_strings", [50, 51]],
  ["voice", [52, 53, 54, 100, 101]], ["orchestra_hit", [55]],
  ["trumpet", [56, 59]], ["trombone", [57]], ["tuba", [58]],
  ["french_horn", [60]], ["brass_section", [61, 62, 63]],
  ["soprano_and_alto_sax", [64, 65]], ["tenor_sax", [66]],
  ["baritone_sax", [67]], ["oboe", [68]], ["english_horn", [69]],
  ["bassoon", [70]], ["clarinet", [71]], ["flutes", range(72, 80)],
  ["synth_lead", range(80, 88)], ["synth_pad", range(88, 96)],
];

function range(start: number, end: number): number[] {
  return Array.from({ length: end - start }, (_, index) => start + index);
}

const programNames = new Map<number, string>();
for (const [name, programs] of PROGRAM_GROUPS) {
  for (const program of programs) programNames.set(program, name);
}

export function instrumentName(note: HendrixNote): string {
  if (note.isDrum) return "drums";
  return programNames.get(note.program) ?? `program_${note.program}`;
}

function decodeToken(token: number): Event {
  if (token === 0) return { type: "PAD", value: 0 };
  if (token === 1) return { type: "EOS", value: 0 };
  if (token === 2) return { type: "UNK", value: 0 };
  if (token <= 208) return { type: "shift", value: token - 3 };
  if (token <= 336) return { type: "pitch", value: token - 209 };
  if (token <= 338) return { type: "velocity", value: token - 337 };
  if (token === 339) return { type: "tie", value: 0 };
  if (token <= 467) return { type: "program", value: token - 340 };
  if (token <= 595) return { type: "drum", value: token - 468 };
  throw new Error(`unknown token ${token}`);
}

function decodeWindow(tokens: number[], start: number): DecodedWindow {
  const events = tokens.map(decodeToken);
  let program: number | null = null;
  const ties: NoteEvent[] = [];
  const lastActivity: Array<[number, number]> = [];
  let tieIndex = -1;
  for (let i = 0; i < events.length; i++) {
    const event = events[i];
    if (event.type === "tie") {
      tieIndex = i;
      break;
    }
    if (event.type === "shift") break;
    if (event.type === "program") program = event.value;
    if (event.type === "pitch" && program !== null) {
      ties.push({ isDrum: false, program, time: null, velocity: 1, pitch: event.value });
      lastActivity.push([program, event.value]);
    }
  }
  if (tieIndex < 0) return { notes: [], ties: [], lastActivity: [], start };

  const notes: NoteEvent[] = [];
  const startTick = Math.round(start * 100);
  let tick = startTick;
  let velocity: number | null = null;
  for (const event of events.slice(tieIndex + 1)) {
    if (event.type === "shift") {
      if (event.value > 0) tick = startTick + event.value;
    } else if (event.type === "drum") {
      notes.push({ isDrum: true, program: 128, time: tick / 100, velocity: 1, pitch: event.value });
    } else if (event.type === "program") {
      program = event.value;
    } else if (event.type === "velocity") {
      velocity = event.value;
    } else if (event.type === "pitch" && program !== null && velocity !== null) {
      const key = `${program}:${event.value}`;
      if (velocity > 0) {
        lastActivity.push([program, event.value]);
      } else {
        const index = lastActivity.findIndex(([p, pitch]) => `${p}:${pitch}` === key);
        if (index < 0) continue;
        lastActivity.splice(index, 1);
      }
      notes.push({ isDrum: false, program, time: tick / 100, velocity, pitch: event.value });
    } else if (event.type === "EOS") {
      break;
    }
  }
  notes.sort((a, b) => (a.time ?? 0) - (b.time ?? 0) || Number(a.isDrum) - Number(b.isDrum) || a.program - b.program || a.velocity - b.velocity || a.pitch - b.pitch);
  ties.sort((a, b) => Number(a.isDrum) - Number(b.isDrum) || a.program - b.program || a.pitch - b.pitch);
  return { notes, ties, lastActivity, start };
}

function notesFromEvents(events: NoteEvent[]): HendrixNote[] {
  const active = new Map<string, NoteEvent>();
  const notes: HendrixNote[] = [];
  events.sort((a, b) => (a.time ?? 0) - (b.time ?? 0) || Number(a.isDrum) - Number(b.isDrum) || a.pitch - b.pitch || a.velocity - b.velocity || a.program - b.program);
  for (const event of events) {
    if (event.time === null) continue;
    if (event.isDrum) {
      notes.push({ onset: event.time, offset: event.time + 0.01, pitch: event.pitch, program: 128, isDrum: true, velocity: 1 });
      continue;
    }
    const key = `${event.pitch}:${event.program}`;
    if (event.velocity === 1) {
      const previous = active.get(key);
      if (previous !== undefined && previous.time !== null) {
        notes.push({ onset: previous.time, offset: event.time, pitch: event.pitch, program: event.program, isDrum: false, velocity: previous.velocity });
      }
      active.set(key, event);
    } else {
      const previous = active.get(key);
      active.delete(key);
      if (previous !== undefined && previous.time !== null) {
        notes.push({ onset: previous.time, offset: event.time, pitch: event.pitch, program: event.program, isDrum: false, velocity: previous.velocity });
      }
    }
  }
  for (const event of active.values()) {
    if (event.time !== null) notes.push({ onset: event.time, offset: event.time + 0.01, pitch: event.pitch, program: event.program, isDrum: false, velocity: 1 });
  }
  return notes;
}

function trimOverlaps(notes: HendrixNote[]): HendrixNote[] {
  const groups = new Map<string, HendrixNote[]>();
  for (const note of notes) {
    const key = `${note.pitch}:${note.program}:${Number(note.isDrum)}`;
    const group = groups.get(key) ?? [];
    group.push(note);
    groups.set(key, group);
  }
  const result: HendrixNote[] = [];
  for (const group of groups.values()) {
    group.sort((a, b) => a.onset - b.onset);
    for (let i = 1; i < group.length; i++) {
      if (group[i - 1].offset > group[i].onset) group[i - 1].offset = group[i].onset;
    }
    result.push(...group.filter((note) => note.onset < note.offset));
  }
  return result.sort((a, b) => a.onset - b.onset || Number(a.isDrum) - Number(b.isDrum) || a.program - b.program || a.velocity - b.velocity || a.pitch - b.pitch);
}

/** Decode [window][channel][token] output with the upstream tie merge rules. */
export function decodeTrackTokens(
  windows: number[][][],
  segmentSeconds: number,
  duration: number,
): HendrixNote[] {
  const channelNotes: HendrixNote[][] = [];
  for (let channel = 0; channel < 13; channel++) {
    const decoded = windows.map((window, index) => decodeWindow(window[channel], index * segmentSeconds));
    const merged: NoteEvent[] = [];
    let previousActivity: Array<[number, number]> | null = null;
    for (const current of decoded) {
      if (previousActivity !== null) {
        const ties = new Set(current.ties.map((event) => `${event.program}:${event.pitch}`));
        for (const [program, pitch] of previousActivity) {
          if (!ties.has(`${program}:${pitch}`)) {
            merged.push({ isDrum: false, program, time: current.start, velocity: 0, pitch });
          }
        }
      }
      merged.push(...current.notes);
      previousActivity = current.lastActivity;
    }
    channelNotes.push(notesFromEvents(merged));
  }
  return trimOverlaps(channelNotes.flat())
    .filter((note) => note.onset < duration)
    .map((note) => ({ ...note, offset: Math.min(duration, Math.max(note.offset, note.onset + 0.01)) }));
}