import argparse
import copy
import fractions
import json
import pathlib
import tempfile
import unittest
import unittest.mock
import xml.etree.ElementTree

import generator.__main__
import generator.accompaniment
import generator.arrange.generate
import generator.model
import generator.musicxml
import generator.parts
import generator.pitch
import generator.playback
import tests.arrangement_fixtures
import tests.rhythm_fixtures
import tests.test_musicxml

def xml_harmony(root='C', kind='major', offset=0):
  return f'<harmony><root><root-step>{root}</root-step></root><kind>{kind}</kind><offset>{offset}</offset></harmony>'

class PartGenerationTests(unittest.TestCase):
  def assert_complete_measure(self, measure):
    cursor = fractions.Fraction()
    for note in measure['voices'][0]['notes']:
      self.assertEqual(generator.model.fraction(note['onset']), cursor)
      self.assertGreater(generator.model.fraction(note['duration']), 0)
      cursor += generator.model.fraction(note['duration'])
    self.assertEqual(cursor, generator.model.fraction(measure['duration']))
  def test_dense_melody_keeps_rhythmic_accompaniment_without_sustaining_across_changes(self):
    source = tests.test_musicxml.parse(tests.test_musicxml.xml_score(xml_harmony() + xml_harmony('F', offset=3) + ''.join(tests.test_musicxml.xml_note(duration=1, kind='eighth') for _ in range(8)), divisions=2))
    for result in (tests.arrangement_fixtures.section_document(source, 'acoustic', rhythms=tests.rhythm_fixtures.library()), generator.arrange.generate.arrange_document(source, 'electric', rhythms=tests.rhythm_fixtures.library())):
      generator.model.validate_document(result)
      for measure in result['staves'][0]['measures']:
        events = measure['voices'][0]['notes']
        if measure.get('section_label') in ('Rhythm Lead-in', 'Guitar Melody', 'Guitar Solo'):
          continue
        self.assert_complete_measure(measure)
        sounding = [n for n in events if n['pitches']]
        self.assertTrue(sounding)
        self.assertLess(len(sounding), 8)
        for note in sounding:
          onset = generator.model.fraction(note['onset'])
          if onset < fractions.Fraction(3, 2):
            self.assertLessEqual(onset + generator.model.fraction(note['duration']), fractions.Fraction(3, 2))
          self.assertLessEqual({generator.model.pitch_number(p) % 12 for p in note['pitches']}, {0, 4, 7} if onset < fractions.Fraction(3, 2) else {0, 5, 9})
      self.assertTrue(generator.playback.compile_timeline(result)['events'])
  def test_harmony_carry_and_explicit_silence(self):
    content = xml_harmony(offset=1) + xml_harmony(kind='none', offset=2) + xml_harmony('F', offset=3) + tests.test_musicxml.xml_note()
    source = tests.test_musicxml.xml_score(content).replace('</measure></part>', '</measure><measure number="2">' + tests.test_musicxml.xml_note() + '</measure><measure number="3">' + xml_harmony(kind='none') + tests.test_musicxml.xml_note() + '</measure><measure number="4">' + tests.test_musicxml.xml_note() + '</measure></part>')
    result = tests.arrangement_fixtures.section_document(tests.test_musicxml.parse(source), 'acoustic', rhythms=tests.rhythm_fixtures.library())
    for measure in result['staves'][0]['measures']:
      self.assert_complete_measure(measure)
    sounding = [n for n in tests.test_musicxml.notes(result) if n['pitches']]
    for start, end in ((1, 2), (3, 4)):
      self.assertTrue(any(start <= generator.model.fraction(n['onset']) < end for n in sounding))
    for note in sounding:
      onset = generator.model.fraction(note['onset'])
      self.assertTrue(1 <= onset < 2 or 3 <= onset < 4)
      self.assertLessEqual(onset + generator.model.fraction(note['duration']), 2 if onset < 2 else 4)
    carried = [n for n in tests.test_musicxml.notes(result, 1) if n['pitches']]
    self.assertTrue(carried)
    self.assertTrue(all(generator.model.pitch_number(p) % 12 in {0, 5, 9} for n in carried for p in n['pitches']))
    for index in (2, 3):
      self.assertTrue(all(n['kind'] == 'rest' and not n['pitches'] for n in tests.test_musicxml.notes(result, index)))
    generator.model.validate_document(result)
  def test_accompaniment_handles_pickups_meters_and_coincident_markers(self):
    for duration, kind, attributes, implicit in ((1, 'quarter', '', True), (3, 'half', '<time><beats>3</beats><beat-type>4</beat-type></time>', False)):
      source = tests.test_musicxml.parse(tests.test_musicxml.xml_score(xml_harmony() + xml_harmony('F') + tests.test_musicxml.xml_note(duration=duration, kind=kind, extra='<dot/>' if duration == 3 else ''), attributes, implicit=implicit))
      measure = source['staves'][0]['measures'][0]
      plan = generator.accompaniment.plan_section([measure], tests.rhythm_fixtures.library(), electric=False)
      events = generator.accompaniment.realize(measure, plan.acoustic_upper[0], None)
      self.assertEqual(sum(generator.model.fraction(n['duration']) for n in events), duration)
      sounding = [n for n in events if n['kind'] == 'note']
      self.assertTrue(sounding)
      self.assertTrue(all(n['_harmony']['root']['letter'] == 'f' for n in sounding))
    source = tests.test_musicxml.parse(tests.test_musicxml.xml_score(xml_harmony() + tests.test_musicxml.xml_note(duration=1, kind='quarter') + '<attributes><time><beats>6</beats><beat-type>8</beat-type></time></attributes>' + tests.test_musicxml.xml_note(duration=3, kind='half', extra='<dot/>')))
    with self.assertRaisesRegex(ValueError, 'mid-measure meter'):
      tests.arrangement_fixtures.section_document(source, 'acoustic', rhythms=tests.rhythm_fixtures.library())
  def test_dense_melody_has_accompaniment_in_every_measure(self):
    bar = xml_harmony() + ''.join(tests.test_musicxml.xml_note(duration=1, kind='quarter') for _ in range(4))
    xml = tests.test_musicxml.xml_score(bar).replace('</part>', ''.join(f'<measure number="{i}">{bar}</measure>' for i in range(2, 13)) + '</part>')
    result = tests.arrangement_fixtures.section_document(tests.test_musicxml.parse(xml), 'acoustic', rhythms=tests.rhythm_fixtures.library())
    self.assertEqual(len(result['staves'][0]['measures']), 12)
    for measure in result['staves'][0]['measures']:
      self.assertTrue(any(n['pitches'] for v in measure['voices'] for n in v['notes']), measure['id'])
    generator.model.validate_document(result)
    timeline = generator.playback.compile_timeline(result)
    start = fractions.Fraction()
    for measure in result['staves'][0]['measures']:
      end = start + generator.model.fraction(measure['duration'])
      self.assertTrue(any(generator.model.fraction(e['onset']) < end and generator.model.fraction(e['onset']) + generator.model.fraction(e['duration']) > start for e in timeline['events']), measure['id'])
      start = end
  def test_prepared_document_and_parts_have_independent_ownership(self):
    source = tests.test_musicxml.parse(tests.test_musicxml.xml_score(xml_harmony() + tests.test_musicxml.xml_note()))
    source['staves'][0]['octave_shifts'] = [{'id': 'octave', 'start': generator.model.rational(0), 'end': generator.model.rational(4), 'octaves': -1}]
    original = copy.deepcopy(source)
    prepared = generator.parts.prepare_document(source)
    before = copy.deepcopy(prepared)
    plan = generator.arrange.generate.plan_arrangement(prepared, tests.rhythm_fixtures.library())
    plan_before = copy.deepcopy(plan)
    parts = [generator.arrange.generate._arrange_document(prepared, part, plan) for part in generator.parts.PART_TYPES]
    self.assertEqual(prepared, before)
    self.assertEqual(plan, plan_before)
    self.assertEqual(parts, [generator.arrange.generate.arrange_document(source, part, rhythms=tests.rhythm_fixtures.library()) for part in generator.parts.PART_TYPES])
    vocal = copy.deepcopy(parts[2])
    parts[1]['metadata']['title'] = 'changed'
    tests.test_musicxml.notes(parts[1])[0]['pitches'].clear()
    if parts[1]['staves'][0]['octave_shifts']:
      parts[1]['staves'][0]['octave_shifts'][0]['octaves'] = 9
    self.assertEqual(parts[2], vocal)
    self.assertEqual(prepared, before)
    self.assertEqual(source, original)
  def test_melody_written_pitch_rhythm_rests_and_cross_bar_ties_are_preserved(self):
    content = ''
    content += tests.test_musicxml.xml_note(duration=2, kind='quarter', alter=1, extra='<tie type="start"/>')
    content += tests.test_musicxml.xml_note(duration=6, kind='half', alter=1, extra='<dot/><tie type="stop"/><tie type="start"/>')
    root = xml.etree.ElementTree.fromstring(tests.test_musicxml.xml_score(content, divisions=2))
    second = tests.test_musicxml.xml_note(duration=4, kind='half', alter=1, extra='<tie type="stop"/>')
    second += '<note><rest/><duration>1</duration><type>eighth</type></note>'
    second += tests.test_musicxml.xml_note('D', duration=3, kind='quarter', alter=1, extra='<dot/>')
    root.find('part').append(xml.etree.ElementTree.fromstring('<measure number="2">' + second + '</measure>'))
    source = tests.test_musicxml.parse(xml.etree.ElementTree.tostring(root, encoding='unicode'))
    before = copy.deepcopy(source)
    expected = [n for m in source['staves'][0]['measures'] for n in m['voices'][0]['notes']]
    for builder in (tests.arrangement_fixtures.section_document, generator.arrange.generate.arrange_document):
      for part in ('electric', 'vocal'):
        with self.subTest(builder=builder.__module__, part=part):
          result = builder(source, part, rhythms=tests.rhythm_fixtures.library())
          generator.model.validate_document(result)
          measures = result['staves'][0]['measures']
          starts = [0] if builder is tests.arrangement_fixtures.section_document else [i for i, m in enumerate(measures) if m.get('section_label') == ('Guitar Melody' if part == 'electric' else 'Vocal Melody')]
          self.assertTrue(starts)
          for start in starts:
            notes = [n for m in measures[start:start + 2] for n in m['voices'][0]['notes']]
            self.assertEqual(len(notes), len(expected))
            for original, generated in zip(expected, notes):
              for key in ('onset', 'duration', 'kind', 'written'):
                self.assertEqual(generated[key], original[key])
              self.assertEqual(len(generated['pitches']), len(original['pitches']))
              for pitch, tone in zip(original['pitches'], generated['pitches']):
                octave = 12 if part == 'electric' else 0
                self.assertEqual(generator.model.pitch_number(tone), generator.model.pitch_number(pitch) + octave)
                self.assertEqual(tone['written'], pitch['written'])
                for key in ('tie_from', 'tie_to'):
                  self.assertEqual(key in tone, key in pitch)
            self.assertEqual(notes[0]['pitches'][0]['tie_to'], notes[1]['pitches'][0]['id'])
            self.assertEqual(notes[1]['pitches'][0]['tie_to'], notes[2]['pitches'][0]['id'])
          if builder is tests.arrangement_fixtures.section_document:
            events = generator.playback.compile_timeline(result)['events']
            self.assertEqual([(generator.model.fraction(e['onset']), generator.model.fraction(e['duration']), generator.model.fraction(e['pitch'])) for e in events], [(0, 6, 61 if part == 'electric' else 49), (fractions.Fraction(13, 2), fractions.Fraction(3, 2), 63 if part == 'electric' else 51)])
    self.assertEqual(source, before)
  def test_song_slug_collisions_fail_before_writes(self):
    for names in (('one-two', 'one_two'), ('Song', 'song'), ('!!!', '???')):
      with self.subTest(names=names), tempfile.TemporaryDirectory() as directory:
        root = pathlib.Path(directory)
        source, output = root / 'songs', root / 'output'
        source.mkdir()
        output.mkdir()
        (output / 'sentinel').write_text('preserved')
        for name in names:
          (source / (name + '.xml')).write_text(tests.test_musicxml.xml_score(tests.test_musicxml.xml_note()))
        with self.assertRaisesRegex(ValueError, 'conflicting output path') as raised:
          generator.arrange.generate.generate_performance_scores(source, output, rhythms=tests.rhythm_fixtures.library())
        for name in names:
          self.assertIn(str(source / (name + '.xml')), str(raised.exception))
        self.assertEqual([p.name for p in output.iterdir()], ['sentinel'])
  def test_pipeline_parses_and_normalizes_each_source_once(self):
    with tempfile.TemporaryDirectory() as directory:
      root = pathlib.Path(directory)
      for category, name in (('songs', 'song'), ('songs', 'other')):
        (root / 'data' / category).mkdir(parents=True, exist_ok=True)
        (root / 'data' / category / (name + '.xml')).write_text(tests.test_musicxml.xml_score(xml_harmony() + tests.test_musicxml.xml_note()))
      args = argparse.Namespace(data_dir=root / 'data', output_dir=root / 'output', static_assets_dir=pathlib.Path('package/assets'), check=False)
      with (
        unittest.mock.patch('generator.rhythm.load_library', return_value=tests.rhythm_fixtures.library()) as load_library,
        unittest.mock.patch('generator.musicxml.load_musicxml', wraps=generator.musicxml.load_musicxml) as load,
        unittest.mock.patch('generator.pitch.concert_document', wraps=generator.pitch.concert_document) as normalize,
        unittest.mock.patch('generator.engraving.pipeline.engrave_tree'),
        unittest.mock.patch('generator.pico_artwork.generate_titles'),
        unittest.mock.patch('generator.catalog.build_catalog'),
        unittest.mock.patch('generator.playback.compile_tree'),
      ):
        generator.__main__.generate_assets(args)
      load_library.assert_called_once_with(args.data_dir.resolve() / 'rhythms')
      self.assertEqual(load.call_args_list, [unittest.mock.call(args.data_dir.resolve() / 'songs/other.xml'), unittest.mock.call(args.data_dir.resolve() / 'songs/song.xml')])
      self.assertEqual(normalize.call_count, 2)
      original = {p.relative_to(args.output_dir): p.read_bytes() for p in args.output_dir.rglob('*') if p.is_file()}
      (root / 'data/songs/broken.xml').write_text('<broken')
      with unittest.mock.patch('generator.rhythm.load_library', return_value=tests.rhythm_fixtures.library()), self.assertRaisesRegex(ValueError, 'broken.xml'):
        generator.__main__.generate_assets(args)
      self.assertEqual(original, {p.relative_to(args.output_dir): p.read_bytes() for p in args.output_dir.rglob('*') if p.is_file()})
      self.assertEqual(len(list(args.output_dir.rglob('*.json'))), 6)
      self.assertEqual(json.loads((args.output_dir / 'songs/song-vocal.json').read_text())['metadata']['title'], 'Fixture')