main.py

import json
import os
import uuid
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

import cv2
import numpy as np
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import JSONResponse

import omr_neural
import score_preprocess
import staff_rectify

app = FastAPI(title="Stella Score Reader API", version="0.2.0")

ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp"}
MAX_SINGLE_IMAGE_BYTES = 10 * 1024 * 1024
MAX_TOTAL_BYTES = 30 * 1024 * 1024
_MAX_EVENTS_PER_STAFF = max(8, int(os.environ.get("STELLA_MAX_EVENTS_PER_STAFF", "256")))


class SingleStaffAnalyzeError(Exception):
    """Raised when a segment cannot be analyzed under single-staff v1 rules."""

    def __init__(self, code: str, message: str, details: Optional[Dict[str, Any]] = None) -> None:
        super().__init__(message)
        self.code = code
        self.message = message
        self.details = details or {}


@dataclass
class StaffCandidate:
    line_ys: List[int]
    x0: int
    y0: int
    width: int
    height: int


@app.get("/")
def read_root() -> Dict[str, str]:
    return {"status": "running"}


def _get_extension(filename: str) -> str:
    lower_name = filename.lower()
    dot_index = lower_name.rfind(".")
    return lower_name[dot_index:] if dot_index >= 0 else ""


def _validate_and_decode_image(upload_file: UploadFile) -> Dict[str, Any]:
    filename = upload_file.filename or "unknown"
    extension = _get_extension(filename)
    if extension not in ALLOWED_EXTENSIONS:
        raise HTTPException(
            status_code=415,
            detail=f"Unsupported image extension for '{filename}'.",
        )

    contents = upload_file.file.read()
    if len(contents) == 0:
        raise HTTPException(status_code=400, detail=f"Empty file: '{filename}'.")
    if len(contents) > MAX_SINGLE_IMAGE_BYTES:
        raise HTTPException(
            status_code=413,
            detail=f"Image too large: '{filename}' exceeds 10MB limit.",
        )

    np_buffer = np.frombuffer(contents, np.uint8)
    image = cv2.imdecode(np_buffer, cv2.IMREAD_COLOR)
    if image is None:
        raise HTTPException(
            status_code=400,
            detail=f"Unable to decode image: '{filename}'.",
        )

    height, width = image.shape[:2]
    return {
        "filename": filename,
        "bytes": len(contents),
        "width": int(width),
        "height": int(height),
        "image": image,
    }


def _parse_score_context(raw: Optional[str]) -> Dict[str, Any]:
    if raw is None or not str(raw).strip():
        raise HTTPException(status_code=400, detail="score_context is required (JSON string).")
    try:
        obj = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise HTTPException(status_code=400, detail=f"Invalid score_context JSON: {exc}") from exc
    if not isinstance(obj, dict):
        raise HTTPException(status_code=400, detail="score_context must be a JSON object.")

    clef = obj.get("clef")
    if clef not in ("treble", "bass"):
        raise HTTPException(
            status_code=400,
            detail="score_context.clef must be 'treble' or 'bass'.",
        )

    ks = obj.get("key_signature")
    if not isinstance(ks, dict):
        raise HTTPException(status_code=400, detail="score_context.key_signature must be an object.")
    if "fifths" not in ks:
        raise HTTPException(status_code=400, detail="score_context.key_signature.fifths is required.")
    fifths = ks["fifths"]
    if not isinstance(fifths, int) or fifths < -6 or fifths > 6:
        raise HTTPException(
            status_code=400,
            detail="score_context.key_signature.fifths must be an integer in [-6, 6].",
        )

    time_sig = obj.get("time_signature", "4/4")
    if not isinstance(time_sig, str) or not time_sig.strip():
        raise HTTPException(status_code=400, detail="score_context.time_signature must be a non-empty string.")

    tempo = obj.get("tempo_bpm_reference")
    if tempo is not None and (not isinstance(tempo, (int, float)) or tempo <= 0):
        raise HTTPException(
            status_code=400,
            detail="score_context.tempo_bpm_reference must be a positive number or null.",
        )

    divisions = obj.get("divisions", 4)
    if not isinstance(divisions, int) or divisions < 1 or divisions > 64:
        raise HTTPException(status_code=400, detail="score_context.divisions must be an integer in [1, 64].")

    return {
        "clef": clef,
        "key_signature": {"fifths": fifths},
        "time_signature": time_sig.strip(),
        "tempo_bpm_reference": tempo,
        "divisions": divisions,
    }


