import xml.etree.ElementTree as ET
import os

# MusicXML filename to be processed
TARGET_FILENAME = 'GUjianpu.musicxml'

# Mapping from standard musical notes to Jianpu numbers
JIANPU_PITCH_MAP = {
    'C': '1', 'D': '2', 'E': '3', 'F': '4', 'G': '5', 'A': '6', 'B': '7'
}

# Accidentals (flats, sharps) mapping
ACCIDENTALS = {
    1: '♯',  # Sharp
    -1: '♭', # Flat
}

# Time values in MusicXML and their corresponding Jianpu notations (for underlines, dots, etc.)
DURATION_MAP = {
    960: 'half',  # Half note
    1920: 'whole',  # Whole note
    240: 'eighth',  # Eighth note
    120: 'sixteenth',  # Sixteenth note
}

def find_target_file():
    """Look for the target MusicXML file in the current working directory."""
    files_in_directory = os.listdir('.')
    if TARGET_FILENAME in files_in_directory:
        return TARGET_FILENAME
    else:
        print(f"Error: {TARGET_FILENAME} not found in the current directory.")
        return None

def convert_pitch_to_jianpu(pitch_element):
    """Convert a pitch element from MusicXML to Jianpu."""
    try:
        step = pitch_element.find('step').text
        octave = int(pitch_element.find('octave').text)
        alter = pitch_element.find('alter')

        # Convert step (C, D, E, etc.) to Jianpu number (1, 2, 3, etc.)
        jianpu_number = JIANPU_PITCH_MAP.get(step, '')

        # If there's an accidental (sharp/flat), add it to the Jianpu number
        if alter is not None:
            alter_value = int(alter.text)
            jianpu_number += ACCIDENTALS.get(alter_value, '')

        # Add octave marking (dot above or below)
        if octave < 4:
            jianpu_number = '·' * (4 - octave) + jianpu_number  # Dots below for lower octave
        elif octave > 4:
            jianpu_number = jianpu_number + '·' * (octave - 4)  # Dots above for higher octave

        return jianpu_number
    except Exception as e:
        print(f"Error processing pitch: {e}")
        return None

def convert_duration_to_jianpu(duration, dots=0):
    """Handle note duration and convert it to Jianpu equivalent (underlines, dots)."""
    try:
        jianpu_duration = 'normal'
        if duration in DURATION_MAP:
            jianpu_duration = DURATION_MAP[duration]

        # Handle dotted notes (附点音符)
        if dots > 0:
            jianpu_duration += ' ' + '·' * dots  # Add dots to the note

        return jianpu_duration
    except Exception as e:
        print(f"Error processing duration: {e}")
        return 'normal'

def handle_rest(note_element):
    """Handle rests and convert to '0' notation in Jianpu."""
    try:
        rest_element = note_element.find('rest')
        if rest_element is not None:
            return '0'  # Jianpu notation for rests
    except Exception as e:
        print(f"Error processing rest: {e}")
    return None

def handle_tie(note_element):
    """Handle tied notes, where the note continues across bar lines."""
    try:
        tie_start = note_element.find('tie[@type="start"]')
        tie_stop = note_element.find('tie[@type="stop"]')
        if tie_start is not None:
            return True  # Mark as tied note, but no need to duplicate the note
    except Exception as e:
        print(f"Error processing tie: {e}")
    return False

def handle_slur(note_element):
    """Handle slurs (连音线) between notes."""
    try:
        slur_start = note_element.find('notations/slur[@type="start"]')
        slur_stop = note_element.find('notations/slur[@type="stop"]')

        # Slur between notes (连音线)
        if slur_start is not None:
            return "start_slur"
        elif slur_stop is not None:
            return "end_slur"
    except Exception as e:
        print(f"Error processing slur: {e}")
    return None

def process_musicxml(input_file, output_file):
    """Parse MusicXML, convert to Jianpu, and save the modified file."""
    try:
        # Parse the MusicXML file
        tree = ET.parse(input_file)
        root = tree.getroot()

        # Define namespaces for MusicXML
        ns = {'score': 'http://www.musicxml.org/ns/score'}

        # Process all <note> elements
        for note in root.findall('.//score:note', namespaces=ns):
            try:
                # Handle rests
                rest_jianpu = handle_rest(note)
                if rest_jianpu:
                    staff_text = ET.Element('staff-text')
                    staff_text.text = rest_jianpu
                    note.append(staff_text)
                    continue

                # Handle pitch and duration conversion
                pitch = note.find('score:pitch', namespaces=ns)
                duration = note.find('score:duration', namespaces=ns)
                dots = note.find('score:dots', namespaces=ns)
                tie = handle_tie(note)
                slur = handle_slur(note)

                if pitch is not None:
                    # Convert pitch to Jianpu notation
                    jianpu_notation = convert_pitch_to_jianpu(pitch)

                    # Handle duration and dotted notes
                    jianpu_duration = convert_duration_to_jianpu(int(duration.text), int(dots.text) if dots is not None else 0)

                    # Append Jianpu notation and additional marks
                    jianpu_notation += f" {jianpu_duration}"
                    if tie:
                        jianpu_notation += ' (tie)'
                    if slur == "start_slur":
                        jianpu_notation += ' (slur start)'
                    elif slur == "end_slur":
                        jianpu_notation += ' (slur end)'

                    # Create a new XML element to store the Jianpu notation
                    staff_text = ET.Element('staff-text')
                    staff_text.text = jianpu_notation

                    # Insert the Jianpu notation element into the MusicXML tree
                    note.append(staff_text)

            except Exception as e:
                print(f"Error processing note: {e}. Skipping this note but continuing.")

        # Write the modified XML to the output file
        tree.write(output_file, encoding="UTF-8", xml_declaration=True)
        print(f"Jianpu-converted MusicXML saved as {output_file}")

    except ET.ParseError as e:
        print(f"Error parsing the MusicXML file: {e}")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")

def main():
    # Look for the target MusicXML file in the current directory
    input_file = find_target_file()
    if input_file:
        output_file = 'GUjianpu_converted.musicxml'

        # Process MusicXML to Jianpu
        process_musicxml(input_file, output_file)

if __name__ == "__main__":
    main()