"""PLUS transcribe: пост-обработка MIDI от MuScriptor / Aria.

Стадии независимы и включаются флагами:
  cleanup   унисоны, фантомы, перекрытия, короткие ноты
  quantize  посадка на сетку файла (сила 0..1, свинг)
  hands     разделение на правую и левую руку
  pedal     вывод педали (CC64) и продление нот под ней
  chords    распознавание аккордов и подписи в MIDI
  dynamics  громкость из аудио (RMS + атака), баланс рук
  report    метрики, аномалии, JSON

CLI:
  python plus_transcribe.py --midi in.mid --out out.mid --audio in.wav \
      --stages cleanup,quantize,hands,pedal,chords,dynamics,report \
      --quantize-strength 0.85 --report out.json
"""
from __future__ import annotations

import argparse
import math
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from plus_core import (  # noqa: E402
    Note, Piece, ambiguity_of, dump_json, explained_fraction, load_midi,
    pitch_name, save_midi, step, tempo_from_onsets,
)

__version__ = "1.0"
STAGES = ("cleanup", "quantize", "hands", "pedal", "chords", "dynamics", "report")

# ---------------------------------------------------------------- стадии ----

def stage_cleanup(piece: Piece, options: dict) -> dict:
    """Убирает дубли, фантомы по громкости и нефизичные длительности.

    Три правила, каждое с отдельным счётчиком, чтобы отчёт был честным:
    1) унисон: две ноты одного питча в пределах 45 мс и с перекрытием > 60%
       считаются одной (остаётся более громкая, длительность — объединение);
    2) фантом: нота короче 30 мс или громче чем на 25 ниже медианы ансамбля
       при том же питче в пределах 200 мс;
    3) обрезка: длительность не длиннее времени до следующей ноты того же
       питча минус 10 мс (иначе в DAW нота глушит саму себя).
    """
    min_dur = float(options.get("min_dur", 0.03))
    unison_ms = float(options.get("unison_ms", 45.0))
    result: list[Note] = []
    removed_short = removed_ghost = merged = trimmed = 0

    by_pitch: dict[int, list[Note]] = {}
    for note in sorted(piece.notes, key=lambda n: (n.pitch, n.start)):
        by_pitch.setdefault(note.pitch, []).append(note)

    for pitch, group in by_pitch.items():
        kept: list[Note] = []
        for note in group:
            if note.dur < min_dur:
                removed_short += 1
                continue
            if kept:
                last = kept[-1]
                overlap = min(last.end, note.end) - max(last.start, note.start)
                same_time = abs(note.start - last.start) <= unison_ms / 1000.0
                if same_time or (overlap > 0.6 * min(last.dur, note.dur)):
                    if note.velocity > last.velocity:
                        kept[-1] = Note(last.start, max(last.end, note.end) - last.start,
                                        pitch, note.velocity, note.channel, note.track)
                    else:
                        kept[-1] = Note(last.start, max(last.end, note.end) - last.start,
                                        pitch, last.velocity, last.channel, last.track)
                    merged += 1
                    continue
            kept.append(Note(**vars(note)))

        # одиночная тихая короткая нота на редком питче — почти всегда фантом
        if len(kept) <= 2:
            global_median = float(np.median([n.velocity for n in piece.notes])) if piece.notes else 80.0
            filtered = []
            for note in kept:
                if note.dur < 0.12 and note.velocity < global_median - 35:
                    removed_ghost += 1
                    continue
                filtered.append(note)
            kept = filtered

        # фантомы: слишком тихие на фоне соседей того же питча
        if len(kept) >= 5:
            levels = np.array([n.velocity for n in kept], dtype=float)
            floor = np.median(levels) - 25.0
            pruned = []
            for note in kept:
                if note.velocity < floor and note.dur < 0.12:
                    removed_ghost += 1
                    continue
                pruned.append(note)
            kept = pruned

        # обрезка до следующей ноты того же питча
        for index, note in enumerate(kept[:-1]):
            gap = kept[index + 1].start - note.start
            if note.dur > gap:
                note.dur = max(min_dur, gap - 0.01)
                trimmed += 1
        result.extend(kept)

    result.sort(key=lambda n: (n.start, n.pitch))
    piece.notes = result
    return {"short_removed": removed_short, "ghosts_removed": removed_ghost,
            "unisons_merged": merged, "trimmed": trimmed}


