"""Improvisation replaces source-allocated solos, including their pickups."""
import copy
import fractions
import pathlib
import unittest

import generator.arrange.assembly
import generator.arrange.form
import generator.arrange.generate
import generator.composition
import generator.composition_types
import generator.engraving.fonts
import generator.engraving.layout
import generator.engraving.types
import generator.model
import generator.musicxml
import generator.parts
import generator.performance
import generator.source
import generator.validation
import tests.production_fixtures
import tests.solo_fixtures
import tests.test_arrangement_pickup
import tests.test_musicxml

F = fractions.Fraction

def events(document):
  cursor = F()
  result = []
  for measure in document['staves'][0]['measures']:
    for voice in measure['voices']:
      for note in voice['notes']:
        start = cursor + generator.model.fraction(note['onset'])
        result.append((start, start + generator.model.fraction(note['duration']), note))
    cursor += generator.model.fraction(measure['duration'])
  return result

class ImprovisationTests(unittest.TestCase):
  def setUp(self):
    tests.production_fixtures.install(self)

  def check_variants(self, document):
    source = generator.source.SourceAnalysis.read(generator.parts.prepare_document(document))
    results = {r: generator.composition.generate_role(source, r, pathlib.Path('data')) for r in generator.composition_types.ROLES}
    before = copy.deepcopy(results)
    written = generator.arrange.assembly.assemble(source, results)
    improv = generator.arrange.assembly.assemble(source, results, solos_enabled=False)
    self.assertEqual(results, before)
    self.assertEqual(written['vocal'], improv['vocal'])
    for part in ('acoustic', 'electric'):
      windows = [(a, b) for a, b, role in generator.arrange.form.windows(source.document(), part) if role == part + '_solo']
      actual = events(improv[part])
      slashes = [(a, b) for a, b, n in actual if n.get('notehead') == 'slash']
      self.assertTrue(slashes)
      for start, end in windows:
        cursor = start
        for a, b in slashes:
          if start <= a < end:
            self.assertEqual(a, cursor)
            self.assertLessEqual(b, end)
            cursor = b
        self.assertEqual(cursor, end)
      self.assertFalse([n for a, b, n in actual if n['pitches'] and any(a < end and b > start for start, end in windows)])
      outside = lambda doc: [(a, b, n) for a, b, n in events(doc) if not any(a < end and b > start for start, end in windows)]
      self.assertEqual(outside(written[part]), outside(improv[part]))
      for a, b in zip(written[part]['staves'][0]['measures'], improv[part]['staves'][0]['measures']):
        for key in ('harmonies', 'section_label', 'number', 'duration', 'attributes'):
          self.assertEqual(a.get(key), b.get(key), key)
      self.assertEqual(generator.performance.timing(written[part]), generator.performance.timing(improv[part]))
      generator.validation.validate_document(improv[part])
    return improv

  def test_no_pickup_and_simple_meters(self):
    for meter in (3, 4, 5):
      with self.subTest(meter=meter):
        self.check_variants(tests.solo_fixtures.song(count=2, meter=meter, chords=('CM7',)))

  def test_partial_pickup_includes_allocated_rests_and_preserves_earlier_music(self):
    pickup = '<note><rest/><duration>1</duration><type>eighth</type></note>' + tests.test_musicxml.xml_note('D', 1, 'eighth')
    self.check_variants(tests.test_arrangement_pickup.source(pickup, divisions=2))

  def test_repertoire_pickup_allocations(self):
    for song, first in (('all-of-me', (288, 416)), ('autumn_leaves', (284, 412))):
      with self.subTest(song=song):
        variants = self.check_variants(generator.musicxml.load_musicxml(pathlib.Path('data/songs') / (song + '.xml')))
        for part, expected in zip(('acoustic', 'electric'), first):
          slashes = [(a, b) for a, b, n in events(variants[part]) if n.get('notehead') == 'slash']
          self.assertEqual(slashes[0][0], expected)
          self.assertTrue(all(b - a == 1 for a, b in slashes))

  def test_meter_changes_and_compound_beats_render_in_staff_and_tab(self):
    change = '<attributes><time><beats>6</beats><beat-type>8</beat-type></time></attributes>'
    document = tests.test_musicxml.parse(tests.test_musicxml.xml_score(tests.test_musicxml.xml_note(duration=1, kind='quarter') + change + tests.test_musicxml.xml_note(duration=3, kind='half', extra='<dot/>')))
    source = generator.source.SourceAnalysis.read(generator.parts.prepare_document(document))
    template = generator.arrange.generate._arrange_document(source.document(), 'acoustic', owned={'slash'})
    template['staves'][0]['tablature'] = True
    result = generator.arrange.assembly.improvisation(source, template, 'acoustic')
    result['metadata']['solos_enabled'] = False
    solo = next(m for m in result['staves'][0]['measures'] if m.get('section_label') == 'Acoustic Solo')
    self.assertEqual([generator.model.fraction(n['duration']) for n in solo['voices'][0]['notes']], [1, F(3, 2), F(3, 2)])
    generator.validation.validate_document(result)
    fonts = generator.engraving.fonts.Fonts(pathlib.Path('package/assets/fonts/petaluma'))
    layout = generator.engraving.layout.layout_score(result, generator.engraving.types.Style(), fonts)
    primitives = [p for page in layout.pages for p in page.primitives]
    self.assertEqual(sum(p.role == 'slash' for p in primitives), 3)
    self.assertEqual(sum(p.role == 'tab_slash' for p in primitives), 3)
    self.assertFalse(any(p.role in ('tab', 'head') for p in primitives))
    self.assertTrue(any(p.role == 'subtitle' and 'Improvisation' in p.value for p in primitives))