import dataclasses
from fractions import Fraction as F
import unittest

from generator.solo import types, harmonic, enclosures, phrasing, composer, placement, notation, report
import generator.harmony
import generator.model
import tests.scale_fixtures
import tests.solo_fixtures
from tests.test_solo_passages import choices
from tests.test_solo import song


def note(start,end,letter,alter=0,octave=4,identity=''):
  return types.Note(F(start),F(end),types.Pitch(letter,F(alter),octave),event_id=identity or f'{start}:{letter}')


def fixture(symbol='C7',target=None,scale=None):
  target=target or note(4,8,'e',identity='target')
  harmony=harmonic.Harmony(choices((symbol,symbol),scale=scale))
  pair=enclosures.approach(target,harmony)
  contract=types.EnclosureContract('enclosure',*pair,target,harmony.at(target.start).chord,'terminal',types.SustainedTail(target.start,target.end,target.pitch.number()))
  return harmony,(*pair,target),contract


class EnclosureContracts(unittest.TestCase):
  def test_existing_upper_can_continue_without_a_new_attack(self):
    h,notes,c=fixture('Dm',note(4,8,'d'))
    upper=dataclasses.replace(notes[0],start=F(5,2))
    contract=dataclasses.replace(c,upper=upper,approach_start=F(3))
    self.assertTrue(h.evaluate((upper,*notes[1:]),enclosures=(contract,)).valid)
    # Joining a tie must not hide an exposed long non-chord upper neighbor.
    long_upper=dataclasses.replace(upper,start=F(1))
    invalid=dataclasses.replace(contract,upper=long_upper)
    self.assertFalse(h.evaluate((long_upper,*notes[1:]),enclosures=(invalid,)).valid)
    split=(dataclasses.replace(upper,end=F(3)),dataclasses.replace(upper,start=F(3),continuation=True,event_id='tie'),*notes[1:])
    self.assertEqual(h.evaluate(split,enclosures=(contract,)),h.evaluate((upper,*notes[1:]),enclosures=(contract,)))

  def test_worked_examples_and_actual_collection_membership(self):
    hm=tests.scale_fixtures.library().scale('harmonic_minor',1,'D')
    for symbol,target,scale,expected in (
      ('C7',note(4,8,'e'),None,('F4','D#4','E4')),
      ('FM7',note(4,8,'a',octave=3),None,('Bb3','G#3','A3')),
      ('A7',note(4,8,'c',1),hm,('D4','C4','C#4')),
      ('Dm',note(4,8,'d'),None,('E4','C#4','D4'))):
      with self.subTest(symbol=symbol):
        harmony,notes,contract=fixture(symbol,target,scale)
        self.assertEqual(tuple(n.pitch.label() for n in notes),expected)
        self.assertFalse(harmony.evaluate(notes).valid)
        self.assertTrue(harmony.evaluate(notes,enclosures=(contract,),complete=True).valid)
        self.assertEqual(harmony.evaluate(notes,enclosures=(contract,)).exposure,0)
        self.assertEqual(sum(r['scale_exception'] for r in harmony.describe(notes,enclosures=(contract,))),1)

  def test_in_map_enclosure_does_not_need_exception(self):
    h,notes,c=fixture('FM7',note(4,8,'f'))
    self.assertEqual([n.pitch.label() for n in notes],['G4','E4','F4'])
    self.assertTrue(h.evaluate(notes).valid)
    self.assertTrue(h.evaluate(notes,enclosures=(c,)).valid)
    self.assertFalse(any(r['scale_exception'] for r in h.describe(notes,enclosures=(c,))))

  def test_missing_changed_target_wrong_octave_and_repeated_attack_fail(self):
    h,notes,c=fixture()
    for actual in (notes[:-1],(*notes[:-1],dataclasses.replace(notes[-1],pitch=types.Pitch('e',F(),3))),
      (*notes[:-1],dataclasses.replace(notes[-1],start=F(9,2))),
      (*notes[:-1],dataclasses.replace(notes[-1],end=F(6)),dataclasses.replace(notes[-1],start=F(6),event_id='reattack'))):
      self.assertFalse(h.evaluate(actual,enclosures=(c,)).valid)
    self.assertFalse(h.evaluate(notes,enclosures=(c,),until=F(7),complete=True).valid)

  def test_roles_alone_never_authorize_chromatic_pitches(self):
    h,notes,c=fixture()
    for role in ('enclosure_lower','chromatic_approach','quote'):
      actual=tuple(dataclasses.replace(n,role=role) for n in notes)
      self.assertFalse(h.evaluate(actual).valid)

  def test_complete_sounding_target_is_partition_invariant(self):
    h,notes,c=fixture()
    split=(*notes[:-1],dataclasses.replace(notes[-1],end=F(6)),dataclasses.replace(notes[-1],start=F(6),continuation=True,event_id='tie'))
    self.assertEqual(h.evaluate(split,enclosures=(c,)),h.evaluate(notes,enclosures=(c,)))
    self.assertEqual(h.describe(split,enclosures=(c,)),h.describe(notes,enclosures=(c,)))
    self.assertEqual(types.sounding_notes(split),notes)

  def test_orphan_long_approach_order_and_duplicate_contracts_fail(self):
    h,notes,c=fixture()
    variants=(dataclasses.replace(c,upper=dataclasses.replace(c.upper,start=F(2))),
      dataclasses.replace(c,lower=dataclasses.replace(c.lower,pitch=types.Pitch('f',F(1),4))),
      dataclasses.replace(c,target=dataclasses.replace(c.target,event_id='missing')),
      dataclasses.replace(c,chord='Dm'))
    for changed in variants:
      actual=(changed.upper,changed.lower,changed.target)
      if changed.target.event_id=='missing':actual=notes
      self.assertFalse(h.evaluate(actual,enclosures=(changed,)).valid)
    self.assertFalse(h.evaluate(notes,enclosures=(c,c)).valid)

  def test_boundary_crossing_uses_arrival_chord_and_each_actual_scale(self):
    hm=tests.scale_fixtures.library().scale('harmonic_minor',1,'D')
    cs=choices(('A7','Dm'))
    cs=(dataclasses.replace(cs[0],scale=hm),cs[1])
    h=harmonic.Harmony(cs);target=note(4,8,'d')
    upper,lower=enclosures.approach(target,h)
    c=types.EnclosureContract('boundary',upper,lower,target,'Dm','terminal',types.SustainedTail(F(4),F(8),62))
    self.assertTrue(h.evaluate((upper,lower,target),enclosures=(c,)).valid)
    self.assertTrue(h.evaluate((upper,lower,target)).valid) # C# still belongs to A7's map
    self.assertFalse(h.evaluate((upper,lower,target),enclosures=(dataclasses.replace(c,chord='A7'),)).valid)
    missing=harmonic.Harmony((dataclasses.replace(cs[0],scale=None),cs[1]))
    self.assertFalse(missing.evaluate((upper,lower,target),enclosures=(c,)).valid)

  def test_held_target_must_remain_chordal_after_boundary(self):
    h,notes,c=fixture('Gm7',note(2,6,'f'))
    h=harmonic.Harmony(choices(('Gm7','C7')))
    self.assertFalse(h.evaluate(notes,enclosures=(c,)).valid)

  def test_transposition_preserves_contract_validation(self):
    h,notes,c=fixture()
    for octave in (-1,0,1):
      moved=tuple(dataclasses.replace(n,pitch=dataclasses.replace(n.pitch,octave=n.pitch.octave+octave)) for n in notes)
      contract=dataclasses.replace(c,upper=moved[0],lower=moved[1],target=moved[2],protected_tail=dataclasses.replace(c.protected_tail,pitch=moved[2].pitch.number()))
      self.assertTrue(h.evaluate(moved,enclosures=(contract,)).valid)

  def test_spelling_and_notation_ties_tab_and_playback_survive(self):
    import generator.playback
    import generator.musicxml_writer
    h,notes,c=fixture(target=note(3,7,'e',identity='target'))
    placed=placement.place(notes,F(120))[1]
    doc=song(count=2,chords=('C7',))
    measures=notation.measures(doc['staves'][0]['measures'],placed,'enclosure_test')
    notation.replace_region(doc,0,measures)
    generator.model.validate_document(doc)
    import generator.fretboard
    generator.fretboard.validate_document(doc)
    pitches=[p for m in measures for v in m['voices'] for n in v['notes'] for p in n['pitches']]
    self.assertTrue(any(p.get('tie_from') for p in pitches))
    self.assertTrue(any(p['letter']=='d' and generator.model.fraction(p['alter'])==1 for p in pitches))
    self.assertTrue(all(1<=p['tablature']['fret']<=12 for p in pitches))
    # The guitar's written octave must be removed exactly once at playback.
    doc['staves'][0]['measures'][0]['transpose']={'diatonic':-7,'chromatic':-12,'octave_change':0}
    doc['staves'][0]['measures'][1]['transpose']={'diatonic':-7,'chromatic':-12,'octave_change':0}
    timeline=generator.playback.compile_timeline(doc)
    self.assertEqual([generator.model.fraction(e['pitch']) for e in timeline['events']],[65,63,64])
    self.assertEqual([generator.model.fraction(e['duration']) for e in timeline['events']],[F(1,2),F(1,2),F(4)])
    tree=generator.musicxml_writer.musicxml_tree(doc)
    self.assertTrue(any(e.text=='1' for e in tree.findall('.//pitch/alter')))