def _parse_options(options_raw: Optional[str]) -> Dict[str, Any]:
    defaults = {
        "return_debug": False,
        "quantization": "1/8",
    }
    if not options_raw:
        return defaults

    try:
        parsed = json.loads(options_raw)
        if not isinstance(parsed, dict):
            raise ValueError("options must be a JSON object")
    except (json.JSONDecodeError, ValueError) as exc:
        raise HTTPException(status_code=400, detail=f"Invalid options JSON: {exc}") from exc

    options = {**defaults, **parsed}
    if options["quantization"] not in {"1/4", "1/8", "1/16"}:
        raise HTTPException(
            status_code=400,
            detail="Invalid options.quantization. Allowed: 1/4, 1/8, 1/16.",
        )
    if not isinstance(options["return_debug"], bool):
        raise HTTPException(status_code=400, detail="options.return_debug must be a boolean.")
    return options


def _cluster_peaks(peaks: np.ndarray, min_gap: int = 2) -> List[int]:
    if len(peaks) == 0:
        return []
    grouped: List[List[int]] = [[int(peaks[0])]]
    for y in peaks[1:]:
        y_int = int(y)
        if y_int - grouped[-1][-1] <= min_gap:
            grouped[-1].append(y_int)
        else:
            grouped.append([y_int])
    return [int(sum(group) / len(group)) for group in grouped]


def _synthetic_staff_fallback(image_bgr: np.ndarray) -> Optional[StaffCandidate]:
    """When morphology-based 5-line clustering fails (thin crops), estimate one staff from ink projection."""
    gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (3, 3), 0)
    binary = cv2.adaptiveThreshold(
        blur,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        31,
        15,
    )
    h, w = binary.shape[:2]
    row_density = np.sum(binary > 0, axis=1).astype(np.float32)
    mx = float(row_density.max()) if row_density.size else 0.0
    if mx < 1.0:
        return None
    active = np.where(row_density > max(8.0, mx * 0.18))[0]
    if active.size < 3:
        return None
    y_top = int(active[0])
    y_bottom = int(active[-1])
    span = y_bottom - y_top
    if span < 6:
        return None
    margin = max(2, min(span // 5, h // 8))
    y0 = max(0, y_top - margin)
    y1 = min(h - 1, y_bottom + margin)
    line_ys = [int(round(y_top + i * (y_bottom - y_top) / 4.0)) for i in range(5)]
    x_nonzero = np.where(np.sum(binary[y0 : y1 + 1, :] > 0, axis=0) > 0)[0]
    if len(x_nonzero) == 0:
        x0, width = 0, w
    else:
        x0 = int(x_nonzero[0])
        x1 = int(x_nonzero[-1])
        width = max(30, x1 - x0 + 1)
    height = max(20, y1 - y0 + 1)
    return StaffCandidate(line_ys, x0, y0, width, height)


def _detect_staff_candidates(image_bgr: np.ndarray) -> Tuple[List[StaffCandidate], List[str]]:
    gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (3, 3), 0)
    binary = cv2.adaptiveThreshold(
        blur,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        31,
        15,
    )
    h, w = binary.shape
    kernel_width = max(25, w // 12)
    horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_width, 1))
    horizontal = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel, iterations=1)

    row_density = np.sum(horizontal > 0, axis=1)
    threshold = max(20, int(np.max(row_density) * 0.35))
    peak_rows = np.where(row_density > threshold)[0]
    clustered_rows = _cluster_peaks(peak_rows)
    det_warnings: List[str] = []
    if len(clustered_rows) < 5:
        fb = _synthetic_staff_fallback(image_bgr)
        if fb:
            return [fb], ["staff_geometry_fallback_projection"]
        return [], []

    def _collect_from_threshold(threshold: int) -> List[StaffCandidate]:
        peak_rows = np.where(row_density > threshold)[0]
        clustered = _cluster_peaks(peak_rows)
        if len(clustered) < 5:
            return []
        out: List[StaffCandidate] = []
        j = 0
        while j + 4 < len(clustered):
            window = clustered[j : j + 5]
            gaps = np.diff(window)
            median_gap = int(np.median(gaps))
            if median_gap < 3:
                j += 1
                continue
            if max(abs(int(g) - median_gap) for g in gaps) <= max(3, int(median_gap * 0.8)):
                y_top = max(0, window[0] - 4 * median_gap)
                y_bottom = min(h - 1, window[-1] + 4 * median_gap)
                x_nonzero = np.where(np.sum(horizontal[y_top : y_bottom + 1, :] > 0, axis=0) > 0)[0]
                if len(x_nonzero) == 0:
                    x0 = 0
                    width = w
                else:
                    x0 = int(x_nonzero[0])
                    x1 = int(x_nonzero[-1])
                    width = max(30, x1 - x0 + 1)
                height = max(20, y_bottom - y_top + 1)
                out.append(StaffCandidate(window, x0, y_top, width, height))
                j += 5
            else:
                j += 1
        return out

    threshold_primary = max(20, int(np.max(row_density) * 0.35))
    candidates = _collect_from_threshold(threshold_primary)
    if not candidates:
        threshold_loose = max(12, int(np.max(row_density) * 0.22))
        if threshold_loose != threshold_primary:
            candidates = _collect_from_threshold(threshold_loose)

    if not candidates:
        fb = _synthetic_staff_fallback(image_bgr)
        if fb:
            return [fb], ["staff_geometry_fallback_projection"]

    if len(candidates) > 1:
        candidates = [max(candidates, key=lambda c: int(c.width) * int(c.height))]

    return candidates, det_warnings


