import contextlib
import copy
import dataclasses
import fractions
import io
import pathlib
import shutil
import tempfile
import unittest

import generator.arrange.form
import generator.arrange.generate
import generator.engraving.cache
import generator.engraving.fonts
import generator.engraving.layout
import generator.engraving.types
import generator.harmony
import generator.model
import generator.musicxml
import generator.parts
import generator.playback
import generator.scales.mapping
import generator.scales.library
import generator.scales.selection
import generator.scales.timeline
import tests.rhythm_fixtures
import tests.scale_fixtures
import tests.test_accompaniment
import tests.test_musicxml

def timeline(symbols, key=None, duration=4):
  key = key or generator.model.key_from_fifths(-1, 'minor')
  result = []
  for index, symbol in enumerate(symbols):
    result.append(generator.scales.timeline.ChordSpan(fractions.Fraction(index * duration), fractions.Fraction((index + 1) * duration), index + 1, str(index + 1), fractions.Fraction(1), key, generator.harmony.parse_symbol(symbol) if symbol else None))
  return tuple(result)

def choose(symbols, key=None):
  spans = timeline(symbols, key)
  return generator.scales.selection.select(spans, generator.scales.timeline.Solo(spans[0].start, spans[-1].end, 'Solo'), tests.scale_fixtures.library())

class ScaleLibraryTests(unittest.TestCase):
  def test_all_patterns_and_transpositions_preserve_spelling_and_rotations(self):
    library = tests.scale_fixtures.library()
    self.assertEqual(len(library.patterns), 21)
    self.assertEqual(len({scale.pitches for scale in library.collections()}), 36)
    for pattern in library.patterns:
      for root in generator.scales.library.ROOTS:
        scale = library.scale(pattern.family, pattern.mode, root)
        self.assertEqual(len(scale.pitches), 7)
        self.assertEqual(scale.notes[0], root)
        self.assertEqual(library.chord_mode(scale, root), (1, scale.label()))
    self.assertEqual(library.scale('harmonic_minor', 3, 'C').notes, ('C', 'D', 'E', 'F', 'G#', 'A', 'B'))
    self.assertEqual(library.scale('harmonic_minor', 7, 'C').notes, ('C', 'Db', 'Eb', 'Fb', 'Gb', 'Ab', 'Bbb'))
    self.assertEqual(library.scale('major', 1, 'C#').pitches, library.scale('major', 1, 'Db').pitches)
    self.assertNotEqual(library.scale('major', 1, 'C#').notes, library.scale('major', 1, 'Db').notes)
  def test_incomplete_unexpected_and_inconsistent_xml_patterns_are_rejected(self):
    with tempfile.TemporaryDirectory() as directory:
      root = pathlib.Path(directory) / 'scales'
      shutil.copytree('data/scales', root)
      path = root / 'harmonic_minor_scale_modes/0003.xml'
      original = path.read_text()
      path.unlink()
      with self.assertRaisesRegex(ValueError, '21 scale patterns'):
        generator.scales.library.load_library(root)
      path.write_text(original.replace('<step>A</step>', '<step>G</step>'))
      with self.assertRaisesRegex(ValueError, '0003.xml.*C through B'):
        generator.scales.library.load_library(root)
      path.write_text(original.replace('<alter>1</alter>', '<alter>0</alter>').replace('<accidental>sharp</accidental>', '<accidental>natural</accidental>'))
      with self.assertRaisesRegex(ValueError, '0003.xml.*not mode 3'):
        generator.scales.library.load_library(root)
      path.write_text(original)
      extra = root / 'unknown.xml'
      extra.write_text(original)
      with self.assertRaisesRegex(ValueError, 'unexpected'):
        generator.scales.library.load_library(root)

