import generator.validation
import copy
import fractions
import pathlib
import unittest

import generator.arrange.generate
import generator.arrange.pickup
import generator.model
import generator.musicxml
import generator.parts
import tests.score_events
import generator.score
import tests.test_musicxml
import tests.test_parts

def source(pickup, body=None, *, implicit=True, divisions=1):
  body = body or tests.test_parts.xml_harmony() + tests.test_musicxml.xml_note(duration=4 * divisions)
  xml = tests.test_musicxml.xml_score(pickup, implicit=implicit, divisions=divisions)
  return tests.test_musicxml.parse(xml.replace('</measure></part>', '</measure><measure number="2">' + body + '</measure></part>'))

class ArrangementPickupTests(unittest.TestCase):
  def test_catalog_body_lengths_and_part_timelines(self):
    for path in sorted(pathlib.Path('data/songs').glob('*.xml')):
      document = generator.musicxml.load_musicxml(path)
      before = copy.deepcopy(document)
      prepared = generator.parts.prepare_document(document)
      timelines = []
      for part in generator.parts.PART_TYPES:
        with self.subTest(song=path.stem, part=part):
          result = generator.arrange.generate.arrange_document(document, part)
          generator.validation.validate_document(result)
          measures = result['staves'][0]['measures']
          self.assertEqual(len(measures), 171)
          self.assertEqual([m['number'] for m in measures], [str(i) for i in range(1, 172)])
          self.assertEqual([m['number'] for m in measures if m.get('section_label')], ['1', '9', '41', '73', '105', '137', '169'])
          self.assertEqual(sum(generator.model.fraction(m['duration']) for m in measures), 684)
          timelines.append([m['duration'] for m in measures])
          if path.stem == 'autumn_leaves' and part in ('vocal', 'electric'):
            for boundary in ([8, 136] if part == 'vocal' else [40]):
              notes = [n for n in generator.score.measure_notes(measures[boundary - 1]) if generator.model.fraction(n['onset']) >= 1 and n['pitches']]
              self.assertEqual([generator.model.fraction(n['onset']) for n in notes], [1, 2, 3])
              self.assertEqual([n['pitches'][-1]['letter'] for n in notes], ['d', 'e', 'f'])
              self.assertTrue(all(n['pitches'][-1]['octave'] == (5 if part == 'electric' else 4) for n in notes))
              self.assertFalse(measures[boundary - 1]['attributes'])
      self.assertTrue(all(timeline == timelines[0] for timeline in timelines))
      self.assertEqual(document, before)
  def test_unmarked_leading_rests_are_not_pickups(self):
    opening = '<note><rest/><duration>1</duration><type>quarter</type></note>' + tests.test_musicxml.xml_note(duration=3, kind='half', extra='<dot/>')
    document = source(opening, implicit=False)
    result = generator.arrange.generate.arrange_document(document, 'vocal')
    self.assertEqual(len(result['staves'][0]['measures']), 15)
    self.assertEqual(generator.score.body_start_measure_index(document['staves'][0]['measures']), 0)
  def test_short_pickup_ties_and_octave_spans_survive_each_melody_entry(self):
    document = source(tests.test_musicxml.xml_note(duration=1, kind='quarter', extra='<tie type="start"/>'), tests.test_musicxml.xml_note(extra='<tie type="stop"/>'))
    document['staves'][0]['octave_shifts'] = [{'id': 'pickup-shift', 'start': generator.model.rational(0), 'end': generator.model.rational(1), 'octaves': -1}]
    for part, boundaries in (('vocal', (1, 5)), ('electric', (2,))):
      with self.subTest(part=part):
        result = generator.arrange.generate._arrange_document(generator.parts.prepare_document(document), part, owned={'melody', 'vocal'})
        generator.validation.validate_document(result)
        measures = result['staves'][0]['measures']
        self.assertEqual(len(measures), 9)
        for boundary in boundaries:
          pickup = generator.score.measure_notes(measures[boundary - 1])[-1]
          downbeat = generator.score.measure_notes(measures[boundary])[0]
          self.assertEqual(generator.model.fraction(pickup['onset']), 3)
          self.assertEqual(pickup['pitches'][0]['tie_to'], downbeat['pitches'][0]['id'])
          self.assertEqual(downbeat['pitches'][0]['tie_from'], pickup['pitches'][0]['id'])
          self.assertTrue(any(generator.model.fraction(s['start']) == boundary * 4 - 1 and generator.model.fraction(s['end']) == boundary * 4 for s in result['staves'][0]['octave_shifts']))
        self.assertTrue(tests.score_events.timeline(result)['events'])
  def test_pickup_truncates_outgoing_music_without_discarding_earlier_notes(self):
    opening = '<note><rest/><duration>1</duration><type>quarter</type></note>' + tests.test_musicxml.xml_note('D', 3, 'half', extra='<dot/><lyric><text>doo</text></lyric>')
    document = source(opening)
    pickup, destination = copy.deepcopy(document['staves'][0]['measures'])
    before = copy.deepcopy(destination['voices'][0]['notes'][0])
    start = generator.arrange.pickup.splice(destination, pickup)
    self.assertEqual(start, 1)
    notes = generator.score.measure_notes(destination)
    self.assertEqual(len(notes), 2)
    self.assertEqual(notes[0]['id'], before['id'])
    self.assertEqual(notes[0]['pitches'][0]['id'], before['pitches'][0]['id'])
    self.assertEqual(generator.model.fraction(notes[0]['duration']), 1)
    self.assertEqual(notes[0]['written']['type'], 'quarter')
    self.assertEqual(notes[1]['lyrics'][0]['syllables'][0]['text'], 'doo')
    self.assertEqual([generator.model.fraction(n['onset']) for n in notes], [0, 1])
    document['staves'][0]['measures'] = [destination]
    generator.validation.validate_document(document)
  def test_pickup_harmonies_are_planned_for_every_part_at_the_same_offset(self):
    document = source(tests.test_parts.xml_harmony('F') + tests.test_musicxml.xml_note('F', 1, 'quarter'))
    prepared = generator.parts.prepare_document(document)
    for part in generator.parts.PART_TYPES:
      result = generator.arrange.generate.arrange_document(document, part)
      generator.validation.validate_document(result)
      for index in (0, 1, 4):
        chords = result['staves'][0]['measures'][index]['harmonies']
        self.assertEqual([(generator.model.fraction(h['onset']), h['root']['letter']) for h in chords], [(0, 'c'), (3, 'f')])
  def test_pickup_tuplets_keep_unique_identities_across_iterations(self):
    notes = ''
    for index, pitch in enumerate(('C', 'D', 'E')):
      tuplet = '<notations><tuplet type="start"/></notations>' if index == 0 else '<notations><tuplet type="stop"/></notations>' if index == 2 else ''
      notes += tests.test_musicxml.xml_note(pitch, 1, 'eighth', extra='<time-modification><actual-notes>3</actual-notes><normal-notes>2</normal-notes><normal-type>eighth</normal-type></time-modification>' + tuplet)
    document = source(notes, divisions=3)
    for part in ('vocal', 'electric'):
      result = generator.arrange.generate.arrange_document(document, part)
      generator.validation.validate_document(result)
      spans = result['staves'][0]['tuplets']
      self.assertEqual(len(spans), 2 if part == 'vocal' else 1)
      self.assertEqual(len({s['id'] for s in spans}), len(spans))
      for span in spans:
        members = [n for m in result['staves'][0]['measures'] for n in generator.score.measure_notes(m) if span['id'] in n.get('tuplet_ids', [])]
        self.assertEqual(len(members), 3)
        self.assertEqual([generator.model.fraction(n['onset']) for n in members], [3, fractions.Fraction(10, 3), fractions.Fraction(11, 3)])
  def test_overlong_pickup_is_rejected_with_measure_context(self):
    document = source(tests.test_musicxml.xml_note(duration=4), '<attributes><time><beats>3</beats><beat-type>4</beat-type></time></attributes>' + tests.test_musicxml.xml_note(duration=3, kind='half', extra='<dot/>'))
    with self.assertRaisesRegex(ValueError, 'measure 2: pickup is longer'):
      generator.arrange.generate.arrange_document(document, 'vocal')