import copy
from dataclasses import replace
from fractions import Fraction
from pathlib import Path
import unittest

from generator.engraving import EngravingError, Fonts, Style, layout_score, rasterize_page
from generator.engraving.audit import audit_layout
from generator.engraving.notation import ScoreIndex
from generator.model import NOTE_TYPES, TYPE_LENGTHS, fraction, rational
from generator.musicxml import load_musicxml
from generator.practice.generate import practice_document
from tests.test_musicxml import notes, parse, xml_note, xml_score

class EngravingTests(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
    cls.fonts = Fonts(Path('package/assets/fonts/petaluma'))
  def layout(self, source, style=None):
    before = copy.deepcopy(source)
    layout = layout_score(source, style or Style(), self.fonts)
    self.assertEqual(source, before)
    return layout
  def primitives(self, layout, role):
    return [p for page in layout.pages for p in page.primitives if p.role == role]
  def test_notation_fixture_visual_golden(self):
    from PIL import Image, ImageChops
    layout = self.layout(load_musicxml(Path('tests/fixtures/notation.xml')))
    self.assertEqual(len(layout.pages), 1)
    actual = rasterize_page(layout.pages[0], self.fonts)
    with Image.open('tests/fixtures/notation-python.png') as expected:
      self.assertIsNone(ImageChops.difference(actual, expected.convert('RGB')).getbbox())
  def test_accidental_work_is_bounded_by_local_event_count(self):
    class Unscannable(list):
      def __iter__(self):
        raise AssertionError('whole-score event scan')
    class Counted(list):
      visited = 0
      def __iter__(self):
        for value in super().__iter__():
          Counted.visited += 1
          yield value
    for count in (100, 200, 400):
      content = ''.join(xml_note(duration=1, kind='quarter', alter=i % 2) for i in range(4))
      extra = ''.join(f'<measure number="{i + 1}">{content}</measure>' for i in range(1, count))
      document = parse(xml_score(content).replace('</part></score-partwise>', extra + '</part></score-partwise>'))
      index = ScoreIndex(document)
      expected = index.accidentals
      index.events = Unscannable(index.events)
      index.events_by_onset = {time: Counted(events) for time, events in index.events_by_onset.items()}
      Counted.visited = 0
      self.assertEqual(index.accidental_state(), expected)
      self.assertEqual(Counted.visited, count * 4)
  def test_indexed_columns_preserve_sustains_and_timed_annotations(self):
    from generator.engraving.layout import column
    class Unscannable(list):
      def __iter__(self):
        raise AssertionError('whole-score scan from column layout')
    content = xml_note(extra='<voice>held</voice>') + '<backup><duration>4</duration></backup>'
    content += ''.join(xml_note('G', 1, 'quarter', extra='<voice>moving</voice>') for _ in range(4))
    content += '<harmony><root><root-step>F</root-step></root><kind>major</kind><offset>-2</offset></harmony>'
    content += '<direction><direction-type><metronome><beat-unit>quarter</beat-unit><per-minute>90</per-minute></metronome></direction-type><offset>-2</offset></direction>'
    index = ScoreIndex(parse(xml_score(content)))
    expected = column(index, Fraction(2), Fraction(2), Fraction(4), self.fonts, Style())
    index.events = Unscannable(index.events)
    index.measures = Unscannable(index.measures)
    actual = column(index, Fraction(2), Fraction(2), Fraction(4), self.fonts, Style())
    self.assertEqual(actual, expected)
    self.assertEqual(len(index.events_at(Fraction(2), include_sustained=True)), 2)
    self.assertTrue(any(p.role == 'harmony' for p in actual[0]))
    self.assertEqual(len(index.events_at(Fraction(4), include_sustained=True)), 0)
  def test_every_written_note_and_rest_value(self):
    for kind in NOTE_TYPES:
      for dots in (0, 1, 3):
        duration = TYPE_LENGTHS[kind] * (2 - Fraction(1, 2 ** dots))
        for rest in (False, True):
          with self.subTest(kind=kind, dots=dots, rest=rest):
            content = xml_note(duration=duration.numerator, kind=kind, extra='<dot/>' * dots)
            if rest:
              content = f'<note><rest/><duration>{duration.numerator}</duration><type>{kind}</type>' + '<dot/>' * dots + '</note>'
            source = parse(xml_score(content, divisions=duration.denominator, implicit=True))
            layout = self.layout(source)
            self.assertEqual(len(self.primitives(layout, 'dot')), dots)
            self.assertEqual(len(self.primitives(layout, 'rest' if rest else 'head')), 1)
  def test_sixteen_dots_are_not_truncated(self):
    duration = TYPE_LENGTHS['quarter'] * (2 - Fraction(1, 2 ** 16))
    source = parse(xml_score(xml_note(duration=duration.numerator, kind='quarter', extra='<dot/>' * 16), divisions=duration.denominator, implicit=True))
    self.assertEqual(len(self.primitives(self.layout(source), 'dot')), 16)
  def test_arbitrary_voice_ids_and_more_than_four_voices(self):
    content = ''
    for i in range(6):
      if i:
        content += '<backup><duration>4</duration></backup>'
      content += xml_note(extra=f'<voice>part-{i}</voice>')
    source = parse(xml_score(content))
    self.assertEqual(len(self.primitives(self.layout(source), 'head')), 6)
  def test_microtonal_numeric_fallback_and_explicit_missing_glyph(self):
    source = parse(xml_score(xml_note(alter='0.1')))
    self.assertEqual(self.primitives(self.layout(source), 'accidental')[0].value, '(+1/10 st)')
    notes(source)[0]['pitches'][0]['accidental_display'] = {'name': 'other', 'smufl': 'missingVampGlyph'}
    with self.assertRaisesRegex(EngravingError, 'measure 1, voice 1.*missingVampGlyph'):
      self.layout(source)
  def test_accidentals_are_shared_by_simultaneous_voices_without_losing_conflicts(self):
    content = xml_note(duration=1, kind='quarter', alter=1, extra='<voice>a</voice>')
    content += xml_note(duration=1, kind='quarter', alter=1, extra='<voice>a</voice>')
    content += xml_note(duration=2, kind='half', extra='<voice>a</voice>')
    source = parse(xml_score(content))
    self.assertEqual([p.value for p in self.primitives(self.layout(source), 'accidental')], ['accidentalSharp', 'accidentalNatural'])
    conflict = parse(xml_score(xml_note(alter=1, extra='<voice>a</voice>') + '<backup><duration>4</duration></backup>' + xml_note(extra='<voice>b</voice>')))
    self.assertEqual(len(self.primitives(self.layout(conflict), 'accidental')), 2)
  def test_dense_measure_splits_exactly_without_modifying_music(self):
    content = ''.join(xml_note('G', 1, '64th', extra='<voice>fast</voice>') for _ in range(64))
    content += '<backup><duration>64</duration></backup>' + xml_note('C', 64, 'whole', extra='<voice>held</voice>')
    source = parse(xml_score(content, divisions=16))
    layout = self.layout(source)
    systems = [s for p in layout.pages for s in p.systems]
    self.assertGreater(len(systems), 1)
    held = source['staves'][0]['measures'][0]['voices'][1]['notes'][0]['pitches'][0]['id']
    pieces = [f for s in systems for f in s.fragments if f.source == held]
    self.assertEqual(sum(f.length() for f in pieces), 4)
    self.assertGreater(len(pieces), 1)
    self.assertTrue(self.primitives(layout, 'tie'))
    for page in layout.pages:
      self.assertEqual(rasterize_page(page, self.fonts).size, (976, 484))
  def test_hidden_pitch_does_not_suppress_a_visible_accidental(self):
    source = parse(xml_score(xml_note(duration=2, kind='half', alter=1) + xml_note(duration=2, kind='half', alter=1)))
    notes(source)[0]['hidden'] = True
    self.assertEqual([p.value for p in self.primitives(self.layout(source), 'accidental')], ['accidentalSharp'])
  def test_octave_shift_changes_display_not_pitch(self):
    source = parse(xml_score(xml_note(octave=5)))
    original = ScoreIndex(source).pitch_position(notes(source)[0]['pitches'][0], Fraction())
    source['staves'][0]['octave_shifts'] = [{'id': 'octave', 'start': rational(0), 'end': rational(4), 'octaves': -1}]
    index = ScoreIndex(source)
    self.assertEqual(index.pitch_position(notes(source)[0]['pitches'][0], Fraction()), original - 7)
    self.assertTrue(any(p.value == '8va' for p in self.primitives(self.layout(source), 'octave-shift')))
  def test_wide_pitch_range_can_break_inside_a_measure_vertically(self):
    source = parse(xml_score(xml_note(duration=2, kind='half', octave=10) + xml_note(duration=2, kind='half', octave=-1)))
    layout = self.layout(source)
    self.assertEqual(len(layout.pages), 2)
    self.assertEqual(len(self.primitives(layout, 'head')), 2)
  def test_cue_grace_and_tremolo_are_drawn(self):
    grace = '<note><grace slash="yes"/><pitch><step>D</step><octave>4</octave></pitch><type>16th</type></note>'
    source = parse(xml_score(grace + xml_note(duration=4, kind='whole', prefix='<cue/>')))
    layout = self.layout(source)
    self.assertEqual(len(self.primitives(layout, 'head')), 2)
    self.assertEqual(len(self.primitives(layout, 'grace-slash')), 1)
    source = parse(xml_score(xml_note(duration=2, kind='half', extra='<notations><ornaments><tremolo type="single">8</tremolo></ornaments></notations>'), implicit=True))
    self.assertEqual(len(self.primitives(self.layout(source), 'tremolo')), 8)
  def test_percussion_and_tab_use_explicit_display_positions(self):
    percussion = '<note><unpitched><display-step>C</display-step><display-octave>5</display-octave></unpitched><duration>4</duration><type>whole</type><notehead>x</notehead></note>'
    source = parse(xml_score(percussion, attributes='<clef><sign>percussion</sign></clef>'))
    self.assertEqual(self.primitives(self.layout(source), 'head')[0].value, 'noteheadXWhole')
    tab = practice_document(parse(xml_score(xml_note('C', octave=7))), 'electric')
    self.assertEqual(self.primitives(self.layout(tab), 'tab')[0].value, '20')
  def test_generated_tab_shares_columns_and_systems_without_duplicating_music(self):
    content = ''.join(xml_note('G', 1, '64th') for _ in range(64))
    source = practice_document(parse(xml_score(content, divisions=16)), 'electric')
    layout = self.layout(source)
    systems = [s for page in layout.pages for s in page.systems]
    self.assertGreater(len(systems), 1)
    self.assertEqual(sum(len(s.fragments) for s in systems), 64)
    for system in systems:
      heads = [p for p in system.primitives if p.role == 'head']
      tabs = [p for p in system.primitives if p.role == 'tab']
      self.assertEqual(len(heads), len(tabs))
      self.assertEqual(len({p.points[0][1] for p in system.primitives if p.role == 'tab-staff'}), 6)
      self.assertLess(max(p.box.bottom for p in system.primitives if not p.role.startswith('tab')), min(p.box.top for p in system.primitives if p.role.startswith('tab')))
      self.assertEqual({p.sources[-1] for p in heads}, {p.sources[-1] for p in tabs})
      self.assertTrue(all(any(abs(head.points[0][0] - tab.points[0][0]) < 0.01 for head in heads if head.sources[-1] == tab.sources[-1]) for tab in tabs))
    from generator.playback import compile_timeline
    self.assertEqual(len(compile_timeline(source)['events']), 64)
  def test_generated_tab_ties_and_barline_styles(self):
    content = xml_note('E', 2, 'half', extra='<tie type="start"/>') + xml_note('E', 2, 'half', extra='<tie type="stop"/>') + '<barline location="right"><bar-style>dashed</bar-style></barline>'
    source = practice_document(parse(xml_score(content)), 'electric')
    layout = self.layout(source)
    self.assertEqual(len(self.primitives(layout, 'tab-tie')), 1)
    self.assertEqual(len(self.primitives(layout, 'barline')), len(self.primitives(layout, 'tab-barline')))
    for system in (s for page in layout.pages for s in page.systems):
      baselines = {role: max(p.points[0][1] for p in system.primitives if p.role == role) for role in ('staff', 'tab-staff')}
      notation = [p for p in system.primitives if p.role == 'barline']
      tabs = [p for p in system.primitives if p.role == 'tab-barline']
      for bar, tab in zip(notation, tabs):
        for (x1, y1), (x2, y2) in zip(bar.points, tab.points):
          self.assertAlmostEqual(x1, x2)
          self.assertAlmostEqual((y1 - baselines['staff']) * 5 / 4, y2 - baselines['tab-staff'])
  def test_adjacent_string_tab_numerals_do_not_overlap(self):
    from generator.guitar import place_guitar
    content = ''.join(xml_note(step, octave=octave, prefix='<chord/>' if index else '') for index, (step, octave) in enumerate((('G', 3), ('C', 4), ('E', 4), ('G', 4), ('C', 5), ('E', 5))))
    source = parse(xml_score(content))
    place_guitar(source)
    layout = self.layout(source)
    tabs = self.primitives(layout, 'tab')
    self.assertEqual(len(tabs), 6)
    for i, first in enumerate(tabs):
      for second in tabs[i + 1:]:
        self.assertFalse(first.box.overlaps(second.box, 0.1))
  def test_lyrics_hyphens_and_extenders(self):
    begin = '<lyric number="1"><syllabic>begin</syllabic><text>hel</text></lyric>'
    end = '<lyric number="1"><syllabic>end</syllabic><text>lo</text><extend type="start"/></lyric>'
    source = parse(xml_score(xml_note(duration=1, kind='quarter', extra=begin) + xml_note(duration=1, kind='quarter', extra=end) + xml_note(duration=2, kind='half')))
    layout = self.layout(source)
    self.assertTrue(self.primitives(layout, 'lyric-hyphen'))
    self.assertTrue(self.primitives(layout, 'lyric-extension'))
  def test_nested_and_cross_bar_tuplets_keep_span_ownership(self):
    from xml.etree import ElementTree
    modulation = '<time-modification><actual-notes>3</actual-notes><normal-notes>2</normal-notes></time-modification>'
    content = ''.join(xml_note(duration=2, kind='quarter', extra=modulation) for _ in range(6))
    tree = ElementTree.fromstring(xml_score(content, divisions=3))
    second = ElementTree.SubElement(tree.find('part'), 'measure', number='2')
    for element in ElementTree.fromstring(xml_score(content, divisions=3)).findall('part/measure/note'):
      second.append(element)
    source = parse(ElementTree.tostring(tree, encoding='unicode'))
    events = notes(source) + notes(source, 1)
    spans = [{'id': 'outer', 'voice': '1', 'start': events[0]['id'], 'end': events[-1]['id'], 'note_count': 3, 'space_count': 2, 'depth': 0, 'bracket': 'yes', 'show_number': 'both', 'show_type': 'both', 'line_shape': 'curved'}, {'id': 'inner', 'voice': '1', 'start': events[1]['id'], 'end': events[3]['id'], 'note_count': 3, 'space_count': 2, 'depth': 1, 'bracket': 'yes', 'show_number': 'actual'}]
    source['staves'][0]['tuplets'] = spans
    for i, event in enumerate(events):
      event['tuplet_ids'] = ['outer'] + (['inner'] if 1 <= i <= 3 else [])
    layout = self.layout(source, Style(width=600))
    self.assertGreater(len([s for p in layout.pages for s in p.systems]), 1)
    self.assertEqual({p.sources[0] for p in self.primitives(layout, 'tuplet-bracket')}, {'outer', 'inner'})
  def test_paired_tremolo_and_double_stems(self):
    modulation = '<time-modification><actual-notes>2</actual-notes><normal-notes>1</normal-notes></time-modification>'
    content = ''.join(xml_note(step, 2, 'whole', extra=modulation + f'<notations><ornaments><tremolo type="{action}">3</tremolo></ornaments></notations>') for step, action in [('C', 'start'), ('G', 'stop')])
    self.assertEqual(len(self.primitives(self.layout(parse(xml_score(content))), 'tremolo')), 3)
    source = parse(xml_score(xml_note(duration=1, kind='eighth', extra='<stem>double</stem>'), divisions=2, implicit=True))
    layout = self.layout(source)
    self.assertEqual(len(self.primitives(layout, 'stem')), 2)
    self.assertEqual(len(self.primitives(layout, 'flag')), 2)
  def test_meter_symbols_and_timed_attributes(self):
    for symbol in ('common', 'cut', 'single-number', 'note', 'dotted-note', 'normal'):
      source = parse(xml_score(xml_note()).replace('<time>', f'<time symbol="{symbol}">'))
      with self.subTest(symbol=symbol):
        self.assertTrue(self.primitives(self.layout(source), 'meter'))
    change = '<attributes><key><fifths>-3</fifths></key><clef><sign>C</sign><line>3</line></clef></attributes>'
    source = parse(xml_score(xml_note(duration=2, kind='half') + change + xml_note(duration=2, kind='half')))
    layout = self.layout(source)
    self.assertEqual([p.value for p in self.primitives(layout, 'clef')], ['gClef', 'cClef'])
  def test_lyric_extender_continues_on_systems_without_syllables(self):
    lyrics = '<lyric number="1"><text>Ah</text><extend type="start"/></lyric>'
    content = ''.join(xml_note('G', 1, '64th', extra=lyrics if i == 0 else '') for i in range(64))
    source = parse(xml_score(content, divisions=16))
    layout = self.layout(source)
    systems = [s for p in layout.pages for s in p.systems]
    self.assertGreater(len(systems), 1)
    self.assertTrue(all(any(p.role == 'lyric-extension' for p in s.primitives) for s in systems))
  def test_audit_catches_removed_pitch_and_accidental_primitives(self):
    source = parse(xml_score(xml_note(alter=1)))
    layout = self.layout(source)
    for role in ('head', 'accidental'):
      pages = tuple(replace(p, primitives=tuple(g for g in p.primitives if g.role != role)) for p in layout.pages)
      with self.assertRaisesRegex(EngravingError, 'no drawing primitive|accidental was omitted'):
        audit_layout(source, replace(layout, pages=pages), self.fonts)
  def test_fingerprint_includes_style_and_fixed_overflow_errors(self):
    source = parse(xml_score(xml_note()))
    a = self.layout(source)
    b = self.layout(source, Style(column_gap=14))
    self.assertNotEqual(a.fingerprint, b.fingerprint)
    with self.assertRaisesRegex(EngravingError, 'page width'):
      self.layout(source, Style(width=100))