class ScaleSelectionTests(unittest.TestCase):
  def test_autumn_leaves_guide_collections_and_modes(self):
    first = ['Gm7', 'C7', 'FM7', 'BbM7', 'Em7b5', 'A7', 'Dm7', 'Dm7']
    bridge = ['Em7b5', 'A7', 'Dm7', 'Dm7', 'Gm7', 'C7', 'FM7', 'FM7']
    labels = [c.scale.label() for c in choose(first + first + bridge)]
    phrase = ['D Aeolian'] * 4 + ['D Harmonic Minor'] * 2 + ['D Aeolian'] * 2
    self.assertEqual(labels, phrase * 2 + ['D Harmonic Minor'] * 2 + ['D Aeolian'] * 6)
    ending = choose(['Em7b5', 'A7', 'Dm7', 'G7', 'Cm7', 'F7', 'Em7b5', 'A7', 'Dm7'])
    self.assertEqual([c.scale.label() for c in ending], ['D Harmonic Minor'] * 2 + ['D Aeolian', 'C Harmonic Minor', 'Bb Ionian', 'Bb Ionian'] + ['D Harmonic Minor'] * 2 + ['D Aeolian'])
    library = tests.scale_fixtures.library()
    self.assertEqual(library.chord_mode(ending[0].scale, 'E'), (2, 'E Locrian natural 6'))
    self.assertEqual(library.chord_mode(ending[1].scale, 'A'), (5, 'A Phrygian Dominant'))
    self.assertEqual(library.chord_mode(ending[4].scale, 'C'), (2, 'C Dorian'))
  def test_actual_song_uses_minor_tonic_triads_without_forcing_sevenths(self):
    source = generator.musicxml.load_musicxml(pathlib.Path('data/songs/autumn_leaves.xml'))
    spans = generator.scales.timeline.read(source['staves'][0])
    choices = generator.scales.selection.select(spans, generator.scales.timeline.Solo(4, spans[-1].end, 'Solo'), tests.scale_fixtures.library())
    self.assertTrue(all(c.scale for c in choices))
    for choice in choices:
      chord = choice.span.chord
      if chord['symbol'] == 'Dm':
        self.assertNotIn(7, generator.harmony.degrees(chord))
        self.assertEqual(choice.scale.label(), 'D Aeolian')
      self.assertLessEqual(set(generator.harmony.pitch_classes(chord)), choice.scale.pitches)
    half_bar = [(c.span.source_measure, c.span.beat, c.scale.label()) for c in choices if c.span.source_measure in ('28', '29')]
    self.assertEqual(half_bar, [('28', 1, 'D Aeolian'), ('28', 3, 'C Harmonic Minor'), ('29', 1, 'Bb Ionian'), ('29', 3, 'Bb Ionian')])
  def test_extensions_slash_basses_omissions_and_diminished_chords(self):
    for symbol in ('Em9b5', 'DmM7', 'Dm6', 'FM7#5', 'C#dim7', 'Dm7/F', 'C7no5', 'C7b5b9'):
      with self.subTest(symbol=symbol):
        choice = choose([symbol])[0]
        self.assertIsNotNone(choice.scale)
        self.assertLessEqual(set(generator.harmony.pitch_classes(choice.span.chord)), choice.scale.pitches)
    choices = choose(['C#dim7', 'Dm'])
    self.assertEqual([c.scale.label() for c in choices], ['D Harmonic Minor', 'D Aeolian'])
    self.assertIsNone(choose(['A13b9'])[0].scale)
    with self.assertRaises(ValueError):
      generator.harmony.parse_symbol('A7alt')
  def test_shared_collection_wins_over_independent_compatible_colors(self):
    choices = choose(['E7', 'F#7'], generator.model.key_from_fifths(0))
    self.assertEqual([choice.scale.label() for choice in choices], ['B Melodic Minor'] * 2)
    self.assertEqual(choices, choose(['E7', 'F#7'], generator.model.key_from_fifths(0)))
  def test_carried_harmony_and_fractional_key_boundaries_are_exact(self):
    document = tests.test_musicxml.parse(tests.test_accompaniment.xml_bars([tests.test_musicxml.xml_note()] * 3))
    measures = document['staves'][0]['measures']
    measures[0]['harmonies'] = [{'id': 'h1', 'onset': generator.model.rational(0), **generator.harmony.parse_symbol('CM7')}]
    measures[1]['attributes'] = [{'onset': generator.model.rational(fractions.Fraction(3, 2)), 'key': generator.model.key_from_fifths(1)}]
    measures[2]['harmonies'] = [{'id': 'h2', 'onset': generator.model.rational(0), 'kind': 'none', 'degrees': []}]
    spans = generator.scales.timeline.read(document['staves'][0])
    self.assertEqual([(s.start, s.end) for s in spans], [(0, 4), (4, fractions.Fraction(11, 2)), (fractions.Fraction(11, 2), 8), (8, 12)])
    self.assertEqual(spans[1].chord, spans[0].chord)
    self.assertEqual(spans[2].key['tonic']['letter'], 'g')
    self.assertIsNone(spans[3].chord)
  def test_microtonal_harmonies_are_explicitly_unmatched(self):
    spans = timeline(['A7', 'Dm'])
    for span in spans:
      span.chord['root']['alter'] = generator.model.rational(fractions.Fraction(1, 2))
    result = generator.scales.selection.select(spans, generator.scales.timeline.Solo(0, 8, 'Solo'), tests.scale_fixtures.library())
    self.assertTrue(all(c.scale is None and 'No exact match' in c.reason for c in result))
  def test_local_key_and_no_chord_break_continuity_and_missing_tonic_fails(self):
    spans = list(timeline(['Dm', None, 'Cm']))
    spans[-1] = dataclasses.replace(spans[-1], key=generator.model.key_from_fifths(-3, 'minor'))
    solo = generator.scales.timeline.Solo(0, 12, 'Solo')
    result = generator.scales.selection.select(tuple(spans), solo, tests.scale_fixtures.library())
    self.assertEqual([c.scale.label() if c.scale else None for c in result], ['D Aeolian', None, 'C Aeolian'])
    spans[0] = dataclasses.replace(spans[0], key=generator.model.key_from_fifths(0, 'none'))
    with self.assertRaisesRegex(ValueError, 'measure 1.*declared tonic'):
      generator.scales.selection.select(tuple(spans), solo, tests.scale_fixtures.library())

