import copy
import fractions
import itertools
import pathlib
import unittest

import generator.accompaniment
import generator.acoustic
import generator.acoustic_voicing
import generator.arrange.generate
import generator.fretboard
import generator.harmony
import generator.model
import generator.musicxml
import tests.test_acoustic
import tests.test_accompaniment

def planned_document(harmonies, attacks=None):
  if attacks is None:
    attacks = [([(0, 1), (2, 1)], [(1, 1), (3, 1)])] * len(harmonies)
  document = tests.test_acoustic.acoustic_document(attacks)
  previous = None
  for measure, changes in zip(document['staves'][0]['measures'], harmonies):
    measure['harmonies'] = []
    for onset, symbol in changes:
      harmony = {'kind': 'none', 'degrees': []} if symbol == 'N.C.' else generator.harmony.parse_symbol(symbol) if isinstance(symbol, str) else copy.deepcopy(symbol)
      measure['harmonies'].append(harmony | {'id': f'{measure["id"]}_h{len(measure["harmonies"])}', 'onset': generator.model.rational(onset)})
    for voice in measure['voices']:
      for note in voice['notes']:
        if not note['pitches']:
          continue
        note['pitches'] = []
        harmony = generator.accompaniment.active_chord(measure, generator.model.fraction(note['onset']), previous)
        if harmony:
          note.update(_harmony=harmony, _guitar_profile='acoustic', _guitar_lane=voice['id'])
        else:
          note['kind'] = 'rest'
    if measure['harmonies']:
      previous = measure['harmonies'][-1]
  generator.acoustic.prepare(document)
  return document

def shape(note):
  tuning = dict(generator.fretboard.STANDARD_TUNING)
  return tuple(sorted(generator.fretboard.TabPosition(p['tablature']['string'], p['tablature']['fret'], tuning[p['tablature']['string']] + p['tablature']['fret']) for p in note['pitches']))

def voicing(frets, bass_string=5):
  tuning = dict(generator.fretboard.STANDARD_TUNING)
  return generator.acoustic_voicing.Voicing(*(generator.fretboard.TabPosition(s, f, tuning[s] + f) for s, f in zip((bass_string, 4, 3, 2), frets)))

def candidates(chord):
  definition = generator.harmony.parse_symbol(chord) if isinstance(chord, str) else chord
  return generator.acoustic_voicing.candidates(generator.fretboard.chord_requirements(definition), generator.harmony.is_triad(definition))

