import { secondsToTicks } from "../core/time.ts";
import type { Grid, Role, ScoreNote, TimedNote, TruthTrack } from "../core/types.ts";
import { matchNotes, maxAssignment, prf, type MatchResult } from "./notes.ts";
import { truthFromGp } from "../score/gp.ts";
import { alignGrid, rhythmScore } from "./rhythm.ts";

export interface OracleFlags { oracleStems: boolean; oracleTempo: boolean; oracleNotes: boolean }
export interface OutTrack { name: string; role: Role; notes: TimedNote[]; scoreNotes: ScoreNote[] }
export interface PipelineOutput { grid: Grid; tracks: OutTrack[] }
export type DrumClass = "kick" | "snare" | "hihat" | "cymbal" | "tom";
export interface SongScore {
  flags: OracleFlags;
  guitarMerged: MatchResult & { pass: boolean };
  partSplit: { f1: number | null; outTracks: number; refTracks: number } | null;
  perRole: Partial<Record<Role, MatchResult>>;
  drumsPerClass: Record<DrumClass, MatchResult> | null;
  rhythm: { factor: number; beatOffset: number; onBeat: number | null; durationOk: number | null } | null;
  stringAgreement: number | null;
  spuriousRate: number | null;
  /** Merged-guitar pitch match computed from the re-imported output GP file (null without outGp). */
  gpPitch: MatchResult | null;
  /** Truth vs output-GP bar count and time signature (null without outGp). */
  barAgreement: { refBars: number; outBars: number; timeSigMatch: boolean } | null;
}

export const ROLE_ORDER: Role[] = ["guitar", "vocals", "bass", "drums", "keys"];
export const DRUM_CLASSES: Record<DrumClass, number[]> = {
  kick: [35, 36], snare: [37, 38, 40], hihat: [42, 44, 46],
  cymbal: [49, 51, 52, 53, 55, 57, 59], tom: [41, 43, 45, 47, 48, 50],
};
const PASS_F1 = 0.85;

export function drumClass(pitch: number): DrumClass | null {
  for (const [c, ps] of Object.entries(DRUM_CLASSES)) if (ps.includes(pitch)) return c as DrumClass;
  return null;
}
// Drum pitches are compared by class: map to the class's first GM number; unknown pitches stay distinct.
export const toClassPitch = (n: TimedNote): TimedNote => {
  const c = drumClass(n.pitch);
  return { ...n, pitch: c ? DRUM_CLASSES[c][0]! : 1000 + n.pitch };
};

const notesOf = (tracks: { role: Role; notes: TimedNote[] }[], role: Role) =>
  tracks.filter((t) => t.role === role).flatMap((t) => t.notes);

export function scoreSong(
  truth: { bpm: number; timeSig: [number, number]; bars?: number; tracks: TruthTrack[] },
  out: PipelineOutput, flags: OracleFlags, outGp?: Uint8Array,
): SongScore {
  const refGuitar = notesOf(truth.tracks, "guitar");
  const outGuitar = notesOf(out.tracks, "guitar");
  const gm = matchNotes(refGuitar, outGuitar);
  const guitarMerged = { ...gm, pass: gm.f1 !== null && gm.f1 >= PASS_F1 };

  // partSplit (R3): null only with <2 ref guitar tracks
  const refG = truth.tracks.filter((t) => t.role === "guitar");
  const outG = out.tracks.filter((t) => t.role === "guitar");
  let partSplit: SongScore["partSplit"] = null;
  if (refG.length >= 2) {
    const gain = refG.map((r) => outG.map((o) => matchNotes(r.notes, o.notes).tp));
    const asg = outG.length ? maxAssignment(gain) : refG.map(() => -1);
    let tp = 0;
    asg.forEach((j, i) => { if (j >= 0) tp += gain[i]![j]!; });
    const fn = refG.reduce((s, t) => s + t.notes.length, 0) - tp;
    const fp = outG.reduce((s, t) => s + t.notes.length, 0) - tp;
    partSplit = { f1: prf(tp, fp, fn).f1, outTracks: outG.length, refTracks: refG.length };
  }

  // per role (drums compared by class)
  const perRole: Partial<Record<Role, MatchResult>> = {};
  for (const role of ROLE_ORDER) {
    if (role === "guitar") { perRole.guitar = gm; continue; }
    let r = notesOf(truth.tracks, role), e = notesOf(out.tracks, role);
    if (role === "drums") { r = r.map(toClassPitch); e = e.map(toClassPitch); }
    perRole[role] = matchNotes(r, e);
  }

  let drumsPerClass: SongScore["drumsPerClass"] = null;
  const refDrums = notesOf(truth.tracks, "drums"), outDrums = notesOf(out.tracks, "drums");
  if (refDrums.length + outDrums.length > 0) {
    drumsPerClass = {} as Record<DrumClass, MatchResult>;
    for (const c of Object.keys(DRUM_CLASSES) as DrumClass[]) {
      drumsPerClass[c] = matchNotes(
        refDrums.filter((n) => drumClass(n.pitch) === c).map(toClassPitch),
        outDrums.filter((n) => drumClass(n.pitch) === c).map(toClassPitch),
      );
    }
  }

  // rhythm on merged guitar (R2: truth seconds -> ticks)
  let rhythm: SongScore["rhythm"] = null;
  const outGuitarScore = outG.flatMap((t) => t.scoreNotes);
  if (refGuitar.length > 0 && outGuitarScore.length > 0) {
    const refScore: ScoreNote[] = refGuitar.map((n) => ({
      tick: Math.round(secondsToTicks(n.onset, truth.bpm)),
      dur: Math.round(secondsToTicks(n.offset - n.onset, truth.bpm)),
      pitch: n.pitch, string: n.string, fret: n.fret,
    }));
    const align = alignGrid(refScore, outGuitarScore);
    rhythm = { factor: align.factor, beatOffset: align.beatOffset, ...rhythmScore(refScore, outGuitarScore, align) };
  }

  let withString = 0, agree = 0;
  for (const [i, j] of gm.matches) {
    const a = refGuitar[i]!.string, b = outGuitar[j]!.string;
    if (a !== undefined && b !== undefined) { withString++; if (a === b) agree++; }
  }

  let gpPitch: SongScore["gpPitch"] = null;
  let barAgreement: SongScore["barAgreement"] = null;
  if (outGp) {
    const re = truthFromGp(outGp);
    const shift = out.grid.firstDownbeat; // the GP's time zero is the grid's first downbeat
    const gpGuitar = notesOf(re.tracks, "guitar").map((n) => ({ ...n, onset: n.onset + shift, offset: n.offset + shift }));
    gpPitch = matchNotes(refGuitar, gpGuitar);
    const barSec = (60 / truth.bpm) * truth.timeSig[0] * (4 / truth.timeSig[1]);
    const lastEnd = Math.max(0, ...truth.tracks.flatMap((t) => t.notes.map((n) => n.offset)));
    const refBars = truth.bars ?? Math.ceil(lastEnd / barSec - 1e-9);
    barAgreement = { refBars, outBars: re.bars, timeSigMatch: truth.timeSig[0] === re.timeSig[0] && truth.timeSig[1] === re.timeSig[1] };
  }

  return {
    flags, guitarMerged, partSplit, perRole, drumsPerClass, rhythm,
    stringAgreement: withString ? agree / withString : null,
    spuriousRate: outGuitar.length ? gm.fp / outGuitar.length : null,
    gpPitch, barAgreement,
  };
}