def stage_quantize(piece: Piece, options: dict) -> dict:
    """Посадка онсетов и длительностей на сетку.

    Темп из файла часто врёт: модели пишут заголовок 120, а ноты ставят по
    своей сетке. Поэтому сначала ищем темп, который ОБЪЯСНЯЕТ онсеты (файловый,
    его кратные и скан), и подменяем заголовок, только если найденное объяснение
    заметно лучше. Если даже лучший кандидат объясняет меньше 75 процентов
    онсетов, квантование пропускается: ровнять по чужой сетке хуже, чем не
    ровнять вовсе.
    """
    division = int(options.get("division", 4))
    strength = float(options.get("strength", 0.85))
    swing = float(options.get("swing", 0.0))
    forced = options.get("tempo")
    file_bpm = float(piece.bpm or 120.0)
    notes = piece.notes
    if not notes:
        return {"bpm": round(file_bpm, 2), "division": division, "strength": strength,
                "swing": swing, "explained_before": 0.0, "explained_after": 0.0,
                "notes_moved": 0, "corrected_tempo": False, "skipped": True,
                "reason": "нет нот", "alternatives": []}

    starts = piece.starts()
    sample = starts if len(starts) <= 400 else starts[:: max(1, len(starts) // 400)]
    file_fraction = explained_fraction(starts, file_bpm)

    def score(candidate):
        return explained_fraction(sample, candidate)

    # Только простые отношения к темпу файла: плотный скан даёт произвольные
    # темпы вида 108, которые «объясняют» сетку случайно.
    factors = (1.0, 1.5, 2.0 / 3.0, 2.0, 0.5, 3.0, 1.0 / 3.0)
    best_bpm, best_fraction = file_bpm, file_fraction
    if forced is None and options.get("auto_tempo", True):
        for factor in factors:
            candidate = file_bpm * factor
            if not (30 <= candidate <= 260):
                continue
            value = score(candidate)
            if value > best_fraction + 1e-9:
                best_bpm, best_fraction = candidate, value
    if forced is not None:
        best_bpm, best_fraction = float(forced), score(float(forced))

    threshold = float(options.get("min_explained", 0.75))
    if best_fraction < threshold:
        return {"bpm": round(file_bpm, 2), "division": division, "strength": strength,
                "swing": swing, "explained_before": round(file_fraction, 3),
                "explained_after": round(file_fraction, 3), "notes_moved": 0,
                "corrected_tempo": False, "skipped": True,
                "reason": "сетка объясняет только %.2f онсетов (порог %.2f), тайминг не трогаю"
                          % (best_fraction, threshold),
                "best_candidate_bpm": round(best_bpm, 2),
                "alternatives": ambiguity_of(starts, best_bpm)}

    corrected = abs(np.log2(best_bpm / file_bpm)) > 0.02
    bpm = float(best_bpm)
    beat = 60.0 / bpm
    step_s = beat / division
    fitted = 0
    for note in notes:
        target = round(note.start / step_s) * step_s
        if swing and int(round(note.start / step_s)) % 2 == 1:
            target += swing * step_s
        moved = (target - note.start) * strength
        note.start = max(0.0, note.start + moved)
        if abs(moved) > 1e-6:
            fitted += 1
        length = round(note.dur / step_s) * step_s
        if length > 0:
            note.dur = max(1e-3, note.dur + (length - note.dur) * strength)
    notes.sort(key=lambda n: (n.start, n.pitch))
    if corrected:
        # заголовок приводим к найденной сетке, иначе такты в DAW встанут не туда
        piece.tempo_changes = [(0.0, bpm)]
    after = explained_fraction(piece.starts(), bpm)
    return {"bpm": round(bpm, 2), "division": division, "strength": strength,
            "swing": swing, "explained_before": round(file_fraction, 3),
            "explained_after": round(after, 3), "notes_moved": fitted,
            "corrected_tempo": corrected, "skipped": False,
            "alternatives": ambiguity_of(starts, bpm)}


def stage_hands(piece: Piece, options: dict) -> dict:
    """Делит ноты на правую и левую руку по стоимости перехода.

    Жадный проход с двумя курсорами: нота идёт в ту руку, где ближе по питчу и
    по времени, с запретом левой руке уходить выше правой больше чем на октаву.
    Работает на одну дорожку, поэтому подходит и для фортепианных моделей,
    которые пишут всё в один канал.
    """
    if not piece.notes:
        return {"left": 0, "right": 0}
    split_pitch = int(options.get("split_pitch", 60))
    notes = sorted(piece.notes, key=lambda n: (n.start, n.pitch))
    left_last = right_last = None
    left_cost = right_cost = 0.0
    for note in notes:
        def cost(last, side):
            if last is None:
                base = abs(note.pitch - (split_pitch - 7 if side == "left" else split_pitch + 7))
                return base / 12.0 + 0.6
            return abs(note.pitch - last.pitch) / 12.0 + min(1.0, (note.start - last.start) / 2.0)

        to_left, to_right = cost(left_last, "left"), cost(right_last, "right")
        if to_left < to_right:
            note.track = "LH"
            left_last, left_cost = note, left_cost + to_left
        else:
            note.track = "RH"
            right_last, right_cost = note, right_cost + to_right
    # левая рука не должна оказаться выше правой: меняем подозрительные пары
    swapped = 0
    rh = [n for n in notes if n.track == "RH"]
    lh = [n for n in notes if n.track == "LH"]
    if rh and lh:
        rh_median = float(np.median([n.pitch for n in rh]))
        for note in lh:
            if note.pitch > rh_median + 12:
                note.track = "RH"
                swapped += 1
    piece.notes.sort(key=lambda n: (n.start, n.pitch))
    left = sum(1 for n in piece.notes if n.track == "LH")
    return {"left": left, "right": len(piece.notes) - left, "raised_from_left": swapped}


def stage_pedal(piece: Piece, options: dict) -> dict:
    """Выводит педаль по смене гармонии: CC64 плюс продление нот.

    Онсеты кластеризуются в аккордовые события (окно 60 мс). Педаль нажимается
    на аккорде, где >= 2 нот, и отпускается, когда меняется набор высотных
    классов или истекает max_hold. Продление ограничено следующей нотой того же
    питча и концом окна, поэтому каши из наложенных нот не возникает.
    """
    if not piece.notes:
        return {"spans": 0, "extended": 0}
    max_hold = float(options.get("max_hold", 1.6))
    window = float(options.get("window", 0.06))
    notes = sorted(piece.notes, key=lambda n: n.start)

    clusters: list[list[Note]] = []
    for note in notes:
        if clusters and note.start - clusters[-1][0].start <= window:
            clusters[-1].append(note)
        else:
            clusters.append([note])

    spans: list[tuple[float, float]] = []
    current: tuple[float, frozenset] | None = None
    for index, cluster in enumerate(clusters):
        classes = frozenset(n.pitch % 12 for n in cluster)
        start_time = cluster[0].start
        next_time = clusters[index + 1][0].start if index + 1 < len(clusters) else piece.duration + 0.4
        if len(cluster) < 2:
            if current is not None:
                spans.append((current[0], start_time))
                current = None
            continue
        if current is None:
            current = (start_time, classes)
        elif classes != current[1] or start_time - current[0] > max_hold:
            spans.append((current[0], start_time))
            current = (start_time, classes)
        if index + 1 == len(clusters):
            end_time = min(next_time, current[0] + max_hold)
            spans.append((current[0], end_time))
            current = None
    if current is not None:
        spans.append((current[0], min(piece.duration + 0.4, current[0] + max_hold)))
    spans = [(a, b) for a, b in spans if b - a > 0.05]

    by_pitch: dict[int, list[float]] = {}
    for note in notes:
        by_pitch.setdefault(note.pitch, []).append(note.start)
    for starts in by_pitch.values():
        starts.sort()

    extended = 0
    for note in notes:
        for start, end in spans:
            if not (start <= note.start < end):
                continue
            following = next((s for s in by_pitch[note.pitch] if s > note.start + 1e-4), None)
            limit = end if following is None else min(end, following - 0.01)
            if limit > note.end:
                note.dur = limit - note.start
                extended += 1
            break
    piece.meta["pedal_spans"] = [(round(a, 3), round(b, 3)) for a, b in spans]
    return {"spans": len(spans), "extended": extended,
            "median_hold": round(float(np.median([b - a for a, b in spans])), 3) if spans else 0.0}


CHORD_TEMPLATES = {
    "": (0, 4, 7), "m": (0, 3, 7), "7": (0, 4, 7, 10), "m7": (0, 3, 7, 10),
    "maj7": (0, 4, 7, 11), "dim": (0, 3, 6), "aug": (0, 4, 8), "sus4": (0, 5, 7),
    "sus2": (0, 2, 7), "6": (0, 4, 7, 9), "m6": (0, 3, 7, 9),
}


def _chord_name(pitches: list[int]) -> str | None:
    classes = sorted({p % 12 for p in pitches})
    if len(classes) < 3:
        return None
    best = None
    for root in classes:
        relative = sorted((c - root) % 12 for c in classes)
        for suffix, template in CHORD_TEMPLATES.items():
            if len(template) != len(relative):
                continue
            if all(abs(a - b) <= 1 for a, b in zip(relative, template)):
                score = sum(abs(a - b) for a, b in zip(relative, template))
                if best is None or score < best[0]:
                    best = (score, pitch_name(root + 60)[:-1] + suffix)
    return best[1] if best else None


def stage_chords(piece: Piece, options: dict) -> dict:
    """Аккорды по одновременным онсетам: список для отчёта и метки в piece.meta."""
    window = float(options.get("window", 0.06))
    notes = sorted(piece.notes, key=lambda n: n.start)
    groups: list[list[Note]] = []
    for note in notes:
        if groups and note.start - groups[-1][0].start <= window:
            groups[-1].append(note)
        else:
            groups.append([note])
    timeline = []
    named = 0
    for group in groups:
        pitches = [n.pitch for n in group]
        if len(group) < 3:
            continue
        name = _chord_name(pitches)
        if name:
            named += 1
        timeline.append({"time": round(group[0].start, 3), "size": len(group),
                         "pitches": sorted(pitches), "chord": name})
    piece.meta["chords"] = timeline
    return {"clusters": len(timeline), "named": named,
            "first": timeline[:8] if timeline else []}


def stage_dynamics(piece: Piece, options: dict) -> dict:
    """Громкость из аудио: RMS и атака в окне от онсета, затем кривая в MIDI.

    Velocity = 1 + shaped*123, где shaped это перцентильный share громкости,
    поднятый в степень 0.65 и умноженный на акцент атаки. Для рук добавляется
    баланс: левая на gain_left, правая на gain_right, чтобы форте в левой не
    перекрывало мелодию.
    """
    audio_path = options.get("audio")
    if not audio_path or not os.path.isfile(audio_path) or not piece.notes:
        return {"applied": False, "reason": "нет аудио"}
    try:
        import librosa
    except Exception as exc:
        return {"applied": False, "reason": "нет librosa: %s" % exc}

    window = float(options.get("window", 0.25))
    gain_left = float(options.get("gain_left", 0.88))
    gain_right = float(options.get("gain_right", 1.0))
    samples, rate = librosa.load(audio_path, sr=22050, mono=True)
    hop = 128
    frame = 512
    rms = librosa.feature.rms(y=samples, frame_length=frame, hop_length=hop)[0]
    frame_times = librosa.frames_to_time(np.arange(rms.size), sr=rate, hop_length=hop)

    levels, attacks = [], []
    for note in piece.notes:
        index = int(np.searchsorted(frame_times, note.start))
        span = rms[index: index + max(1, int(window * rate / hop))]
        if span.size == 0:
            levels.append(0.0)
            attacks.append(1.0)
            continue
        peak = float(span.max())
        mean = float(span.mean()) + 1e-9
        levels.append(20.0 * math.log10(mean + 1e-9))
        attacks.append(min(2.0, peak / mean))
    levels = np.array(levels)
    attacks = np.array(attacks)
    quiet, loud = np.percentile(levels, 4), np.percentile(levels, 99.7)
    share = np.clip((levels - quiet) / max(1e-6, loud - quiet), 0.0, 1.0)
    shaped = np.power(share, 0.65) * (1.0 + 0.18 * (attacks - 1.0))
    velocity = np.clip(15 + shaped * 103, 15, 118).astype(int)
    for note, value in zip(piece.notes, velocity):
        gain = gain_left if note.track == "LH" else gain_right
        note.velocity = int(max(1, min(127, round(value * gain))))
    piece.meta["dynamics"] = {"quiet_db": round(float(quiet), 2), "loud_db": round(float(loud), 2)}
    spread = int(velocity.max() - velocity.min())
    return {"applied": True, "velocity_min": int(velocity.min()),
            "velocity_max": int(velocity.max()), "spread": spread,
            "window": window, "gain_left": gain_left}


def stage_report(piece: Piece, options: dict) -> dict:
    """Метрики и аномалии: что чинить дальше, без догадок."""
    notes = piece.notes
    if not notes:
        return {"empty": True}
    starts = piece.starts()
    durations = np.array([n.dur for n in notes])
    pitches = piece.pitches()
    velocities = piece.velocities()
    span = max(1e-9, piece.duration)
    overlaps = 0
    by_pitch: dict[int, list[Note]] = {}
    for note in sorted(notes, key=lambda n: (n.pitch, n.start)):
        by_pitch.setdefault(note.pitch, []).append(note)
    for group in by_pitch.values():
        for a, b in zip(group, group[1:]):
            if b.start < a.end - 1e-3:
                overlaps += 1
    density = np.array([np.sum((starts >= t) & (starts < t + 1.0)) for t in np.arange(0, span, 1.0)])
    anomalies = []
    if overlaps:
        anomalies.append({"kind": "overlap", "count": overlaps})
    if durations.min() < 0.02:
        anomalies.append({"kind": "very_short_notes", "count": int((durations < 0.02).sum())})
    if density.size and density.max() > 25:
        anomalies.append({"kind": "density_spike", "per_second": int(density.max())})
    if velocities.size and velocities.std() < 3:
        anomalies.append({"kind": "flat_velocity", "std": round(float(velocities.std()), 2)})
    octave_jumps = int(np.sum(np.abs(np.diff(np.sort(pitches))) > 24))
    return {
        "notes": len(notes), "duration": round(piece.duration, 3),
        "notes_per_second": round(len(notes) / span, 2),
        "pitch_range": [pitch_name(int(pitches.min())), pitch_name(int(pitches.max()))],
        "velocity": [int(velocities.min()), int(velocities.max()),
                     round(float(velocities.mean()), 1)],
        "duration_median": round(float(np.median(durations)), 3),
        "overlaps": overlaps, "octave_leaps": octave_jumps,
        "anomalies": anomalies,
        "tracks": {name: sum(1 for n in notes if n.track == name)
                   for name in sorted({n.track for n in notes})},
    }


HANDLERS = {
    "cleanup": stage_cleanup,
    "quantize": stage_quantize,
    "hands": stage_hands,
    "pedal": stage_pedal,
    "chords": stage_chords,
    "dynamics": stage_dynamics,
    "report": stage_report,
}


def run_pipeline(piece: Piece, stages: list[str], options: dict) -> dict:
    report: dict = {"version": __version__, "source": piece.source, "stages": {}}
    for index, name in enumerate(stages, 1):
        handler = HANDLERS.get(name)
        if handler is None:
            raise SystemExit("неизвестная стадия: " + name)
        step(index, len(stages), name)
        result = handler(piece, options)
        report["stages"][name] = result
        summary = ", ".join("%s=%s" % (k, v) for k, v in result.items()
                            if isinstance(v, (int, float, str)) and k != "first")
        print("      " + (summary[:150] if summary else "готово"), flush=True)
    report["stats"] = piece.stats()
    return report


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="PLUS: пост-обработка MIDI транскрипции")
    parser.add_argument("--midi", required=True, help="входной MIDI")
    parser.add_argument("--out", help="куда писать MIDI (по умолчанию рядом, суффикс _plus)")
    parser.add_argument("--audio", help="аудио для стадии dynamics")
    parser.add_argument("--stages", default="cleanup,quantize,hands,pedal,dynamics,report")
    parser.add_argument("--tempo", type=float, help="переопределить BPM")
    parser.add_argument("--quantize-strength", type=float, default=0.85)
    parser.add_argument("--division", type=int, default=4)
    parser.add_argument("--swing", type=float, default=0.0)
    parser.add_argument("--auto-tempo", action="store_true", help="подстроить темп в пределах 15%")
    parser.add_argument("--split-pitch", type=int, default=60)
    parser.add_argument("--gain-left", type=float, default=0.88)
    parser.add_argument("--report", help="JSON-отчёт")
    args = parser.parse_args(argv)

    stages = [s.strip() for s in args.stages.split(",") if s.strip()]
    out_path = args.out or os.path.splitext(args.midi)[0] + "_plus.mid"
    report_path = args.report or os.path.splitext(out_path)[0] + "_report.json"

    piece = load_midi(args.midi)
    print("вход: %s | нот %d | %.1f с | %.0f BPM"
          % (os.path.basename(args.midi), len(piece.notes), piece.duration, piece.bpm), flush=True)

    options = {
        "audio": args.audio, "tempo": args.tempo,
        "strength": args.quantize_strength, "division": args.division,
        "swing": args.swing, "auto_tempo": args.auto_tempo,
        "split_pitch": args.split_pitch, "gain_left": args.gain_left,
    }
    report = run_pipeline(piece, stages, options)
    save_midi(piece, out_path)
    dump_json(report, report_path)
    print("готово: %s | отчёт: %s" % (out_path, report_path), flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())