class AcousticVoicingTests(unittest.TestCase):
  def place(self, document):
    spans = generator.acoustic.harmony_spans(document)
    lanes = {n['id']: n['_guitar_lane'] for span in spans for n in span.notes}
    before = [tests.test_accompaniment.timing(v['notes']) for m in document['staves'][0]['measures'] for v in m['voices']]
    generator.acoustic_voicing.place(document)
    self.assertEqual(before, [tests.test_accompaniment.timing(v['notes']) for m in document['staves'][0]['measures'] for v in m['voices']])
    for span in spans:
      basses = {shape(n) for n in span.notes if lanes[n['id']] == 'bass'}
      uppers = {shape(n) for n in span.notes if lanes[n['id']] == 'upper'}
      self.assertLessEqual(len(basses), 1)
      self.assertLessEqual(len(uppers), 1)
      grip = tuple(sorted({p for n in span.notes for p in shape(n)}, reverse=True))
      requirements = generator.fretboard.chord_requirements(span.harmony)
      if basses and uppers:
        self.assertTrue(generator.acoustic_voicing.playable(generator.acoustic_voicing.Voicing(*grip)))
        self.assertLessEqual(max(p.fret for p in grip) - min(p.fret for p in grip) + 1, 4)
        self.assertTrue(requirements.covers({p.midi % 12 for p in grip}))
        self.assertTrue({p.midi % 12 for p in grip} <= set(requirements.allowed))
      for note in span.notes:
        strings = sorted(p['tablature']['string'] for p in note['pitches'])
        if lanes[note['id']] == 'upper':
          self.assertEqual(strings, [2, 3, 4])
        else:
          self.assertIn(strings, ([5], [6]))
        self.assertFalse(any(k.startswith('_') for k in note))
        self.assertTrue(all('tie_from' not in p and 'tie_to' not in p for p in note['pitches']))
    generator.model.validate_document(document)
    return spans
  def test_independent_attacks_keep_one_complete_voicing(self):
    document = planned_document([[(0, 'FM7')]])
    spans = self.place(document)
    self.assertEqual(len(spans), 1)
    self.assertEqual(len(spans[0].notes), 4)
    before = tests.test_acoustic.sounding(document)
    generator.acoustic.combine(document)
    self.assertEqual(tests.test_acoustic.sounding(document), before)
    generator.model.validate_document(document)
  def test_four_strings_allow_all_defined_bass_inversions(self):
    for symbol in ('CM', 'Dm', 'Caug', 'Cdim', 'C7', 'Cdim7', 'C7no5', 'C7b9', 'C7no1', 'C6', 'C13', 'Csus', 'Csus2'):
      with self.subTest(chord=symbol):
        chord = generator.harmony.parse_symbol(symbol)
        requirements = generator.fretboard.chord_requirements(chord)
        options = candidates(chord)
        self.assertTrue(options)
        self.assertFalse(requirements.basses)
        self.assertTrue(all(v.bass.midi % 12 in requirements.allowed for v in options))
        self.assertTrue(all(requirements.covers({p.midi % 12 for p in v}) for v in options))
        self.assertTrue(all(generator.acoustic_voicing.playable(v) for v in options))
    self.assertEqual({v.bass.midi % 12 for v in candidates('C7')}, {0, 4, 7, 10})
    self.assertIn(1, {v.bass.midi % 12 for v in candidates('C7b9')})
    self.assertIn(5, {v.bass.midi % 12 for v in candidates('Dm')})
  def test_same_voicing_survives_repeated_symbols_and_rest_bars(self):
    attacks = [([(0, 1), (2, 1)], [(0, 1), (3, 1)]), ([], []), ([(2, 1)], [(1, 1), (3, 1)])]
    document = planned_document([[(0, 'C7')], [], [(0, 'C7'), (2, 'C7')]], attacks)
    self.assertEqual(len(self.place(document)), 1)
  def test_iterations_are_independent_even_with_identical_boundary_chords(self):
    first, second = [[(0, s)] for s in ('C7', 'G7')], [[(0, s)] for s in ('G7', 'FM7')]
    document = planned_document(first + second)
    document['staves'][0]['measures'][2]['section_label'] = 'Next iteration'
    groups = [list(group) for _, group in itertools.groupby(generator.acoustic.harmony_spans(document), key=lambda s: s.iteration)]
    self.assertEqual([len(g) for g in groups], [2, 2])
    expected = []
    for harmonies in (first, second):
      separate = planned_document(harmonies)
      self.place(separate)
      expected.extend(separate['staves'][0]['measures'])
    self.place(document)
    def positions(measures):
      return [[shape(n) for v in m['voices'] for n in v['notes']] for m in measures]
    self.assertEqual(positions(document['staves'][0]['measures']), positions(expected))
  def test_changes_during_silence_end_the_harmony_span_not_the_journey(self):
    for middle in ('G7', 'N.C.'):
      with self.subTest(middle=middle):
        document = planned_document([[(0, 'C7'), (1, middle), (2, 'C7')]], [([(0, 1), (2, 1)], [(0, 1), (3, 1)])])
        spans = self.place(document)
        self.assertEqual(len(spans), 2)
        self.assertEqual([len(s.notes) for s in spans], [2, 2])
        self.assertEqual(len({s.iteration for s in spans}), 1)
  def test_returning_to_a_chord_is_a_new_span(self):
    document = planned_document([[(0, 'C7')], [(0, 'G7')], [(0, 'C7')]])
    self.assertEqual(len(self.place(document)), 3)
  def test_slash_bass_inversion_and_alteration_changes_are_boundaries(self):
    inversion = generator.harmony.parse_symbol('C7') | {'inversion': 2}
    changes = [(0, 'C7'), (1, 'C7/E'), (2, inversion), (3, 'C7b5')]
    document = planned_document([changes], [([(i, 1) for i in range(4)], [(i, 1) for i in range(4)])])
    spans = self.place(document)
    self.assertEqual(len(spans), 4)
    basses = [min(p.midi for n in span.notes for p in shape(n)) % 12 for span in spans]
    self.assertEqual(basses[1:3], [4, 7])
  def test_non_accompaniment_passage_ends_the_span(self):
    document = planned_document([[(0, 'C7')], [], []])
    measure = document['staves'][0]['measures'][1]
    slash = generator.model.make_event('slash', fractions.Fraction(), generator.model.fraction(measure['duration']), 'note') | {'voice': 'voice_arranged', 'notehead': 'slash'}
    measure['voices'] = [{'id': 'voice_arranged', 'notes': [slash]}]
    self.assertEqual(len(self.place(document)), 2)
  def test_carried_harmony_can_begin_an_arrangement_excerpt(self):
    document = planned_document([[(0, 'C7')], []])
    document['staves'][0]['measures'][0]['harmonies'] = []
    self.assertEqual(len(self.place(document)), 1)
  def test_no_chord_and_empty_iterations_do_not_generate_notes(self):
    document = planned_document([[(0, 'N.C.')], [(0, 'C7')], [(0, 'N.C.')]])
    for index, measure in enumerate(document['staves'][0]['measures']):
      measure['section_label'] = f'Iteration {index}'
    self.assertEqual(len(self.place(document)), 1)
    for measure in (document['staves'][0]['measures'][0], document['staves'][0]['measures'][2]):
      self.assertTrue(all(n['kind'] == 'rest' for v in measure['voices'] for n in v['notes']))
    self.assertEqual(generator.acoustic_voicing.select([]), [])
  def test_unsounded_voices_still_constrain_the_full_voicing(self):
    for lane in ('bass', 'upper'):
      document = planned_document([[(0, 'C7')]], [([(0, 1)] if lane == 'bass' else [], [(1, 1)] if lane == 'upper' else [])])
      span = generator.acoustic.harmony_spans(document)[0]
      selected = generator.acoustic_voicing.select([span])[0]
      self.assertTrue(generator.acoustic_voicing.playable(selected))
      self.place(document)
      self.assertEqual(shape(span.notes[0]), tuple(sorted((selected.bass,) if lane == 'bass' else selected[1:])))
  def test_rhythm_density_does_not_multiply_movement_cost(self):
    harmonies = [[(0, symbol)] for symbol in ('FM7', 'Gm7', 'C7')]
    sparse = planned_document(harmonies, [([(0, 1)], [(2, 1)])] * 3)
    dense = planned_document(harmonies, [([(i, 1) for i in range(4)], [(i, 1) for i in range(4)])] * 3)
    selected = [tuple(tuple(sorted({p for n in span.notes for p in shape(n)})) for span in self.place(document)) for document in (sparse, dense)]
    self.assertEqual(*selected)
  def test_bass_switch_counts_both_string_and_fret_travel(self):
    previous = voicing((7, 8, 7, 8))
    for current, expected in ((voicing((7, 8, 7, 8), 6), 1), (voicing((8, 8, 7, 8), 6), 2), (voicing((9, 8, 7, 8), 6), 3), (voicing((6, 8, 7, 8)), 1)):
      self.assertEqual(generator.acoustic_voicing.movement(previous, current), expected)
      self.assertEqual(generator.acoustic_voicing.movement(current, previous), expected)
    self.assertEqual(generator.acoustic_voicing.movement(previous, previous), 0)
  def test_each_upper_string_is_tracked_not_just_the_grip_center(self):
    before, after = voicing((5, 7, 5, 7)), voicing((7, 5, 7, 5))
    self.assertEqual(generator.acoustic_voicing.movement(before, after), 8)
    for index in (1, 2, 3):
      frets = [5, 7, 5, 7]
      frets[index] += 1
      self.assertEqual(generator.acoustic_voicing.movement(before, voicing(frets)), 1)
  def test_fret_and_register_limits_are_hard_constraints(self):
    for frets, valid in (((1, 1, 1, 1), True), ((1, 3, 1, 3), True), ((1, 4, 1, 4), True), ((1, 5, 1, 4), False), ((0, 1, 1, 1), False), ((12, 14, 14, 12), True), ((12, 15, 14, 12), True), ((13, 13, 13, 13), False), ((12, 16, 14, 12), False), ((12, 21, 12, 12), False)):
      with self.subTest(frets=frets):
        self.assertEqual(generator.acoustic_voicing.playable(voicing(frets)), valid)
    for symbol in ('CM', 'C7b9', 'C13', 'Cdim7'):
      self.assertTrue(all(generator.acoustic_voicing.playable(v) for v in candidates(symbol)))
  def test_four_fret_cm7_and_f7_are_fallbacks_not_preferred_shapes(self):
    for symbol, frets in (('Cm7', (3, 1, 3, 4)), ('F7', (3, 1, 2, 4))):
      with self.subTest(symbol=symbol):
        self.assertTrue(generator.acoustic_voicing.playable(voicing(frets)))
        self.assertIn(voicing(frets), candidates(symbol))
        document = planned_document([[(0, symbol)]])
        selected = generator.acoustic_voicing.select(generator.acoustic.harmony_spans(document))
        self.assertTrue(all(generator.fretboard.span_penalty(v) == 0 for v in selected))
        self.place(document)
  def test_acoustic_cost_has_no_register_span_or_shape_penalties(self):
    options = (voicing((1, 1, 1, 1)), voicing((1, 3, 2, 3)), voicing((12, 12, 12, 12)), voicing((12, 14, 13, 14)))
    states = [generator.acoustic_voicing.Placement(generator.acoustic_voicing.Cost())]
    result = generator.acoustic_voicing.advance(options, generator.fretboard.ChordRequirements(), states)
    self.assertEqual([p.cost for p in result], [generator.acoustic_voicing.Cost()] * len(options))
  def test_optional_tones_break_equal_movement_costs_only(self):
    requirements = generator.fretboard.chord_requirements(generator.harmony.parse_symbol('C7'))
    states = [generator.acoustic_voicing.Placement(generator.acoustic_voicing.Cost())]
    result = generator.acoustic_voicing.advance(candidates('C7'), requirements, states)
    self.assertTrue(all(p.cost.travel == 0 for p in result))
    self.assertEqual({p.cost.missing_tones for p in result}, {0, 1})
    best = min(result, key=lambda p: (p.cost, p.voicing))
    self.assertEqual({p.midi % 12 for p in best.voicing}, {0, 4, 7, 10})
  def test_dynamic_program_matches_exhaustive_paths_and_is_deterministic(self):
    harmonies = ['FM7', 'Gm7', 'C7']
    spans = generator.acoustic.harmony_spans(planned_document([[(0, s)] for s in harmonies]))
    options = [candidates(s) for s in harmonies]
    requirements = [generator.fretboard.chord_requirements(span.harmony) for span in spans]
    def cost(path):
      return generator.acoustic_voicing.Cost(sum(generator.fretboard.span_penalty(v) for v in path), sum(generator.acoustic_voicing.movement(a, b) for a, b in zip(path, path[1:])), sum(len(set(r.allowed) - {p.midi % 12 for p in v}) for r, v in zip(requirements, path)))
    selected = generator.acoustic_voicing.select(spans)
    self.assertEqual(cost(selected), min(cost(path) for path in itertools.product(*options)))
    self.assertEqual(selected, generator.acoustic_voicing.select(spans))
  def test_complete_triads_and_doubling_use_only_authored_chord_tones(self):
    for symbol in ('CM', 'Dm', 'Cdim', 'Caug', 'Dm/F', 'Dm7no7'):
      with self.subTest(chord=symbol):
        classes = set(generator.harmony.pitch_classes(generator.harmony.parse_symbol(symbol)))
        options = candidates(symbol)
        self.assertTrue(options)
        self.assertTrue(all(len(v) == 4 and {p.midi % 12 for p in v} == classes for v in options))
    for symbol, omitted in (('C7no1', 0), ('C7no5', 7), ('Dmno1', 2), ('CMno5', 7)):
      self.assertTrue(candidates(symbol))
      self.assertTrue(all(omitted not in {p.midi % 12 for p in v} for v in candidates(symbol)))
  def test_unplayable_omission_is_not_restored_to_fill_upper_strings(self):
    for symbol in ['Dmno5']:
      with self.subTest(symbol=symbol):
        self.assertFalse(candidates(symbol))
        document = planned_document([[(0, symbol)]])
        before = copy.deepcopy(document)
        with self.assertRaisesRegex(ValueError, f'{symbol}: no acoustic grip.*4 consecutive fret positions'):
          generator.acoustic_voicing.place(document)
        self.assertEqual(document, before)
  def test_explicit_inversions_and_nonchord_slash_basses_are_preserved(self):
    for chord, bass in ((generator.harmony.parse_symbol('Dm') | {'inversion': 1}, 5), (generator.harmony.parse_symbol('C7') | {'inversion': 3}, 10), ('C7/E', 4), ('C7b9/Db', 1), ('Dm/B', 11)):
      with self.subTest(chord=chord):
        span = self.place(planned_document([[(0, chord)]]))[0]
        self.assertEqual({p.midi % 12 for n in span.notes for p in shape(n) if p.string_number >= 5}, {bass})
  def test_impossible_later_iteration_does_not_partially_materialize(self):
    document = planned_document([[(0, 'C7')], [(0, 'C7b9#11')]])
    document['staves'][0]['measures'][1]['section_label'] = 'Later iteration'
    before = copy.deepcopy(document)
    with self.assertRaisesRegex(ValueError, 'measure 2, event .*no acoustic grip.*1-20'):
      generator.acoustic_voicing.place(document)
    self.assertEqual(document, before)
  def test_invalid_inversion_and_tuning_fail_with_context(self):
    invalid = generator.harmony.parse_symbol('C7') | {'inversion': 4}
    document = planned_document([[(0, invalid)]])
    with self.assertRaisesRegex(ValueError, 'measure 1.*inversion'):
      generator.acoustic_voicing.place(document)
    document = planned_document([[(0, 'C7')]])
    document['staves'][0]['measures'][0]['staff_details']['capo'] = 1
    with self.assertRaisesRegex(ValueError, 'capo'):
      generator.acoustic_voicing.place(document)