#!/usr/bin/env python3
"""
grand_staff_merge.py

Post-process homr's dual-part MusicXML output into a proper piano grand staff.

What it does:
1. Merges P1 (treble/G clef) and P2 (bass/F clef) into a single Part
2. Adds <staves>2</staves> and <bracket type="brace"/> for piano grand staff
3. Bass stays on staff=1/voice=1, treble moves to staff=2/voice=2
4. Removes ALL <print> elements (MuseScore ignores them and causes extra systems)
5. Preserves note durations, pitches, dynamics — only structural changes

Usage:
    python grand_staff_merge.py input.musicxml [output.musicxml]
"""

import sys
import xml.etree.ElementTree as ET

# Register namespaces to preserve them in output
ET.register_namespace("m", "http://www.musescore.org/ns/mscore")
ET.register_namespace("xlink", "http://www.w3.org/1999/xlink")
ET.register_namespace(
    "mx", "http://schemas.openxmlformats.org/markup-compatibility/2006"
)


def remove_print_elements(part_elem):
    """Remove all <print> elements from a part (they cause extra system breaks in MuseScore)."""
    removed = 0
    for print_elem in part_elem.findall(".//print"):
        for parent in part_elem.iter():
            if print_elem in list(parent):
                parent.remove(print_elem)
                removed += 1
                break
    return removed


def make_attributes_grand_staff(attrs_elem):
    """Modify attributes element for grand staff: add staves, bracket, correct clefs."""
    # Remove existing clef (will add both below)
    existing_clef = attrs_elem.find("clef")
    if existing_clef is not None:
        attrs_elem.remove(existing_clef)

    # Add <staves>2</staves> after <divisions>
    divisions = attrs_elem.find("divisions")
    staves_elem = ET.Element("staves")
    staves_elem.text = "2"
    if divisions is not None:
        idx = list(attrs_elem).index(divisions) + 1
        attrs_elem.insert(idx, staves_elem)
    else:
        attrs_elem.append(staves_elem)

    # MusicXML clef order = visual top-to-bottom order
    # Grand staff: G clef (treble) is visually at the TOP, so it comes FIRST in XML
    # F clef (bass) is visually at the BOTTOM, so it comes SECOND
    #
    # CRITICAL: Staff numbering in MuseScore grand staff:
    #   staff=1 = upper staff = treble (G clef)
    #   staff=2 = lower staff = bass (F clef)
    treble_clef = ET.SubElement(attrs_elem, "clef")
    treble_sign = ET.SubElement(treble_clef, "sign")
    treble_sign.text = "G"
    treble_line = ET.SubElement(treble_clef, "line")
    treble_line.text = "2"
    treble_staff_elem = ET.SubElement(treble_clef, "staff")
    treble_staff_elem.text = "1"  # Upper staff = staff=1

    bass_clef = ET.SubElement(attrs_elem, "clef")
    bass_sign = ET.SubElement(bass_clef, "sign")
    bass_sign.text = "F"
    bass_line = ET.SubElement(bass_clef, "line")
    bass_line.text = "4"
    bass_staff_elem = ET.SubElement(bass_clef, "staff")
    bass_staff_elem.text = "2"  # Lower staff = staff=2

    # Add bracket for piano brace (at end of attributes, after clefs)
    bracket_elem = ET.SubElement(attrs_elem, "bracket")
    bracket_elem.set("type", "brace")
    bracket_elem.set("number", "1")
    bracket_elem.text = ""


def make_attributes_no_print(attrs_elem):
    """Keep existing attributes but ensure no print element sibling."""
    pass  # print elements handled at part level


def process_part_measure(measure_elem, voice_map):
    """
    Process a measure from P1 or P2.
    voice_map: {'staff': 1 or 2, 'voice': 1 or 2}
    """
    # Remove <print> elements from this measure
    print_elem = measure_elem.find("print")
    if print_elem is not None:
        measure_elem.remove(print_elem)

    for note in measure_elem.findall("note"):
        # Update staff attribute
        staff_elem = note.find("staff")
        if staff_elem is not None:
            staff_elem.text = str(voice_map["staff"])

        # Update voice attribute
        voice_elem = note.find("voice")
        if voice_elem is not None:
            voice_elem.text = str(voice_map["voice"])


