import argparse
import copy
from fractions import Fraction
from pathlib import Path
import random
import tempfile
import unittest

from generator import harmony
from generator.__main__ import generate_assets
from generator.arrange.generate import arrange_document
from generator.guitar import Placement, TabPosition, chord_positions, chord_requirements, nearest_positions, place_guitar, position_center, positions_for_pitch, within_fret_span
from generator.model import fraction, pitch_number, rational, validate_document
from generator.parts import PART_TYPES
from generator.pitch import concert_document
from generator.playback import compile_timeline
from generator.practice.generate import generate_practice_scores, practice_document
from generator.summary import validate_practice_parts
from tests.test_musicxml import notes, parse, xml_note, xml_score

class GuitarTests(unittest.TestCase):
  def test_fret_boundaries_and_span(self):
    self.assertIn(TabPosition(6, 0, 40), positions_for_pitch(Fraction(40)))
    self.assertIn(TabPosition(1, 20, 84), positions_for_pitch(Fraction(84)))
    self.assertFalse(positions_for_pitch(Fraction(85)))
    self.assertFalse(positions_for_pitch(Fraction(129, 2)))
    self.assertTrue(within_fret_span((TabPosition(4, 11, 61), TabPosition(3, 12, 67), TabPosition(2, 14, 73))))
    self.assertTrue(within_fret_span((TabPosition(6, 0, 40), TabPosition(1, 20, 84))))
    self.assertFalse(within_fret_span((TabPosition(4, 10, 60), TabPosition(1, 20, 84))))
  def test_theory_controls_omissions_and_bass(self):
    for symbol in ('CM', 'G7b9', 'FM7#11/A', 'CmM7', 'Cdim7', 'C13', 'Csus', 'Cmadd9'):
      with self.subTest(symbol=symbol):
        required, allowed, bass = chord_requirements(harmony.parse_symbol(symbol))
        candidates = chord_positions(required, allowed, bass)
        self.assertTrue(candidates)
        for candidate in candidates:
          covered = {p.midi % 12 for p in candidate}
          self.assertTrue(set(required) <= covered <= set(allowed))
          self.assertEqual(len(candidate), len({p.string_number for p in candidate}))
          self.assertTrue(within_fret_span(candidate))
          if bass is not None:
            self.assertEqual(min(p.midi for p in candidate) % 12, bass)
    self.assertIn(2, chord_requirements(harmony.parse_symbol('Cmadd9'))[0])
    self.assertNotIn(7, chord_requirements(harmony.parse_symbol('C7'))[0])
    self.assertIn(7, chord_requirements(harmony.parse_symbol('CM'))[0])
    self.assertEqual(chord_requirements({**harmony.parse_symbol('CM'), 'inversion': 0})[2], 0)
    self.assertEqual(chord_requirements({**harmony.parse_symbol('CM'), 'inversion': 1})[2], 4)
    self.assertTrue(any(p.fret > 12 for c in chord_positions(*chord_requirements(harmony.parse_symbol('CM'))) for p in c))
  def test_melody_placement_preserves_notes_timing_and_sound(self):
    source = parse(xml_score(xml_note('E', 2, 'half', octave=4) + xml_note('F', 2, 'half', octave=4)))
    before = copy.deepcopy(source)
    electric = practice_document(source, 'electric')
    vocal = practice_document(source, 'vocal')
    validate_document(electric)
    self.assertEqual(source, before)
    self.assertEqual(compile_timeline(electric)['events'], compile_timeline(vocal)['events'])
    self.assertEqual([(e['pitch'], e['onset'], e['duration']) for e in compile_timeline(electric)['events']], [(rational(fraction(e['pitch']) - 12), e['onset'], e['duration']) for e in compile_timeline(source)['events']])
    for original, generated in zip(notes(source), notes(electric)):
      self.assertEqual(pitch_number(generated['pitches'][0]), pitch_number(original['pitches'][0]))
      self.assertIn('tablature', generated['pitches'][0])
  def test_instrument_transposition_and_octave_clef_applied_once(self):
    for attributes, expected in (('', 48), ('<transpose><chromatic>-12</chromatic><diatonic>-7</diatonic></transpose>', 48), ('<transpose><chromatic>0</chromatic></transpose>', 60), ('<clef><sign>G</sign><line>2</line><clef-octave-change>-1</clef-octave-change></clef>', 60), ('<clef><sign>F</sign><line>4</line></clef>', 48)):
      source = parse(xml_score(xml_note(), attributes))
      for part in ('electric', 'vocal'):
        document = practice_document(source, part)
        self.assertEqual(fraction(compile_timeline(document)['events'][0]['pitch']), expected)
  def test_written_middle_c_and_c_sharp_six_have_matching_tab_and_audio(self):
    for octave, alter, expected, tab in ((4, 0, 48, {'string': 5, 'fret': 3}), (6, 1, 73, {'string': 1, 'fret': 9}), (7, 0, 84, {'string': 1, 'fret': 20})):
      source = parse(xml_score(xml_note(octave=octave, alter=alter)))
      document = practice_document(source, 'electric')
      tone = notes(document)[0]['pitches'][0]
      self.assertEqual(pitch_number(tone), 12 * (octave + 1) + alter)
      self.assertEqual(tone['tablature'], tab)
      self.assertEqual(fraction(compile_timeline(document)['events'][0]['pitch']), expected)
      validate_document(document)
  def test_guitar_normalization_is_idempotent_and_respects_timed_transposition(self):
    change = '<attributes><transpose><chromatic>0</chromatic></transpose></attributes>'
    source = parse(xml_score(xml_note(duration=2, kind='half') + change + xml_note(duration=2, kind='half')))
    before = copy.deepcopy(source)
    concert = concert_document(source, octave_transposing=True)
    self.assertEqual(concert_document(concert, octave_transposing=True), concert)
    self.assertEqual(source, before)
    document = practice_document(source, 'electric')
    self.assertEqual([fraction(e['pitch']) for e in compile_timeline(document)['events']], [48, 60])
  def test_all_parts_share_written_octave_convention_in_songs_and_practice(self):
    chord = '<harmony><root><root-step>C</root-step></root><kind>major</kind></harmony>'
    source = parse(xml_score(chord + xml_note()))
    for generate in (practice_document, arrange_document):
      for part in PART_TYPES:
        with self.subTest(generator=generate.__name__, part=part):
          document = generate(source, part)
          validate_document(document)
          written = {p['id']: pitch_number(p) for m in document['staves'][0]['measures'] for v in m['voices'] for n in v['notes'] for p in n['pitches']}
          events = compile_timeline(document)['events']
          self.assertTrue(events)
          self.assertTrue(all(fraction(e['pitch']) == written[e['id']] - 12 for e in events))
          if part == 'vocal':
            self.assertEqual(set(written.values()), {60})
          self.assertTrue(all(m['transpose']['octave_change'] == -1 for m in document['staves'][0]['measures']))
  def test_ties_and_overlapping_voices_keep_strings(self):
    content = xml_note('E', 1, 'quarter', octave=4, extra='<tie type="start"/>') + xml_note('E', 3, 'half', octave=4, extra='<dot/><tie type="stop"/>')
    content += '<backup><duration>4</duration></backup><forward><duration>1</duration></forward>' + xml_note('F', 3, 'half', octave=4, extra='<dot/><voice>2</voice>')
    document = parse(xml_score(content))
    place_guitar(document)
    first, continuation, other = [p['tablature'] for m in document['staves'][0]['measures'] for v in m['voices'] for n in v['notes'] for p in n['pitches']]
    self.assertEqual(first, continuation)
    self.assertNotEqual(continuation['string'], other['string'])
    self.assertEqual(len(compile_timeline(document)['events']), 2)
  def test_proximity_chooses_a_path_without_relocating_tied_notes(self):
    source = parse(xml_score(xml_note('B', 2, 'half', octave=5) + xml_note('C', 2, 'half', octave=6)))
    place_guitar(source)
    tabs = [n['pitches'][0]['tablature'] for n in notes(source)]
    self.assertEqual(tabs, [{'string': 1, 'fret': 19}, {'string': 1, 'fret': 20}])
  def test_grid_search_matches_exhaustive_costs_and_tie_breaks(self):
    randomizer = random.Random(17)
    states = {}
    for index in range(40):
      positions = tuple(TabPosition(s, randomizer.randrange(21), 0) for s in randomizer.sample(range(1, 7), 3))
      states[index] = Placement(tuple(randomizer.randrange(8) for _ in range(3)), tuple((str(i), p) for i, p in enumerate(positions)), None, ())
    grid = nearest_positions(states)
    for x, row in enumerate(grid):
      for y, actual in enumerate(row):
        candidates = []
        for index, state in enumerate(states.values()):
          a, b = position_center(tuple(p for _, p in state.active))
          candidates.append(((state.cost[0] + abs(x - a) + abs(y - b), *state.cost[1:]), index))
        self.assertEqual(actual[:2], min(candidates))
  def test_generated_chord_retains_strings_when_a_melody_enters(self):
    content = xml_note() + '<backup><duration>4</duration></backup><forward><duration>2</duration></forward>' + xml_note('C', 2, 'half', octave=7, extra='<voice>2</voice>')
    document = parse(xml_score(content))
    measure = document['staves'][0]['measures'][0]
    measure['transpose']['octave_change'] = -1
    chord, melody = [n for v in measure['voices'] for n in v['notes']]
    chord['pitches'] = []
    chord['_harmony'] = harmony.parse_symbol('CM')
    place_guitar(document)
    validate_document(document)
    self.assertEqual(melody['pitches'][0]['tablature'], {'string': 1, 'fret': 20})
    self.assertNotIn(1, [p['tablature']['string'] for p in chord['pitches']])
    self.assertEqual(len(compile_timeline(document)['events']), 4)
  def test_more_than_six_simultaneous_tones_fail_without_omission(self):
    content = '<backup><duration>4</duration></backup>'.join(xml_note(extra=f'<voice>{i}</voice>') for i in range(7))
    document = parse(xml_score(content))
    before = copy.deepcopy(document)
    with self.assertRaisesRegex(ValueError, 'no guitar placement'):
      place_guitar(document)
    self.assertEqual(before, document)
  def test_source_contract_and_unplayable_music_fail(self):
    cases = [(xml_note(octave=7, alter=1), '', '0-20'), (xml_note(alter='0.5'), '', '0-20'),
      (xml_note() + xml_note('E', prefix='<chord/>'), '', 'monophonic'),
      (xml_note() + '<backup><duration>4</duration></backup>' + xml_note('E', extra='<voice>2</voice>'), '', 'one monophonic voice'),
      (xml_note(), '<staff-details><capo>1</capo></staff-details>', 'capo')]
    for content, attributes, error in cases:
      with self.subTest(error=error), self.assertRaisesRegex(ValueError, error):
        practice_document(parse(xml_score(content, attributes)), 'electric')
    for content, attributes in [(xml_note(), '<clef><sign>TAB</sign></clef>'), (xml_note(extra='<notations><technical><fingering>1</fingering></technical></notations>'), '')]:
      with self.assertRaisesRegex(ValueError, 'authored TAB'):
        parse(xml_score(content, attributes))
  def test_three_parts_are_deterministic_without_transposition(self):
    with tempfile.TemporaryDirectory() as temporary:
      source, output = Path(temporary) / 'source', Path(temporary) / 'output'
      source.mkdir()
      (source / '0000.xml').write_text(xml_score(xml_note('E', alter=-1), '<key><fifths>-3</fifths></key>'))
      self.assertEqual(generate_practice_scores(source, output), 3)
      validate_practice_parts(output)
      first = {p.name: p.read_bytes() for p in output.rglob('*.json')}
      generate_practice_scores(source, output)
      self.assertEqual(first, {p.name: p.read_bytes() for p in output.rglob('*.json')})
      (output / '0000/vocal.json').unlink()
      with self.assertRaisesRegex(ValueError, 'practice generator drift'):
        validate_practice_parts(output)
  def test_unplayable_generation_preserves_previous_assets(self):
    with tempfile.TemporaryDirectory() as temporary:
      root = Path(temporary)
      (root / 'data/songs').mkdir(parents=True)
      (root / 'data/practice').mkdir()
      (root / 'data/practice/0001.xml').write_text(xml_score(xml_note(octave=7, alter=1)))
      output = root / 'output'
      output.mkdir()
      (output / 'last-good').write_text('preserved')
      args = argparse.Namespace(data_dir=root / 'data', output_dir=output, static_assets_dir=Path('package/assets'), check=False)
      with self.assertRaisesRegex(ValueError, '0001.xml, part electric:.*0-20'):
        generate_assets(args)
      self.assertEqual([p.name for p in output.iterdir()], ['last-good'])
      self.assertEqual((output / 'last-good').read_text(), 'preserved')