"""Слияние нескольких проходов транскрипции в один результат.

Зачем: один проход ошибается системно — с широким conditioning модель
«размазывается» и теряет партии (колокольчики, бас), с узким — находит их,
но теряет остальное. Поэтому делаем несколько проходов с РАЗНЫМ фокусом
и сливаем: нота выживает, если её подтвердил хотя бы один проход, а дубликаты
одного питча в пределах допуска склеиваются в одну.

Слияние не трогает высоты и время (медиана подтверждений), поэтому качество
не размывается: это объединение доказательств, а не усреднение.
"""

import json
import math
from collections import defaultdict

TOLERANCE = 0.07          # окно склейки одного питча, секунды


def parse_sse(text):
    """SSE-поток сервера → список нот {pitch, start, end, instrument, velocity}."""
    open_notes = {}
    notes = []
    for line in text.splitlines():
        if not line.startswith("data: "):
            continue
        try:
            event = json.loads(line[6:])
        except Exception:
            continue
        kind = event.get("type")
        if kind == "start":
            key = event.get("index", len(notes))
            open_notes[key] = {"pitch": int(event["pitch"]),
                               "start": float(event["start_time"]),
                               "instrument": event.get("instrument", "unknown"),
                               "index": key}
        elif kind == "end":
            key = event.get("start_event_index")
            note = open_notes.pop(key, None)
            if note is None:
                continue
            note["end"] = float(event.get("end_time", note["start"] + 0.1))
            notes.append(note)
    for note in open_notes.values():          # нота без конца = до конца файла
        note["end"] = note["start"] + 0.1
        notes.append(note)
    notes.sort(key=lambda n: (n["start"], n["pitch"]))
    return notes


def merge(passes, tolerance=TOLERANCE, prefer=None):
    """passes: список списков нот. Возвращает слитый список и статистику."""
    buckets = defaultdict(list)
    for order, notes in enumerate(passes):
        for note in notes:
            key = (note["pitch"], round(note["start"] / tolerance))
            buckets[key].append((order, note))

    # соседние окна того же питча могут относиться к одной ноте — склеиваем
    merged = []
    used = set()
    for (pitch, slot) in sorted(buckets):
        if (pitch, slot) in used:
            continue
        group = list(buckets[(pitch, slot)])
        used.add((pitch, slot))
        for neighbour in ((pitch, slot - 1), (pitch, slot + 1)):
            if neighbour in buckets and neighbour not in used:
                if abs(buckets[neighbour][0][1]["start"] - group[0][1]["start"]) <= tolerance:
                    group += buckets[neighbour]
                    used.add(neighbour)
        starts = sorted(n["start"] for _, n in group)
        ends = sorted(n["end"] for _, n in group)
        sources = sorted({order for order, _ in group})
        instrument = prefer or max(
            (n["instrument"] for _, n in group),
            key=lambda name: sum(1 for _, n in group if n["instrument"] == name))
        merged.append({
            "pitch": pitch,
            "start": starts[len(starts) // 2],
            "end": max(ends[len(ends) // 2], starts[len(starts) // 2] + 0.05),
            "instrument": instrument,
            "support": len(sources),
            "passes": sources,
        })
    merged.sort(key=lambda n: (n["start"], n["pitch"]))
    stats = {
        "passes": len(passes),
        "input_notes": [len(p) for p in passes],
        "merged": len(merged),
        "confirmed_by_two": sum(1 for n in merged if n["support"] >= 2),
        "added_by_second": sum(1 for n in merged if n["passes"] == [1]),
    }
    return merged, stats


def to_sse(notes, ping_note=None):
    """Слитые ноты → SSE-поток того же формата, что отдаёт сервер."""
    frames = [{"type": "progress", "completed": 0, "total": 1}]
    ends = []
    for index, note in enumerate(notes):
        frames.append({"type": "start", "pitch": note["pitch"],
                       "start_time": round(note["start"], 4), "index": index,
                       "instrument": note["instrument"]})
        ends.append({"type": "end", "end_time": round(note["end"], 4),
                     "start_event_index": index})
    frames += ends
    frames.append({"type": "progress", "completed": 1, "total": 1})
    body = "".join("data: %s\n\n" % json.dumps(frame, ensure_ascii=False) for frame in frames)
    if ping_note:
        body += ": %s\n\n" % ping_note
    return body.encode("utf-8")


def notes_to_export(notes):
    """Формат /export (JSON ExportRequest)."""
    return {"notes": [{"pitch": n["pitch"], "start": round(n["start"], 4),
                       "end": round(n["end"], 4),
                       "instrument": n["instrument"]} for n in notes]}