class DestinationPriorities(unittest.TestCase):
  def test_actual_symbols_define_members_without_slash_or_omitted_degrees(self):
    cases={'C7':{4,10},'C':{4},'Cm':{3},'C6':{4,9},'Csus4':{5},'Csus2':{2},'C7b5':{4,6,10},'C7no3':{10},'C/D':{4},'C9':{4,10}}
    for symbol,expected in cases.items():
      with self.subTest(symbol=symbol):self.assertEqual(phrasing.defining_members(symbol),expected)

  def test_guide_tones_compete_before_nearest_register_cutoff(self):
    from generator.solo.themes import pitches
    h=harmonic.Harmony(choices(('C7','FM7')))
    notes=pitches(h,note(2,4,'c'),60,members=True,limit=4)
    self.assertEqual({p.number()%12 for p in notes},{0,4,7,10})
    cost=lambda p:phrasing.destination_cost(p,F(2),F(4),h,goal=60)
    self.assertLess(cost(types.Pitch('e',F(),4)),cost(types.Pitch('c',F(),4)))
    self.assertLess(cost(types.Pitch('b',F(-1),3)),cost(types.Pitch('g',F(),3)))


def solo_fixture():
  h=harmonic.Harmony(choices(('C7','C7')))
  notes=(note(0,1,'g'),note(1,2,'e'),note(2,3,'g'),note(3,4,'f'),note(4,8,'e',identity='target'))
  notes=placement.place(notes,F(120))[1]
  anchor=types.Anchor('close',F(4),F(4),notes[-1].pitch,frozenset({4}),(41,76),F(4),'terminal','C7',hold_until=F(8))
  motif=types.Motif('theme',notes[:3],F(3),())
  plan=types.PhrasePlan('test',F(),F(8),motif,'develop','thematic_recomposition',None,(41,76),F(1),F(),ending='sustain',closure='closed')
  transform=types.Transformation('develop','','theme','develop',F(),F(8),tuple(n.event_id for n in notes),tuple((n.start,n.end) for n in notes),preserves=())
  draft=types.MelodicDraft('draft',F(),F(8),notes,(transform,),(),(),(anchor,),None,'develop',sustained_tails=(types.SustainedTail(F(4),F(8),64),))
  decision=types.Decision(F(),F(8),'fixture','develop',notes)
  state=types.ContinuationState('test',notes=notes,last=notes[-1])
  phrase=types.Phrase(plan,(decision,),state=state,anchors=(anchor,),draft=draft)
  return types.Solo('acoustic',h.choices,(phrase,),F(120),F(8)),h


