import { Course, Bar, Note, HitNote } from 'tja-parser';

export namespace Feature {
    export function filterHitNotes(course: Course) {
        const notes: HitNote[] = [];
        course.noteGroups.forEach((g) => {
            if (g instanceof Bar) {
                for (const note of g.getNotes()) {
                    if (note instanceof HitNote) {
                        notes.push(note);
                    }
                }
            }
        });
        return notes;
    }

    export function getAllNotes(course: Course) {
        const notes: { type: number, time: number }[] = [];
        for (const group of course.noteGroups) {
            if (group instanceof Bar) {
                notes.push(...group.getNotes().map((note) => ({
                    //@ts-expect-error
                    type: note.type,
                    time: note.getTimingMS()
                })));
            }
        }
        return notes.sort((a, b) => a.time - b.time).filter(n => [1, 2, 3, 4].includes(n.type));
    }

    // 밀도 관련
    export function getAverageDensity(notes: HitNote[]) {
        if (notes.length === 0) {
            return 0;
        }
        return notes.length / (notes[notes.length - 1].getTimingMS() - notes[0].getTimingMS()) * 1000
    }

    export function getPeakDensity(notes: HitNote[]) {
        let peakDensity = 0;
        for (let i = 0; i < notes.length; i++) {
            let noteCount = 1;
            for (let j = i + 1; j < notes.length && (notes[j].getTimingMS() - notes[i].getTimingMS() <= 1000); j++) {
                noteCount++;
            }
            if (peakDensity < noteCount) {
                peakDensity = noteCount;
            }
        }
        return peakDensity;
    }

    // BPM 관련
    export function getAverageBPM(notes: HitNote[]) {
        if(notes.length === 0) return 0;
        const averageBPM = notes.reduce((p, note) => p + note.getBPM().valueOf(), 0) / notes.length;
        return averageBPM;
    }

    export function getBpmChange(notes: HitNote[]) {
        let bpmChange = 0;
        for (let i = 1; i < notes.length; i++) {
            if (Math.abs(notes[i].getBPM() - notes[i - 1].getBPM()) > 1.5) {
                bpmChange++;
            }
        };
        return bpmChange;
    }

    // 스크롤 변화
    export function getScrollChange(notes: HitNote[]) {
        let bpmChange = 0;
        for (let i = 1; i < notes.length; i++) {
            if (Math.abs(notes[i].getBPM() * notes[i].getScroll() - notes[i - 1].getBPM() * notes[i - 1].getScroll()) > 1.5) {
                bpmChange++;
            }
        };
        return bpmChange;
    }

    export function getRhythmComplexity(notes: HitNote[]) {
        let complexity = 0;
        for (let i = 2; i < notes.length; i++) {
            const d1 = notes[i].getTimingMS() - notes[i - 1].getTimingMS()
            const d2 = notes[i - 1].getTimingMS() - notes[i - 2].getTimingMS()

            const ratio = d1 / d2;
            const log = Math.log2(ratio);

            if (Math.abs(log - Math.round(log)) < 1e-3) {
                complexity++;
            }
        }
        return complexity;
    }

    export function getColorComplexity(notes: HitNote[]) {
        let complexity = 0;
        for (let i = 2; i < notes.length; i++) {
            if (notes[i].type % 2 != notes[i - 2].type % 2) {
                complexity += 1 / ((notes[i].getTimingMS() - notes[i - 2].getTimingMS()) ** 2)
            }
        }
        return complexity;
    }
}

export function featurize(course: Course) {
    const notes = Feature.filterHitNotes(course)
    return {
        note_count: notes.length,
        density_avg: Feature.getAverageDensity(notes),
        density_peak: Feature.getPeakDensity(notes),
        bpm_avg: Feature.getAverageBPM(notes),
        bpm_change: Feature.getBpmChange(notes),
        scroll_change: Feature.getScrollChange(notes),
        rhythm_complexity: Feature.getRhythmComplexity(notes),
        color_complexity: Feature.getColorComplexity(notes)
    };
}