import copy
import pathlib
import unittest

import generator.arrange.generate
import generator.guitar
import generator.harmony
import generator.model
import generator.musicxml
import generator.rhythm
import tests.rhythm_fixtures
import tests.test_musicxml
import tests.test_parts

def grip(frets, strings=(4, 3, 2, 1)):
  tuning = dict(generator.guitar.STANDARD_TUNING)
  return tuple(generator.guitar.TabPosition(s, f, tuning[s] + f) for s, f in zip(strings, frets))

def ranked(symbol, melody=(), profile='electric', prune=False):
  definition = generator.harmony.parse_symbol(symbol)
  requirements = generator.guitar.chord_requirements(definition, profile)
  candidates = generator.guitar.chord_positions(requirements, prune=prune, profile=profile, melody=melody)
  return sorted(candidates, key=lambda shape: (generator.guitar.voicing_cost(shape, profile, requirements.allowed, melody), shape))

class VoicingTests(unittest.TestCase):
  def test_requested_shapes_win_without_a_shape_dictionary(self):
    for symbol, melody, expected in (('FM7', 69, (3, 5, 5, 5)), ('Gm7', 67, (3, 3, 3, 3))):
      with self.subTest(symbol=symbol):
        self.assertEqual(ranked(symbol, (melody,))[0], grip(expected))
  def test_four_strings_precede_melody_height_and_root_presence(self):
    candidates = ranked('C#M7', (65,))
    four = next(c for c in candidates if len(c) == 4 and max(p.midi for p in c) > 65)
    three = next(c for c in candidates if len(c) == 3 and max(p.midi for p in c) == 65)
    self.assertLess(candidates.index(four), candidates.index(three))
    candidates = ranked('FM7', (69,))
    rootless = next(c for c in candidates if len(c) == 4 and 5 not in {p.midi % 12 for p in c})
    rooted = next(c for c in candidates if len(c) == 3 and 5 in {p.midi % 12 for p in c})
    self.assertLess(candidates.index(rootless), candidates.index(rooted))
  def test_melody_on_top_precedes_root_among_four_string_shapes(self):
    candidates = ranked('C#M7', (65,))
    best = candidates[0]
    self.assertEqual(len(best), 4)
    self.assertEqual(max(p.midi for p in best), 65)
    self.assertNotIn(1, {p.midi % 12 for p in best})
    self.assertTrue(any(1 in {p.midi % 12 for p in c} and max(p.midi for p in c) > 65 for c in candidates))
  def test_compact_chord_can_beat_lower_frets_without_covering_melody(self):
    candidates = ranked('CM7', (67,))
    self.assertEqual(candidates[0], grip((5, 4, 5, 3)))
    self.assertIn(grip((2, 4, 1, 3)), candidates)
    rooted = generator.guitar.voicing_cost(grip((2, 4, 1, 3)), 'electric', (0, 4, 7, 11), (67,))
    rootless = generator.guitar.voicing_cost(grip((5, 4, 5, 3)), 'electric', (0, 4, 7, 11), (67,))
    self.assertEqual((rooted.effort, rootless.effort), (19, 15))
    self.assertLess(rootless, rooted)
  def test_essential_extensions_are_never_relaxed_to_add_the_root(self):
    for symbol in ('C7b9', 'FM7#11', 'C13', 'Cdim7'):
      definition = generator.harmony.parse_symbol(symbol)
      requirements = generator.guitar.chord_requirements(definition, 'electric')
      for candidate in ranked(symbol):
        classes = {p.midi % 12 for p in candidate}
        self.assertTrue(requirements.covers(classes))
        self.assertTrue(classes <= set(requirements.allowed))
  def test_explicitly_removed_root_is_not_restored_or_penalized(self):
    document = tests.test_musicxml.parse(tests.test_musicxml.xml_score(tests.test_parts.xml_harmony(kind='major-seventh') + tests.test_musicxml.xml_note('E', octave=5)))
    measure = document['staves'][0]['measures'][0]
    chord = measure['harmonies'][0]
    chord['degrees'].append({'value': 1, 'alter': 0, 'type': 'subtract'})
    note = measure['voices'][0]['notes'][0]
    measure['transpose']['octave_change'] = -1
    note.update(_harmony=chord, _guitar_profile='electric', _harmonize=True)
    _, tones = generator.guitar.guitar_events(document)
    candidates = list(generator.guitar.attack_candidates([note], tones, {}, {}, False))
    self.assertTrue(candidates)
    self.assertTrue(all(all(p.midi % 12 != 0 for p in c.shape) and any(p.midi % 12 == 7 for p in c.shape) for c in candidates))
  def test_chord_height_cost_increases_throughout_the_fretboard(self):
    for profile in ('electric', 'acoustic'):
      for fret in range(1, 21):
        with self.subTest(profile=profile, fret=fret):
          self.assertEqual(generator.guitar.voicing_cost(grip((fret,) * 4), profile).effort, fret + 2 * max(0, fret - 12))
      self.assertEqual(generator.guitar.voicing_cost((), profile).effort, 0)
  def test_exact_melody_height_scoring_remains_unchanged(self):
    for fret in range(1, 21):
      self.assertEqual(generator.guitar.voicing_cost(grip((fret,) * 4)).effort, 2 * max(0, fret - 12))
    crossing = grip((11, 12, 13, 14))
    self.assertEqual(generator.guitar.voicing_cost(crossing).effort, 7)
    for profile in ('electric', 'acoustic'):
      self.assertEqual(generator.guitar.voicing_cost(crossing, profile).effort, 33)
  def test_chord_shape_and_height_share_one_effort_cost(self):
    shape = grip((1, 3, 4, 5))
    self.assertEqual(generator.guitar.voicing_cost(shape).effort, 4)
    for profile in ('electric', 'acoustic'):
      self.assertEqual(generator.guitar.voicing_cost(shape, profile).effort, 23)
    for profile in ('', 'electric', 'acoustic'):
      self.assertLess(generator.guitar.voicing_cost(grip((4, 4, 4, 4)), profile), generator.guitar.voicing_cost(grip((2, 4, 1, 3)), profile))
  def test_distinct_frets_add_difficulty_even_with_the_same_span(self):
    simple, varied = grip((3, 3, 5, 5)), grip((3, 4, 5, 5))
    for profile in ('electric', 'acoustic'):
      self.assertEqual(generator.guitar.voicing_cost(simple, profile).effort, 13)
      self.assertEqual(generator.guitar.voicing_cost(varied, profile).effort, 15)
    self.assertEqual(generator.guitar.voicing_cost(simple).effort, generator.guitar.voicing_cost(varied).effort)
  def test_compact_mid_neck_acoustic_pair_beats_high_and_stretched_low_pairs(self):
    high = [grip((12, 12, 12, 11), (6, 4, 3, 2)), grip((12, 11, 12, 11), (6, 4, 3, 2))]
    requested = [grip((7, 8, 7, 8), (5, 4, 3, 2)), grip((7, 8, 6, 8), (5, 4, 3, 2))]
    low = [grip((6, 5, 3, 3), (6, 4, 3, 2)), grip((5, 5, 3, 2), (6, 4, 3, 2))]
    for symbol, upper, middle, lower in zip(('Em7b5', 'A7b9'), high, requested, low):
      requirements = generator.guitar.chord_requirements(generator.harmony.parse_symbol(symbol), 'acoustic')
      candidates = generator.guitar.chord_positions(requirements, profile='acoustic')
      for shape in (upper, middle, lower):
        self.assertIn(shape, candidates)
      self.assertEqual({p.midi % 12 for p in upper}, {p.midi % 12 for p in middle})
      self.assertEqual(min(p.midi for p in upper), min(p.midi for p in middle))
    costs = [sum(generator.guitar.voicing_cost(s, 'acoustic').effort for s in pair) + 2 * generator.guitar.acoustic_movement(*pair) for pair in (requested, high, low)]
    self.assertEqual(costs, [33, 36, 41])
  def test_real_autumn_leaves_lead_in_uses_requested_compact_grips(self):
    source = generator.musicxml.load_musicxml(pathlib.Path('data/songs/autumn-leaves.xml'))
    before = copy.deepcopy(source)
    result = generator.arrange.generate.arrange_document(source, 'acoustic', rhythms=generator.rhythm.load_library(pathlib.Path('data/rhythms')))
    generator.model.validate_document(result)
    measures = result['staves'][0]['measures'][:2]
    self.assertEqual(measures[0]['section_label'], 'Rhythm Lead-in')
    for measure, symbol, expected in zip(measures, ('Em7b5', 'A7b9'), (((5, 7), (4, 8), (3, 7), (2, 8)), ((5, 7), (4, 8), (3, 6), (2, 8)))):
      with self.subTest(chord=symbol):
        self.assertEqual(generator.harmony.symbol(measure['harmonies'][0]), symbol)
        positions = {(p['tablature']['string'], p['tablature']['fret']) for v in measure['voices'] for n in v['notes'] for p in n['pitches']}
        self.assertEqual(tuple(sorted(positions, reverse=True)), expected)
    self.assertEqual(source, before)
  def test_acoustic_simpler_three_note_shape_can_beat_four_notes(self):
    candidates = ranked('Csus', profile='acoustic')
    best = candidates[0]
    self.assertEqual(len(best), 3)
    self.assertTrue(any(len(c) == 4 for c in candidates))
    requirements = generator.guitar.chord_requirements(generator.harmony.parse_symbol('Csus'), 'acoustic')
    self.assertTrue(requirements.covers({p.midi % 12 for p in best}))
  def test_nearby_position_beats_lower_frets_within_comfort_range(self):
    low, nearby, previous = (grip((f,) * 4) for f in (3, 8, 9))
    for profile in ('', 'electric', 'acoustic'):
      self.assertLess(generator.guitar.voicing_cost(nearby, profile).moved(generator.guitar.movement(previous, nearby)), generator.guitar.voicing_cost(low, profile).moved(generator.guitar.movement(previous, low)))
  def test_above_threshold_is_a_soft_cost_not_a_range_priority(self):
    low, high = (grip((f,) * 4) for f in (3, 13))
    for profile in ('', 'electric', 'acoustic'):
      self.assertLess(generator.guitar.voicing_cost(high, profile), generator.guitar.voicing_cost(low, profile).moved(generator.guitar.movement(high, low)))
  def test_candidate_pruning_uses_full_cost_and_preserves_each_center(self):
    for symbol, profile, melody in (('FM7', 'electric', (69,)), ('Gm7', 'electric', (67,))):
      best = {}
      for shape in ranked(symbol, melody, profile):
        best.setdefault(generator.guitar.position_center(shape), shape)
      self.assertEqual(set(ranked(symbol, melody, profile, prune=True)), set(best.values()))
  def test_tie_continuations_do_not_charge_a_second_voicing_cost(self):
    note = tests.test_musicxml.xml_note('E', 2, 'half', octave=5, extra='<tie type="start"/>') + tests.test_musicxml.xml_note('E', 2, 'half', octave=5, extra='<tie type="stop"/>')
    document = tests.test_musicxml.parse(tests.test_musicxml.xml_score(note))
    events, tones = generator.guitar.guitar_events(document)
    states = {(): generator.guitar.Placement(generator.guitar.VoicingCost(), (), None, ())}
    first = generator.guitar.advance_positions([events[0][3]], events[0][0], tones, states, False)
    second = generator.guitar.advance_positions([events[1][3]], events[1][0], tones, first, False)
    self.assertEqual(sorted(p.cost for p in first.values()), sorted(p.cost for p in second.values()))
  def test_empty_and_positioned_states_are_compared_by_the_same_grid_cost(self):
    empty = generator.guitar.Placement(generator.guitar.VoicingCost(effort=20), (), None, ())
    placed = generator.guitar.Placement(generator.guitar.VoicingCost(), (('tone', grip((3,), (1,))[0]),), None, ())
    grid = generator.guitar.nearest_positions({(): empty, 1: placed})
    self.assertIs(grid[6][0][2], placed)
    self.assertIs(grid[40][10][2], empty)
  def test_optimized_passage_search_matches_exhaustive_candidate_ranking(self):
    contents = ''.join(tests.test_musicxml.xml_note('E', 1, 'quarter', octave=5) for _ in range(4))
    document = tests.test_musicxml.parse(tests.test_musicxml.xml_score(contents))
    measure = document['staves'][0]['measures'][0]
    for note, symbol in zip(measure['voices'][0]['notes'], ('FM7', 'Gm7', 'C7b9', 'FM7')):
      note.update(pitches=[], _harmony=generator.harmony.parse_symbol(symbol), _guitar_profile='electric')
    events, tones = generator.guitar.guitar_events(document)
    states = {(): generator.guitar.Placement(generator.guitar.VoicingCost(), (), None, ())}
    for onset, _, _, note in events:
      expected = {}
      for candidate in generator.guitar.attack_candidates([note], tones, {}, {}, False):
        choices = []
        for index, previous in enumerate(states.values()):
          distance = generator.guitar.movement(tuple(p for _, p in previous.active), candidate.shape)
          choices.append((previous.cost.moved(distance).plus(candidate.cost), index))
        expected[candidate.active] = min(choices)
      actual = generator.guitar.advance_positions([note], onset, tones, states, False)
      predecessors = {id(p): i for i, p in enumerate(states.values())}
      self.assertEqual({key: (p.cost, predecessors[id(p.previous)]) for key, p in actual.items()}, expected)
      states = actual
  def test_real_autumn_leaves_melody_has_expected_shapes_and_unchanged_notes(self):
    source = generator.musicxml.load_musicxml(pathlib.Path('data/songs/autumn-leaves.xml'))
    before = copy.deepcopy(source)
    result = generator.arrange.generate.arrange_document(source, 'electric', rhythms=tests.rhythm_fixtures.library())
    generator.model.validate_document(result)
    section, checked = '', set()
    for measure in result['staves'][0]['measures']:
      section = measure.get('section_label') or section
      if section != 'Guitar Melody':
        continue
      original = next(m for m in source['staves'][0]['measures'] if m['number'] == measure['number'])
      for note in measure['voices'][0]['notes']:
        if not note['pitches']:
          continue
        onset = generator.model.fraction(note['onset'])
        authored = next(n for n in original['voices'][0]['notes'] if generator.model.fraction(n['onset']) <= onset < generator.model.fraction(n['onset']) + generator.model.fraction(n['duration']))
        self.assertEqual(generator.model.pitch_number(note['pitches'][0]), generator.model.pitch_number(authored['pitches'][0]) + 12)
        if measure['number'] not in ('4', '12', '22', '24') or len(note['pitches']) == 1:
          continue
        expected = (3, 3, 3, 3) if measure['number'] == '22' else (3, 5, 5, 5)
        frets = tuple(p['tablature']['fret'] for p in sorted(note['pitches'], key=lambda p: -p['tablature']['string']))
        self.assertEqual(frets, expected)
        self.assertEqual(generator.model.pitch_number(note['pitches'][0]), max(generator.model.pitch_number(p) for p in note['pitches']))
        checked.add(measure['number'])
    self.assertEqual(checked, {'4', '12', '22', '24'})
    self.assertEqual(source, before)