class EnclosureRefinement(unittest.TestCase):
  def test_atomic_candidates_preserve_target_timing_and_protections(self):
    solo,h=solo_fixture();phrase=solo.phrases[0]
    options=enclosures.candidates(phrase,h,tests.solo_fixtures.library(),solo.tempo,[])
    terminal=[p for _,p in options if p.draft.enclosures[0].target_role=='terminal']
    self.assertTrue(terminal)
    for p in terminal:
      self.assertEqual(p.draft.enclosures[0].target,phrase.draft.notes[-1])
      self.assertTrue(h.evaluate(p.draft.notes,enclosures=p.draft.enclosures).valid)
      actual=placement.place(p.draft.notes,solo.tempo)[1]
      placed=dataclasses.replace(p,decisions=(dataclasses.replace(p.decisions[0],notes=actual),),state=dataclasses.replace(p.state,notes=actual))
      composer.validate(dataclasses.replace(solo,phrases=(placed,)))
    protected=dataclasses.replace(phrase,draft=dataclasses.replace(phrase.draft,protected=tuple(n.event_id for n in phrase.draft.notes)))
    self.assertFalse(enclosures.candidates(protected,h,tests.solo_fixtures.library(),solo.tempo,[]))

  def test_intermediate_enclosure_leaves_continuation_and_terminal_hold_intact(self):
    solo,h=solo_fixture();phrase=solo.phrases[0]
    options=enclosures.candidates(phrase,h,tests.solo_fixtures.library(),solo.tempo,[])
    intermediate=[p for _,p in options if p.draft.enclosures[0].target_role=='intermediate']
    self.assertTrue(intermediate)
    for candidate in intermediate:
      target=candidate.draft.enclosures[0].target
      self.assertTrue(any(n.start>target.start for n in candidate.draft.notes))
      self.assertEqual(candidate.draft.notes[-1],phrase.draft.notes[-1])
      self.assertEqual(candidate.state.home,phrase.state.home)

  def test_short_score_can_choose_direct_arrivals_without_a_quota(self):
    solo,h=solo_fixture();phrase=solo.phrases[0]
    protected=dataclasses.replace(phrase,draft=dataclasses.replace(phrase.draft,protected=tuple(n.event_id for n in phrase.draft.notes)))
    result=enclosures.refine(dataclasses.replace(solo,phrases=(protected,)),tests.solo_fixtures.library())
    self.assertEqual(result.diagnostics['enclosures_selected'],0)
    self.assertEqual([(n.pitch,n.start,n.end) for n in result.notes()],[(n.pitch,n.start,n.end) for n in solo.notes()])
    self.assertTrue(result.diagnostics['enclosure_rejections'])

  def test_already_sounding_entrance_is_not_detached_for_ornamentation(self):
    solo,h=solo_fixture();p=solo.phrases[0]
    held=dataclasses.replace(p.draft.notes[0],continuation=True)
    p=dataclasses.replace(p,draft=dataclasses.replace(p.draft,notes=(held,*p.draft.notes[1:])))
    self.assertFalse(any(n.event_id==held.event_id for n,_ in enclosures.declared_targets(p,h)))
    rejected=[]
    self.assertFalse(enclosures.candidates(p,h,tests.solo_fixtures.library(),solo.tempo,rejected))
    self.assertIn('sustain retained',rejected[0]['reason'])
    self.assertTrue(p.draft.notes[0].continuation)

  def test_no_gap_filling_and_no_unplayable_enclosure(self):
    solo,h=solo_fixture();p=solo.phrases[0]
    draft=dataclasses.replace(p.draft,notes=tuple(n for n in p.draft.notes if n.start!=3))
    candidates=enclosures.candidates(dataclasses.replace(p,draft=draft),h,tests.solo_fixtures.library(),solo.tempo,[])
    self.assertFalse(any(c.draft.enclosures[0].target.event_id=='target' for _,c in candidates))
    self.assertIsNone(enclosures.approach(note(4,8,'e',octave=5),h))

  def test_disabled_is_exact_strict_input_and_selection_is_deterministic(self):
    solo,h=solo_fixture();library=tests.solo_fixtures.library()
    disabled=enclosures.refine(solo,library,False)
    self.assertEqual(disabled.notes(),solo.notes())
    first=enclosures.refine(solo,library);second=enclosures.refine(solo,library)
    self.assertEqual(first.notes(),second.notes())
    self.assertEqual(enclosures.strict(first).notes(),solo.notes())
    composer.validate(first);composer.validate(disabled)
    self.assertTrue(report.structure(first)['harmonic_valid'])