def merge_parts_to_grand_staff(xml_path, output_path=None):
    """Main function: merge homr dual-part output into grand staff."""
    if output_path is None:
        output_path = xml_path.replace(".musicxml", "_grandstaff.musicxml")

    tree = ET.parse(xml_path)
    root = tree.getroot()

    # Find part-list and part elements
    part_list = root.find("part-list")
    parts = root.findall("part")

    if len(parts) < 2:
        print(f"Warning: Expected 2 parts, found {len(parts)}. Nothing to merge.")
        return

    p1 = parts[0]  # treble (G clef)
    p2 = parts[1]  # bass (F clef)

    # Determine which is which by clef sign
    def get_clef(part):
        for m in part.findall("measure"):
            clef = m.find(".//clef")
            if clef is not None:
                sign = clef.find("sign")
                if sign is not None:
                    return sign.text
        return None

    p1_clef = get_clef(p1)
    p2_clef = get_clef(p2)

    print(f"P1 clef: {p1_clef}, P2 clef: {p2_clef}")

    # Identify treble and bass parts
    if p1_clef == "G" and p2_clef == "F":
        treble_part = p1
        bass_part = p2
        print("Detected: P1=treble, P2=bass")
    elif p1_clef == "F" and p2_clef == "G":
        treble_part = p2
        bass_part = p1
        print("Detected: P1=bass, P2=treble")
    else:
        # Default: P1=treble, P2=bass (homr's typical output)
        treble_part = p1
        bass_part = p2
        print(f"Could not determine clefs definitively, defaulting: P1=treble, P2=bass")

    # === Step 1: Update part-list to single Part ===
    # Replace with single score-part
    new_part_list = ET.Element("part-list")
    score_part = ET.SubElement(new_part_list, "score-part")
    score_part.set("id", "P1")
    part_name = ET.SubElement(score_part, "part-name")
    part_name.text = "Piano"
    part_name.set("print-object", "yes")

    # Add midi-instrument
    midi_inst = ET.SubElement(score_part, "midi-instrument")
    midi_inst.set("id", "P1-I1")
    midi_ch = ET.SubElement(midi_inst, "midi-channel")
    midi_ch.text = "1"
    midi_prog = ET.SubElement(midi_inst, "midi-program")
    midi_prog.text = "1"
    vol = ET.SubElement(midi_inst, "volume")
    vol.text = "100"
    pan = ET.SubElement(midi_inst, "pan")
    pan.text = "0"

    # Replace old part-list
    if part_list is not None:
        root.remove(part_list)
    root.insert(list(root).index(p1), new_part_list)

    # === Step 2: Update bass_part measure 1 attributes ===
    measures_bass = bass_part.findall("measure")
    if measures_bass:
        first_measure_bass = measures_bass[0]
        attrs = first_measure_bass.find("attributes")
        if attrs is not None:
            make_attributes_grand_staff(attrs)

    # Remove <print> from ALL measures in both parts
    for p in [bass_part, treble_part]:
        for m in p.findall("measure"):
            print_e = m.find("print")
            if print_e is not None:
                m.remove(print_e)

    # === Step 3: Change treble notes to staff=1 (upper), voice=1 ===
    for m in treble_part.findall("measure"):
        process_part_measure(m, {"staff": 1, "voice": 1})

    # === Step 4: Change bass notes to staff=2 (lower), voice=2 ===
    for m in bass_part.findall("measure"):
        process_part_measure(m, {"staff": 2, "voice": 2})

    # === Step 5: Interleave measures from both parts ===
    # Create a new combined part
    combined = ET.Element("part")
    combined.set("id", "P1")

    bass_measures = bass_part.findall("measure")
    treble_measures = treble_part.findall("measure")

    # Also remove print from treble first measure if it wasn't already removed
    if treble_measures:
        print_e = treble_measures[0].find("print")
        if print_e is not None:
            treble_measures[0].remove(print_e)

    # Add first measure of treble (with attributes) to combined first
    # But we already added grand staff attributes to bass first measure
    # So treble first measure's attributes should be removed
    treble_attrs_removed = 0
    if treble_measures:
        treble_first_attrs = treble_measures[0].find("attributes")
        if treble_first_attrs is not None:
            treble_measures[0].remove(treble_first_attrs)
            treble_attrs_removed += 1

    # For measures that appear in both, interleave notes
    # Strategy: for each measure number, combine notes from both parts
    bass_by_num = {m.get("number"): m for m in bass_measures}
    treble_by_num = {m.get("number"): m for m in treble_measures}

    all_measure_nums = sorted(
        set(list(bass_by_num.keys()) + list(treble_by_num.keys())),
        key=lambda x: int(x) if x is not None and x.isdigit() else 0,
    )

    for mnum in all_measure_nums:
        bass_m = bass_by_num.get(mnum)
        treble_m = treble_by_num.get(mnum)

        # Create new combined measure
        new_measure = ET.Element("measure")
        if mnum is not None:
            new_measure.set("number", mnum)

        # Copy attributes from bass (which has grand staff setup)
        if bass_m is not None:
            attrs = bass_m.find("attributes")
            if attrs is not None:
                new_measure.append(ET.fromstring(ET.tostring(attrs)))

        # Copy direction/direction-type from treble first measure (for tempo, etc.)
        if treble_m is not None and mnum == "1":
            for child in list(treble_m):
                if child.tag not in (
                    "note",
                    "backup",
                    "forward",
                    "attributes",
                    "barline",
                    "print",
                ):
                    new_measure.append(ET.fromstring(ET.tostring(child)))

        # Add notes from treble (staff=1, upper) — first in XML
        if treble_m is not None:
            for note in treble_m.findall("note"):
                new_measure.append(ET.fromstring(ET.tostring(note)))
            for backup in treble_m.findall("backup"):
                new_measure.append(ET.fromstring(ET.tostring(backup)))
            for fwd in treble_m.findall("forward"):
                new_measure.append(ET.fromstring(ET.tostring(fwd)))

        # Add notes from bass (staff=2, lower) — after treble
        if bass_m is not None:
            for note in bass_m.findall("note"):
                new_measure.append(ET.fromstring(ET.tostring(note)))
            for backup in bass_m.findall("backup"):
                new_measure.append(ET.fromstring(ET.tostring(backup)))
            for fwd in bass_m.findall("forward"):
                new_measure.append(ET.fromstring(ET.tostring(fwd)))
            for barline in bass_m.findall("barline"):
                new_measure.append(ET.fromstring(ET.tostring(barline)))

        combined.append(new_measure)

    # === Step 6: Replace original parts with combined ===
    root.remove(p1)
    if p2 != p1:
        try:
            root.remove(p2)
        except ValueError:
            pass
    root.append(combined)

    # === Step 7: Write output ===
    # ET.indent for pretty printing (Python 3.9+)
    try:
        ET.indent(root)
    except AttributeError:
        pass  # Python < 3.9

    tree.write(output_path, encoding="UTF-8", xml_declaration=True)

    print(f"\nGrand staff conversion complete!")
    print(f"   Input:  {xml_path}")
    print(f"   Output: {output_path}")
    print(f"   Removed <print> elements to prevent extra system breaks")
    print(f"   Added: <staves>2</staves>, bracket type=brace")
    print(f"   Treble: staff=1 (upper), voice=1 | Bass: staff=2 (lower), voice=2")
    print(f"   Measures: {len(all_measure_nums)}")

    return output_path


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    input_file = sys.argv[1]
    output_file = sys.argv[2] if len(sys.argv) > 2 else None

    merge_parts_to_grand_staff(input_file, output_file)