import { Song, Note, Bar, HitNote } from 'tja-parser';
import * as fs from 'fs';
import * as path from 'path';
import * as iconv from 'iconv-lite';
import { spawn } from 'child_process';

function getFeatures(song: Song, diff: 'oni' | 'edit') {
  const course = song.course[diff];
  if (!course) return null;

  const bars = course.noteGroups.filter(ng => typeof (ng as any).getNotes === 'function') as Bar[];
  const hitNotes = bars.flatMap(b => b.getNotes()).filter(n => [1, 2, 3, 4].includes(n.toJSON().type));
  if (hitNotes.length === 0) return null;

  const timings = hitNotes.map(n => n.getTimingMS() / 1000);
  const startT = Math.min(...timings);
  const endT = Math.max(...timings);
  const duration = endT - startT;
  if (duration <= 0) return null;

  const intervals: number[] = [];
  for (let i = 1; i < timings.length; i++) intervals.push(timings[i] - timings[i-1]);

  const getPeakNPS = (window: number) => {
    let max = 0;
    for (let t = startT; t < endT; t += 0.2) {
      const count = timings.filter(time => time >= t && time < t + window).length;
      if (count / window > max) max = count / window;
    }
    return max;
  };

  const getNPSInRange = (t1: number, t2: number) => {
    const count = timings.filter(t => t >= t1 && t < t2).length;
    return count / ((t2 - t1) || 1);
  };
  const q = [0, 0.25, 0.5, 0.75, 1.0].map(p => startT + duration * p);
  const nps_qs = [
    getNPSInRange(q[0], q[1]), getNPSInRange(q[1], q[2]),
    getNPSInRange(q[2], q[3]), getNPSInRange(q[3], q[4])
  ];

  let max_delta = 0;
  for (let i = 1; i < nps_qs.length; i++) max_delta = Math.max(max_delta, Math.abs(nps_qs[i] - nps_qs[i-1]));

  const streams: number[] = [];
  let currentS = 0, restTime = 0;
  for (const interval of intervals) {
    if (interval <= 0.185) currentS++;
    else {
      if (currentS > 0) streams.push(currentS + 1);
      currentS = 0;
      if (interval >= 1.0) restTime += interval;
    }
  }
  if (currentS > 0) streams.push(currentS + 1);

  const donCount = hitNotes.filter(n => [1, 3].includes(n.toJSON().type)).length;
  const kaCount = hitNotes.length - donCount;

  const svs = bars.map(b => b.getScroll());
  const bpms = bars.map(b => b.getBpm());
  const gimmickCount = svs.filter((v, i) => i > 0 && v !== svs[i-1]).length + bpms.filter((v, i) => i > 0 && v !== bpms[i-1]).length;

  const avg = (arr: number[]) => arr.length ? arr.reduce((a, b) => a + b, 0) / arr.length : 0;
  const std = (arr: number[]) => Math.sqrt(avg(arr.map(v => Math.pow(v - avg(arr), 2))));

  return {
    global_nps: hitNotes.length / duration,
    effective_nps: hitNotes.length / (duration - (intervals.filter(v => v >= 1.0).reduce((a, b) => a + b, 0)) || 1),
    peak_nps_1s: getPeakNPS(1),
    peak_nps_2s: getPeakNPS(2),
    peak_nps_5s: getPeakNPS(5),
    nps_variance: std(intervals),
    nps_spike_index: getPeakNPS(1) / (hitNotes.length / duration || 1),
    nps_q1: nps_qs[0], nps_q2: nps_qs[1], nps_q3: nps_qs[2], nps_q4: nps_qs[3],
    max_nps_delta: max_delta,
    longest_stream: Math.max(0, ...streams),
    avg_stream_length: avg(streams),
    stream_count: streams.length,
    stream_ratio: streams.reduce((a, b) => a + b, 0) / hitNotes.length,
    resting_ratio: restTime / duration,
    don_ratio: donCount / hitNotes.length,
    ka_ratio: kaCount / hitNotes.length,
    color_transition_ratio: hitNotes.filter((n, i) => i > 0 && [1, 3].includes(n.toJSON().type) !== [1, 3].includes(hitNotes[i-1].toJSON().type)).length / hitNotes.length,
    pattern_complexity: hitNotes.filter((n, i) => i > 1 && ([1, 3].includes(hitNotes[i-2].toJSON().type) !== [1, 3].includes(hitNotes[i-1].toJSON().type) || [1, 3].includes(hitNotes[i-1].toJSON().type) !== [1, 3].includes(n.toJSON().type))).length / hitNotes.length,
    hand_switch_ratio: streams.filter(s => s % 2 !== 0).length / (streams.length || 1),
    off_beat_ratio: hitNotes.filter(n => Number((n.getTiming() as any).d) % 3 === 0).length / hitNotes.length,
    rhythm_std_dev: std(intervals),
    rhythmic_entropy: new Set(hitNotes.map(n => Number((n.getTiming() as any).d))).size / 10,
    gimmick_density: gimmickCount / duration,
    sv_variance: std(svs),
    bpm_variance: std(bpms),
    max_bpm: Math.max(...bpms),
    min_bpm: Math.min(...bpms),
    max_sv: Math.max(...svs),
    total_notes: hitNotes.length,
    level: course.getLevel()
  };
}

async function run() {
  const modelPath = process.argv[2] || 'model/constant_predictor.joblib';
  const tjaPath = process.argv[3];
  const diff: any = process.argv[4] || 'oni';

  if (!tjaPath) {
    console.log("Usage: bun run src/predict.ts <model_path> <tja_path> [oni|edit]");
    return;
  }

  const buffer = fs.readFileSync(tjaPath);
  const content = iconv.decode(buffer, 'shift-jis');
  const song = Song.parse(content);
  const features = getFeatures(song, diff);

  if (!features) {
    console.error("Course not found or no notes.");
    return;
  }

  const python = spawn('python3', ['model/predict.py', modelPath]);
  python.stdin.write(JSON.stringify(features));
  python.stdin.end();

  python.stdout.on('data', (data) => {
    console.log(`Predicted Constant: ${data.toString().trim()}`);
  });

  python.stderr.on('data', (data) => {
    console.error(`Python Error: ${data.toString()}`);
  });
}

run();