class EnclosureCompositionAPI(unittest.TestCase):
  def test_disable_joint_expressions_selects_a_valid_strict_alternative(self):
    from generator.solo import analysis
    import tests.rhythm_fixtures
    doc=song(count=8,chords=('Gm7','C7','FM7','BbM7','Em7b5','A7','Dm7','Dm7'))
    a=analysis.analyze(doc)
    args=(a,doc['staves'][0]['measures'],'acoustic',F(32))
    kwargs=dict(scales=tests.scale_fixtures.library(),sequences=tests.solo_fixtures.library(),rhythms=tests.rhythm_fixtures.library())
    decorated=composer.compose(*args,**kwargs)
    strict=composer.compose(*args,**kwargs,enclosures_enabled=False)
    # A jointly composed color/anticipation can change phrase selection. Its
    # disabled alternative need not be the selected decorated phrase minus
    # notes; both versions must retain the same form and valid sounding spans.
    self.assertEqual(decorated.end,strict.end)
    composer.validate(decorated)
    composer.validate(strict)
    self.assertTrue(harmonic.Harmony(strict.choices).evaluate(strict.notes()).valid)
    self.assertTrue(all(not p.draft.enclosures for p in strict.phrases))
    self.assertGreater(sum(v['proposed'] for v in decorated.diagnostics['expression_operations'].values()),0)
    self.assertTrue(all(not p.draft.expressions for p in strict.phrases))

if __name__=='__main__':unittest.main()