def _midi_to_step_octave_alter(midi: int) -> Tuple[str, int, int]:
    note_names = [
        ("C", 0),
        ("C", 1),
        ("D", 0),
        ("D", 1),
        ("E", 0),
        ("F", 0),
        ("F", 1),
        ("G", 0),
        ("G", 1),
        ("A", 0),
        ("A", 1),
        ("B", 0),
    ]
    step, alter = note_names[midi % 12]
    octave = midi // 12 - 1
    return step, octave, alter


def _y_to_midi(y_center: int, staff_lines: List[int], clef: str) -> int:
    if len(staff_lines) < 5:
        return 60
    spacing = max(2.0, float((staff_lines[-1] - staff_lines[0]) / 4.0))
    relative_steps = round((staff_lines[-1] - y_center) / (spacing / 2.0))
    base_midi = 64 if clef == "treble" else 43
    midi = base_midi + int(relative_steps)
    return max(21, min(108, midi))


def _reduce_to_top_note_event(
    ev_base: Dict[str, Any],
    prev_top_midi: Optional[int] = None,
) -> Optional[Dict[str, Any]]:
    """
    Preserve parsed chord candidates, then reduce to one melody note.
    Prefer top note by default, but allow continuity-based override.
    """
    source_raw = ev_base.get("source_pitch_midis")
    source_midis: List[int] = []
    if isinstance(source_raw, list):
        for v in source_raw:
            if isinstance(v, (int, float)):
                source_midis.append(int(v))
    if not source_midis and isinstance(ev_base.get("pitch_midi"), (int, float)):
        source_midis.append(int(ev_base["pitch_midi"]))
    if not source_midis:
        return None

    source_midis = sorted({max(21, min(108, int(m))) for m in source_midis})
    top_midi = int(max(source_midis))
    pitch_midi = top_midi
    continuity_override_applied = False

    # Default is top-note reduction. For 2-part choir-like chords, allow
    # a lower candidate only if it greatly improves melodic continuity.
    if prev_top_midi is not None and len(source_midis) >= 2:
        best_midi = top_midi
        best_score = float(abs(top_midi - int(prev_top_midi)))
        for cand in source_midis:
            leap = float(abs(int(cand) - int(prev_top_midi)))
            top_penalty = 0.0 if cand == top_midi else 2.0
            score = leap + top_penalty
            if score + 1e-6 < best_score or (abs(score - best_score) <= 1e-6 and cand > best_midi):
                best_midi = int(cand)
                best_score = score
        if best_midi != top_midi:
            pitch_midi = int(best_midi)
            continuity_override_applied = True

    step, octave, alter = _midi_to_step_octave_alter(pitch_midi)
    return {
        "pitch_midi": pitch_midi,
        "step": step,
        "octave": int(octave),
        "alter": int(alter),
        "source_pitch_midis": source_midis,
        "source_note_count": len(source_midis),
        "reduced_from_chord": len(source_midis) > 1,
        "continuity_override_applied": continuity_override_applied,
        "top_pitch_midi": top_midi,
    }