class ScaleOutputTests(unittest.TestCase):
  def test_annotations_reports_and_warnings_preserve_music_and_playback(self):
    source = tests.test_musicxml.parse(tests.test_accompaniment.xml_bars([tests.test_musicxml.xml_note()] * 4))
    staff = source['staves'][0]
    staff.update(name='Acoustic', arrangement_part_type='acoustic')
    for measure, symbol in zip(staff['measures'], ('CM7', 'CM7', 'A13b9', 'CM7')):
      measure['harmonies'] = [{'id': measure['id'] + '-h1', 'onset': generator.model.rational(0), **generator.harmony.parse_symbol(symbol)}]
    before = copy.deepcopy(source)
    warnings = io.StringIO()
    with contextlib.redirect_stdout(warnings):
      report = generator.scales.mapping.annotate(source, [generator.scales.timeline.Solo(0, 12, 'Acoustic Solo')], tests.scale_fixtures.library(), 'fixture.xml')
    self.assertTrue(report.startswith('Acoustic: Acoustic Solo\n'))
    self.assertIn('fixture.xml, Acoustic, source measure 3', warnings.getvalue())
    self.assertIn('No exact match within the permitted catalog', report)
    self.assertIn('Parent degree: 1; chord-root mode: C Ionian', report)
    spans = staff.pop('scale_spans')
    self.assertEqual(source, before)
    self.assertEqual([(generator.model.fraction(s['start']), generator.model.fraction(s['end']), s['label']) for s in spans], [(0, 8, 'C Ionian'), (8, 12, 'No exact scale match')])
    staff['scale_spans'] = spans
    generator.model.validate_document(source)
    self.assertNotEqual(generator.engraving.cache.score_fingerprint(source, 'renderer'), generator.engraving.cache.score_fingerprint(before, 'renderer'))
    self.assertEqual(generator.playback.compile_timeline(source), generator.playback.compile_timeline(before))
  def test_layout_keeps_exact_offsets_continuations_and_uniform_size(self):
    source = tests.test_musicxml.parse(tests.test_accompaniment.xml_bars([tests.test_musicxml.xml_note()] * 40))
    staff = source['staves'][0]
    staff['measures'][0]['section_label'] = 'Electric Solo'
    staff['scale_spans'] = [{'id': f'scale-{i}', 'start': generator.model.rational(a), 'end': generator.model.rational(b), 'label': label, 'matched': True} for i, (a, b, label) in enumerate(((0, fractions.Fraction(3, 2), 'C Ionian'), (fractions.Fraction(3, 2), 60, 'D Harmonic Minor')), 1)]
    fonts = generator.engraving.fonts.Fonts(pathlib.Path('package/assets/fonts/petaluma'))
    layout = generator.engraving.layout.layout_score(source, generator.engraving.types.Style(), fonts)
    systems = [s for p in layout.pages for s in p.systems]
    self.assertEqual(len({s.scale for s in systems}), 1)
    self.assertTrue(all(len(p.systems) == 2 for p in layout.pages[:-1]))
    self.assertGreater(len(systems), 2)
    self.assertEqual([p.value for s in systems for p in s.primitives if p.role == 'section'], ['Electric Solo'])
    for system in systems:
      expected = [s['label'] for s in staff['scale_spans'] if generator.model.fraction(s['start']) < system.end and generator.model.fraction(s['end']) > system.start]
      self.assertEqual([p.value for p in system.primitives if p.role == 'scale'], expected)
    self.assertIn(fractions.Fraction(3, 2), generator.engraving.notation.ScoreIndex(source).column_times)
  def test_only_soloing_part_is_annotated_without_revoicing(self):
    content = '<harmony><root><root-step>C</root-step></root><kind>major</kind></harmony>' + tests.test_musicxml.xml_note()
    source = tests.test_musicxml.parse(tests.test_musicxml.xml_score(content))
    roles = {'acoustic': ('slash', 'accompaniment'), 'electric': ('accompaniment', 'slash'), 'vocal': ('rest', 'rest')}
    for part in ('acoustic', 'electric', 'vocal'):
      before = generator.arrange.generate.arrange_document(source, part, rhythms=tests.rhythm_fixtures.library())
      result = generator.arrange.generate.arrange_document(source, part, rhythms=tests.rhythm_fixtures.library(), scales=tests.scale_fixtures.library())
      sections = generator.arrange.form.full_arrangement_sections(source['staves'][0]['measures'], part)
      self.assertEqual([(s.label, s.role) for s in sections[3:5]], list(zip(('Acoustic Solo', 'Electric Solo'), roles[part])))
      spans = result['staves'][0].pop('scale_spans', [])
      self.assertEqual(result, before)
      self.assertEqual(len(spans), int(part != 'vocal'))
      if spans:
        self.assertEqual(generator.model.fraction(spans[0]['start']), 12 if part == 'acoustic' else 16)
        self.assertEqual(generator.model.fraction(spans[0]['end']), 16 if part == 'acoustic' else 20)