def _fallback_staff_line_ys(roi_height: int) -> List[int]:
    margin = int(roi_height * 0.12)
    span = roi_height - 2 * margin
    if span < 10:
        margin = 0
        span = max(1, roi_height - 1)
    step = span / 4.0
    return [int(round(margin + i * step)) for i in range(5)]


def _detect_noteheads(staff_roi: np.ndarray) -> List[Tuple[int, int, int, int]]:
    gray = cv2.cvtColor(staff_roi, cv2.COLOR_BGR2GRAY)
    binary = cv2.adaptiveThreshold(
        gray,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        25,
        11,
    )
    horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 1))
    lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel, iterations=1)
    symbols = cv2.subtract(binary, lines)
    contours, _ = cv2.findContours(symbols, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    boxes: List[Tuple[int, int, int, int]] = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        area = w * h
        if area < 18 or area > 1200:
            continue
        aspect = w / max(1, h)
        if 0.35 <= aspect <= 2.4 and 4 <= h <= 40 and 3 <= w <= 40:
            boxes.append((x, y, w, h))
    boxes.sort(key=lambda b: (b[0], b[1]))
    return boxes


def build_analyze_response_v1(
    decoded_images: List[Dict[str, Any]],
    score_context: Dict[str, Any],
    *,
    return_debug: bool = False,
) -> Dict[str, Any]:
    """
    Core v1 pipeline: preprocess, single-staff check, dewarp, OMR (Flova/omr_transformer) + OpenCV fallback,
    merged melody timeline. Raises SingleStaffAnalyzeError on hard failures.
    """
    warnings: List[str] = []
    preprocess_segments: List[Dict[str, Any]] = []
    debug_segments: List[Dict[str, Any]] = []

    clef = str(score_context["clef"])
    key_fifths = int(score_context["key_signature"]["fifths"])
    divisions = int(score_context["divisions"])
    time_signature = str(score_context["time_signature"])
    tempo_ref = score_context.get("tempo_bpm_reference")

    melody_events: List[Dict[str, Any]] = []
    segment_map: List[Dict[str, Any]] = []
    global_onset = 0
    event_counter = 0
    any_neural = False
    chord_reduction_applied = 0
    chord_candidates_seen = 0
    continuity_override_applied = 0
    prev_note_midi: Optional[int] = None

    def _next_event_id() -> str:
        nonlocal event_counter
        event_counter += 1
        return f"e_{event_counter:04d}"

    for image_idx, image_meta in enumerate(decoded_images):
        seg_order = image_idx + 1
        work_bgr, page_geom, pre_meta = score_preprocess.preprocess_page_bgr(image_meta["image"])
        row: Dict[str, Any] = {
            "segment_order": seg_order,
            "original_size": {"width": page_geom.orig_w, "height": page_geom.orig_h},
            "work_size": {"width": page_geom.work_w, "height": page_geom.work_h},
            "uniform_scale": round(page_geom.scale_x, 5),
            "deskew_deg": round(page_geom.deskew_deg_applied, 3),
            **pre_meta,
        }
        if abs(page_geom.deskew_deg_applied) >= 3.0 and "large_deskew_correction_applied" not in warnings:
            warnings.append("large_deskew_correction_applied")

        candidates, det_ws = _detect_staff_candidates(work_bgr)
        warnings.extend(det_ws)
        if len(candidates) != 1:
            raise SingleStaffAnalyzeError(
                "UNPROCESSABLE_STAFF_LAYOUT",
                "Expected exactly one staff per image.",
                {
                    "segment_order": seg_order,
                    "detected_staff_candidates": len(candidates),
                },
            )

        staff = candidates[0]
        rectified, dewarp_meta = staff_rectify.rectify_staff_crop_bgr(
            work_bgr,
            staff.x0,
            staff.y0,
            staff.width,
            staff.height,
            staff.line_ys,
        )
        row.update(dewarp_meta)
        preprocess_segments.append(row)

        neural_notes, raw_omr, omr_err, _ = omr_neural.staff_image_to_note_events(
            rectified,
            key_fifths=key_fifths,
        )
        used_neural = bool(not omr_err and neural_notes)
        segment_events: List[Dict[str, Any]] = []

        if return_debug and raw_omr:
            raw_cap = raw_omr if len(raw_omr) <= 16000 else raw_omr[:16000] + "…"
            debug_segments.append(
                {
                    "segment_order": seg_order,
                    "filename": image_meta["filename"],
                    "raw": raw_cap,
                }
            )

        staff_bbox_orig = page_geom.work_rect_to_original_aabb(
            int(staff.x0),
            int(staff.y0),
            int(staff.width),
            int(staff.height),
        )

        if used_neural:
            any_neural = True
            for ev_base in neural_notes[:_MAX_EVENTS_PER_STAFF]:
                evt: Dict[str, Any] = {
                    "event_id": _next_event_id(),
                    "type": ev_base["type"],
                    "duration_div": int(ev_base["duration_div"]),
                    "onset_div": global_onset,
                    "segment_order": seg_order,
                    "bbox": staff_bbox_orig,
                }
                global_onset += int(ev_base["duration_div"])
                if ev_base["type"] == "note":
                    reduced = _reduce_to_top_note_event(ev_base, prev_top_midi=prev_note_midi)
                    if reduced is None:
                        continue
                    evt["step"] = reduced["step"]
                    evt["octave"] = int(reduced["octave"])
                    evt["alter"] = int(reduced["alter"])
                    evt["pitch_midi"] = int(reduced["pitch_midi"])
                    evt["confidence"] = 0.72
                    evt["source_note_count"] = int(reduced["source_note_count"])
                    if reduced["reduced_from_chord"]:
                        chord_reduction_applied += 1
                        evt["reduced_from_chord"] = True
                    if reduced["continuity_override_applied"]:
                        continuity_override_applied += 1
                        evt["continuity_override_applied"] = True
                    if reduced["source_note_count"] > 1:
                        chord_candidates_seen += 1
                    prev_note_midi = int(reduced["pitch_midi"])
                segment_events.append(evt)
        else:
            if omr_err:
                if omr_err != "disabled_by_env" and "neural_omr_model_unavailable" not in warnings:
                    warnings.append("neural_omr_model_unavailable")
            elif not neural_notes and "neural_omr_empty_sequence" not in warnings:
                warnings.append("neural_omr_empty_sequence")

            line_local = _fallback_staff_line_ys(rectified.shape[0])
            note_boxes = _detect_noteheads(rectified)
            for x, y, bw, bh in note_boxes[:_MAX_EVENTS_PER_STAFF]:
                center_y = y + bh // 2
                midi = _y_to_midi(center_y, line_local, clef)
                step, octave, alter = _midi_to_step_octave_alter(midi)
                segment_events.append(
                    {
                        "event_id": _next_event_id(),
                        "type": "note",
                        "step": step,
                        "octave": octave,
                        "alter": alter,
                        "pitch_midi": midi,
                        "duration_div": 1,
                        "onset_div": global_onset,
                        "confidence": 0.65,
                        "segment_order": seg_order,
                        "bbox": page_geom.work_rect_to_original_aabb(
                            int(staff.x0 + x),
                            int(staff.y0 + y),
                            int(bw),
                            int(bh),
                        ),
                    }
                )
                prev_note_midi = int(midi)
                global_onset += 1

        if not segment_events:
            raise SingleStaffAnalyzeError(
                "UNPROCESSABLE_SCORE",
                "No note or rest events could be extracted from this image.",
                {"segment_order": seg_order},
            )

        start_index = len(melody_events)
        melody_events.extend(segment_events)
        end_index = len(melody_events) - 1
        segment_map.append(
            {
                "segment_id": f"seg{seg_order}",
                "order": seg_order,
                "filename": image_meta["filename"],
                "width": image_meta["width"],
                "height": image_meta["height"],
                "event_index_range": {"start": start_index, "end": end_index},
            }
        )

    if len(decoded_images) > 1:
        warnings.append("line_break_between_images")
    if any_neural and "timing_from_pixel_gaps_heuristic" not in warnings:
        warnings.append("timing_from_pixel_gaps_heuristic")
    if chord_reduction_applied > 0 and "chord_reduction_applied_post_parse" not in warnings:
        warnings.append("chord_reduction_applied_post_parse")
    if continuity_override_applied > 0 and "top_note_continuity_override_applied" not in warnings:
        warnings.append("top_note_continuity_override_applied")

    resp_score_context = {
        "clef": score_context["clef"],
        "key_signature": dict(score_context["key_signature"]),
        "time_signature": time_signature,
        "tempo_bpm_reference": tempo_ref,
        "divisions": divisions,
        "source": "client",
    }

    meta: Dict[str, Any] = {
        "pipeline_mode": "single_staff_v1",
        "preprocess": {"segments": preprocess_segments},
        "reduction": {
            "rule": "top_note_max_two_parts",
            "chord_candidates_seen": int(chord_candidates_seen),
            "chord_reduction_applied": int(chord_reduction_applied),
            "continuity_override_applied": int(continuity_override_applied),
        },
    }
    if return_debug and debug_segments:
        meta["debug"] = {"omr_lilypond_by_segment": debug_segments}

    return {
        "source": {
            "total_images": len(decoded_images),
            "filenames": [m["filename"] for m in decoded_images],
        },
        "score_context": resp_score_context,
        "timeline": {
            "divisions": divisions,
            "time_signature": time_signature,
            "tempo_bpm_reference": tempo_ref,
        },
        "melody": {
            "voice_id": "melody1",
            "reduction_rule": "top_note_max_two_parts",
            "events": melody_events,
        },
        "segment_map": segment_map,
        "warnings": warnings,
        "meta": meta,
    }


@app.post("/analyze")
async def analyze(
    score_context: str = Form(...),
    file: Optional[UploadFile] = File(default=None),
    files: Optional[List[UploadFile]] = File(default=None),
    options: Optional[str] = Form(default=None),
) -> Any:
    ctx = _parse_score_context(score_context)
    options_parsed = _parse_options(options)
    requested_files = files if files else ([file] if file else [])
    if not requested_files:
        raise HTTPException(status_code=400, detail="Either 'file' or 'files' is required.")

    decoded_images = [_validate_and_decode_image(item) for item in requested_files]
    total_bytes = sum(item["bytes"] for item in decoded_images)
    if total_bytes > MAX_TOTAL_BYTES:
        raise HTTPException(status_code=413, detail="Total upload size exceeds 30MB limit.")

    try:
        body = build_analyze_response_v1(
            decoded_images,
            ctx,
            return_debug=bool(options_parsed["return_debug"]),
        )
    except SingleStaffAnalyzeError as exc:
        return JSONResponse(
            status_code=422,
            content={
                "error": {
                    "code": exc.code,
                    "message": exc.message,
                    "details": exc.details,
                }
            },
        )

    body["request_id"] = str(uuid.uuid4())
    return body