"""Construction contracts, not claims about listener preference."""
import dataclasses
import fractions
import unittest

import generator.rhythm
import generator.solo.drafts
import generator.solo.harmonic
import generator.solo.realization
import generator.solo.scoring
import generator.solo.types
import tests.rhythm_fixtures
import tests.solo_fixtures
import tests.test_solo
import tests.test_solo_development
import tests.test_solo_passages

F = fractions.Fraction
note = tests.test_solo_passages.note

def illustration():
  """User-specified four-bar example; additional fixtures await listening review."""
  events = ((0, 1, 'd', 0, 4), (2, 3, 'f', 0, 4), (3, F(7, 2), 'g', 0, 4), (F(7, 2), 4, 'a', 0, 4),
    (4, 5, 'b', -1, 4), (5, 6, 'g', 0, 4), (6, F(13, 2), 'a', 0, 4), (F(13, 2), 7, 'g', 0, 4), (7, 8, 'e', 0, 4),
    (8, 9, 'f', 0, 4), (9, F(19, 2), 'g', 0, 4), (F(19, 2), 10, 'a', 0, 4), (10, 11, 'c', 0, 5), (11, 13, 'a', 0, 4),
    (13, F(27, 2), 'g', 0, 4), (F(27, 2), 14, 'f', 0, 4), (14, 16, 'd', 0, 4))
  notes = tuple(dataclasses.replace(note(a, b, pitch, alter, octave), event_id=f'example_{i}') for i, (a, b, pitch, alter, octave) in enumerate(events))
  choices = tests.test_solo_passages.choices(('Gm7', 'C7', 'FM7', 'BbM7'))
  harmony = generator.solo.harmonic.Harmony(choices)
  home = generator.solo.realization.make_identity('home', 'user_example', notes[1:5], F(2), F(5), 'statement', harmony)
  groups = []
  for identifier, operation, selected in (('home', 'statement', notes[1:5]), ('extension', 'extend', notes[9:13])):
    start = selected[0].start
    groups.append(generator.solo.types.Transformation(identifier, 'user_example' if identifier == 'home' else 'home', 'home', operation, start, selected[-1].end, tuple(n.event_id for n in selected), tuple((n.start - start, n.end - start) for n in selected), (1, 1, 1)))
  anchor = generator.solo.types.Anchor('terminal', F(14), F(14), notes[-1].pitch, frozenset({2}), (55, 76), F(2), 'terminal', 'BbM7', hold_until=F(16))
  draft = generator.solo.types.MelodicDraft('illustration', F(), F(16), notes, tuple(groups), (notes[-1].event_id,), (), (anchor,), home, 'statement', (F(5), F(8)))
  return draft, choices

class DraftRepresentationTests(unittest.TestCase):
  def test_user_phrase_is_valid_and_contains_seventeen_attacks(self):
    draft, choices = illustration()
    harmony = generator.solo.harmonic.Harmony(choices)
    self.assertEqual(len(draft.notes), 17)
    self.assertEqual(draft.rests(), ((1, 2),))
    self.assertTrue(harmony.evaluate(draft.notes, complete=True).valid)
    self.assertTrue(generator.solo.drafts.contract_valid(draft, draft.notes)[0])
    self.assertTrue(generator.solo.harmonic.anchors_valid(draft.anchors, draft.notes, 16, complete=True, closure='closed')[0])
    self.assertTrue(any(t.start < 4 < t.end for t in draft.transformations))
  def test_partition_preserves_draft_relationships_and_harmonic_roles(self):
    draft, choices = illustration()
    whole = draft.notes
    split = []
    for n in whole:
      points = sorted({n.start, n.end} | {F(t) for t in range(1, 16) if n.start < t < n.end})
      split.extend(dataclasses.replace(n, start=a, end=b, continuation=i > 0) for i, (a, b) in enumerate(zip(points, points[1:])))
    self.assertEqual(generator.solo.types.sounding_notes(split), whole)
    self.assertEqual(generator.solo.drafts.relationship_features(draft, whole), generator.solo.drafts.relationship_features(draft, split))
    self.assertTrue(generator.solo.drafts.contract_valid(draft, split)[0])
    harmony = generator.solo.harmonic.Harmony(choices)
    self.assertEqual(harmony.describe(whole), harmony.describe(split))
  def test_disconnected_alternative_breaks_declared_rise_at_same_density(self):
    draft, _ = illustration()
    changed = list(draft.notes)
    changed[2] = dataclasses.replace(changed[2], pitch=note(0, 1, 'e').pitch)
    self.assertEqual(len(changed), len(draft.notes))
    self.assertEqual(sum(n.end - n.start for n in changed), sum(n.end - n.start for n in draft.notes))
    valid, reason = generator.solo.drafts.contract_valid(draft, changed)
    self.assertFalse(valid)
    self.assertIn('contour', reason)
    self.assertGreater(generator.solo.drafts.relationship_features(draft, changed)['direction'], 0)
  def test_catalog_keeps_cross_bar_sustain_as_one_event(self):
    event = generator.rhythm.Event
    bar = generator.rhythm.Bar
    snippet = generator.rhythm.Snippet('ties', '4/4', (bar((event(F(), F(3), False, 'half'), event(F(3), F(1), True, 'quarter')), F(4)), bar((event(F(), F(1), True, 'quarter', continuation=True), event(F(1), F(3), True, 'half')), F(4))))
    catalog = generator.solo.drafts.catalog(generator.rhythm.Library((snippet,)), '4/4')
    self.assertEqual(catalog[0][1], ((3, 5), (5, 8)))
  def test_extension_reopens_only_an_unprotected_provisional_tail(self):
    slot = generator.solo.drafts.Slot('tail', F(2), F(4), 0, 2)
    result, cursor = generator.solo.drafts.reopen((slot,), F(4))
    self.assertEqual((result[-1].end, cursor), (F(7, 2), F(7, 2)))
    for protected in (dataclasses.replace(slot, protected=True), dataclasses.replace(slot, forced=note(0, 1, 'e').pitch)):
      self.assertEqual(generator.solo.drafts.reopen((protected,), F(4)), ((protected,), F(4)))
  def test_existing_terminal_contract_is_not_retimed(self):
    draft, choices = illustration()
    plan = dataclasses.replace(tests.test_solo_development.plan(), end=F(16), target=generator.solo.types.TargetEvent(F(14), F(14), frozenset({2}), 'BbM7', 'explicit ending'))
    self.assertEqual(generator.solo.drafts.ending_options(draft.anchors, plan, choices), (draft.anchors,))
  def test_ending_alternatives_vary_timing_without_changing_destination(self):
    draft, choices = illustration()
    plan = dataclasses.replace(tests.test_solo_development.plan(), end=F(16))
    endings = generator.solo.drafts.ending_options(draft.anchors, plan, choices)
    self.assertEqual({a[-1].duration for a in endings}, {F(1, 2), F(1), F(2), F(3)})
    self.assertTrue(all(a[-1].pitch == draft.anchors[-1].pitch and a[-1].earliest + a[-1].duration == a[-1].hold_until for a in endings))
    self.assertEqual({16 - a[-1].hold_until for a in endings}, {F(), F(1, 2), F(1)})
  def test_revision_cannot_change_terminal_hold_or_identifying_rhythm(self):
    draft, _ = illustration()
    for index in (2, len(draft.notes) - 1):
      notes = list(draft.notes)
      notes[index] = dataclasses.replace(notes[index], end=notes[index].end - F(1, 4))
      self.assertFalse(generator.solo.drafts.contract_valid(draft, notes)[0])
  def test_unrelated_single_note_is_not_perfect_development(self):
    draft, _ = illustration()
    event = draft.notes[0]
    transform = generator.solo.types.Transformation('short', 'home', 'home', 'extend', event.start, event.end, (event.event_id,), ((F(), F(1)),))
    candidate = dataclasses.replace(draft, notes=(event,), transformations=(transform,), structural=(), anchors=(), joins=())
    self.assertEqual(generator.solo.drafts.relationship_features(candidate, candidate.notes)['development'], 1)
  def test_diminution_and_broadening_keep_exact_timing_and_event_order(self):
    pattern = ((F(), F(1)), (F(3, 2), F(2)), (F(2), F(5, 2)))
    transformed, length = generator.solo.drafts.develop_timing(pattern, F(3), F(1, 2), F(3, 4), True)
    self.assertEqual(len(transformed), len(pattern))
    self.assertTrue(all(b - a >= F(1, 2) for a, b in transformed))
    self.assertTrue(all(b <= c for (_, b), (c, _) in zip(transformed, transformed[1:])))
    self.assertLess(length, 3)
  def test_occupancy_goal_cannot_create_arbitrary_tuplet_durations(self):
    pattern = ((F(), F(1, 2)), (F(1), F(3, 2)), (F(2), F(5, 2)))
    for occupancy in (F(3, 4), F(9, 10)):
      for unit in (F(1, 2), F(1), F(2)):
        transformed, length = generator.solo.drafts.develop_timing(pattern, F(3), unit, occupancy, True)
        self.assertTrue(all((a * 4).denominator == (b * 4).denominator == 1 for a, b in transformed))
        self.assertEqual((length * 4).denominator, 1)
  def test_pending_treatment_constrains_the_new_draft_pitch_search(self):
    choices = tests.test_solo_passages.choices(('Gm7', 'C7'))
    harmony = generator.solo.harmonic.Harmony(choices)
    held = note(2, F(9, 2), 'f')
    treatment = generator.solo.types.HarmonicTreatment('pending', 'suspension', 65, F(4), F(9, 2), 64, F(5))
    incoming = generator.solo.types.ContinuationState('p', (held,), treatments=(treatment,), last=held)
    plan = dataclasses.replace(tests.test_solo_development.plan(), start=F(9, 2), end=F(6))
    slot = generator.solo.drafts.Slot('resolve', F(9, 2), F(5), 0, 0)
    group = generator.solo.drafts.Group('g', 'continue', F(9, 2), F(6), 'prepared treatment')
    rhythm = generator.solo.drafts.RhythmDraft('r', (slot,), (group,), (), plan.motif, 'continue')
    paths, error = generator.solo.realization.draft_pitch_paths(rhythm, plan, choices, harmony, tests.solo_fixtures.library(), incoming, generator.solo.realization.SearchContext())
    self.assertTrue(paths, error)
    self.assertTrue(all(notes[0].pitch.number() == 64 for notes, _, _ in paths))
    self.assertTrue(all(harmony.evaluate((held, *notes), contracts, complete=True).valid for notes, _, contracts in paths))
  def test_join_search_can_rewrite_both_sides_and_preserve_outside_events(self):
    pitches = (('c', 4), ('g', 4), ('e', 4), ('c', 5), ('c', 4), ('f', 4), ('g', 4), ('e', 4))
    notes = tuple(dataclasses.replace(note(i, i + 1, letter, octave=octave), event_id=f'n{i}') for i, (letter, octave) in enumerate(pitches))
    plan = tests.test_solo_development.plan()
    harmony = generator.solo.harmonic.Harmony(tests.test_solo_development.choices())
    transform = generator.solo.types.Transformation('connecting_run', '', '', 'close', F(), F(8), tuple(n.event_id for n in notes), tuple((n.start, n.end) for n in notes), preserves=())
    draft = generator.solo.types.MelodicDraft('join_fixture', F(), F(8), notes, (transform,), ('n0', 'n7'), (), (), None, 'continue', (F(4),))
    incoming = generator.solo.types.ContinuationState('')
    original = generator.solo.realization.checked_phrase(plan, draft, notes, (), incoming, harmony, F(120))
    self.assertIsNotNone(original)
    revised = generator.solo.realization.improve_window(original, incoming, harmony, tests.solo_fixtures.library(), F(120), F(2), F(6), 'join', generator.solo.realization.SearchContext())
    actual = generator.solo.realization.emitted_notes(revised)
    self.assertLess(revised.score, original.score)
    self.assertTrue(any(a.pitch != b.pitch and a.start < 4 for a, b in zip(notes, actual)))
    self.assertTrue(any(a.pitch != b.pitch and a.start >= 4 for a, b in zip(notes, actual)))
    self.assertEqual((actual[0].pitch, actual[-1].pitch), (notes[0].pitch, notes[-1].pitch))
    self.assertTrue(harmony.evaluate(actual, complete=True).valid)
  def test_passage_boundary_does_not_count_a_held_note_as_a_new_attack(self):
    choices = tests.test_solo_development.choices()
    harmony = generator.solo.harmonic.Harmony(choices)
    placement = generator.solo.placement.place((note(0, 4, 'g'),), F(120))[1][0]
    incoming = generator.solo.types.ContinuationState('old', (placement,), last=placement)
    held = dataclasses.replace(note(4, 8, 'g'), continuation=True, event_id='hold')
    plan = dataclasses.replace(tests.test_solo_development.plan(), start=F(4), passage_id='new', passages=(generator.solo.types.ActivitySpan(F(4), F(8), 'active', F(1), F(3, 4)),))
    transform = generator.solo.types.Transformation('hold', '', '', 'close', F(4), F(8), ('hold',), ((F(), F(4)),), preserves=())
    draft = generator.solo.types.MelodicDraft('hold', F(4), F(8), (held,), (transform,), (), (), (), None, 'continue')
    phrase = generator.solo.realization.checked_phrase(plan, draft, (held,), (), incoming, harmony, F(120))
    self.assertIsNotNone(phrase)
    features = generator.solo.scoring.features(plan, (dataclasses.replace(placement, end=F(8)),), choices)
    self.assertEqual(phrase.score, generator.solo.scoring.value(features))
  def test_new_passage_carries_pending_but_not_completed_suspensions(self):
    treatment = generator.solo.types.HarmonicTreatment('prepared', 'suspension', 65, F(4), F(9, 2), 64, F(5))
    held, resolution = note(2, F(9, 2), 'f'), note(F(9, 2), 5, 'e')
    plan = dataclasses.replace(tests.test_solo_development.plan(), start=F(5), passage_id='next')
    incoming = generator.solo.types.ContinuationState('old', (held,), treatments=(treatment,), last=held)
    self.assertEqual(generator.solo.realization.inherited_treatments(plan, incoming), (treatment,))
    resolved = dataclasses.replace(incoming, notes=(held, resolution), last=resolution)
    self.assertEqual(generator.solo.realization.inherited_treatments(plan, resolved), ())
    self.assertEqual(generator.solo.realization.inherited_treatments(dataclasses.replace(plan, passage_id='old'), resolved), (treatment,))

class DraftCompositionTests(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
    cls.solo = tests.test_solo.compose(tests.test_solo.song(count=8, chords=('Gm7', 'C7', 'FM7', 'BbM7', 'Em7b5', 'A7', 'Dm7', 'Dm7')))
  def test_production_candidates_contain_complete_playable_drafts(self):
    harmony = generator.solo.harmonic.Harmony(self.solo.choices)
    previous = None
    for phrase in self.solo.phrases:
      draft = phrase.draft
      self.assertIsNotNone(draft)
      self.assertEqual((draft.start, draft.end), (phrase.plan.start, phrase.plan.end))
      prefix = (dataclasses.replace(previous, position=None),) if previous else ()
      self.assertTrue(harmony.evaluate((*prefix, *draft.notes), phrase.state.treatments, until=draft.end).valid)
      self.assertTrue(generator.solo.placement.place(draft.notes, self.solo.tempo, previous))
      self.assertTrue(draft.transformations)
      previous = phrase.state.last
  def test_source_family_survives_phrase_variations(self):
    first, second = self.solo.phrases
    self.assertEqual(first.state.home.id, second.state.home.id)
    for phrase in self.solo.phrases:
      self.assertTrue(all(t.home == phrase.state.home.id for t in phrase.draft.transformations))
      self.assertTrue(phrase.state.home.chord_degrees)
      self.assertTrue(phrase.draft.source_links)
  def test_groups_cross_barlines_without_reference_timing_reset(self):
    # Exercise an explicitly spanning construction; do not require a changed
    # search objective to pick the same old winner merely to test timing.
    plan = self.solo.phrases[0].plan
    pattern = ((F(), F(1)), (F(1), F(3, 2)), (F(3, 2), F(2)), (F(2), F(3)))
    constructed = generator.solo.drafts.construct(plan, plan.motif, 'statement', ('spanning', pattern, F(3)), (), 0, self.solo.choices)
    self.assertIsNotNone(constructed)
    self.assertTrue(any(g.start < 4 < g.end for g in constructed.groups))
    self.assertTrue(any(s.start == F(4) and s.end == F(9, 2) for s in constructed.slots))
    for phrase in self.solo.phrases:
      by_id = {n.event_id: n for n in phrase.draft.notes}
      for transform in phrase.draft.transformations:
        self.assertEqual(tuple((by_id[e].start - transform.start, by_id[e].end - transform.start) for e in transform.events), transform.rhythm)
  def test_new_passage_candidates_explicitly_recall_or_contrast(self):
    prior = self.solo.phrases[0].state
    plan = dataclasses.replace(self.solo.phrases[0].plan, id='new', passage_id='new_passage')
    rhythms = generator.solo.drafts.rhythms(plan, prior, tests.test_solo.song(count=8)['staves'][0]['measures'], self.solo.choices, tests.rhythm_fixtures.library())
    self.assertTrue(any(r.reference.id == prior.home.id and r.relationship in ('continue', 'answer', 'recall') for r in rhythms))
    self.assertTrue(all(r.relationship != 'statement' for r in rhythms))
  def test_cell_boundaries_are_not_phrase_rhythm_inputs(self):
    prior = generator.solo.types.ContinuationState('')
    plan = self.solo.phrases[0].plan
    measures = tests.test_solo.song(count=8)['staves'][0]['measures']
    first = generator.solo.drafts.rhythms(plan, prior, measures, self.solo.choices, tests.rhythm_fixtures.library())
    second = generator.solo.drafts.rhythms(dataclasses.replace(plan, gestures=(generator.solo.types.Gesture('arbitrary_cell', F(), F(3), 'continue', 'contour_reference'),)), prior, measures, self.solo.choices, tests.rhythm_fixtures.library())
    self.assertEqual(first, second)
  def test_full_realization_revalidates_each_accepted_revision(self):
    generator.solo.composer.validate(self.solo)
    for phrase in self.solo.phrases:
      notes = generator.solo.realization.emitted_notes(phrase)
      self.assertTrue(generator.solo.drafts.contract_valid(phrase.draft, notes)[0])
      self.assertTrue(all(1 <= n.position.fret <= 12 for n in notes))
      for edit in phrase.edits:
        self.assertLessEqual(edit.end - edit.start, 8 if edit.kind in ('replace_idea', 'complete_route') else 4)
  def test_join_comparison_reuses_identical_upstream_draft_pool(self):
    plan = self.solo.phrases[0].plan
    measures = tests.test_solo.song(count=8)['staves'][0]['measures']
    context = generator.solo.realization.SearchContext(revise=False)
    arguments = (plan, measures, self.solo.choices, tests.solo_fixtures.library(), tests.rhythm_fixtures.library(), self.solo.tempo)
    generator.solo.realization.alternatives(*arguments, context=context)
    fingerprints = dict(context.diagnostics['draft_pool_fingerprints'])
    constructed = context.diagnostics['generated_candidates']
    context.revise = True
    generator.solo.realization.alternatives(*arguments, context=context)
    self.assertEqual(context.diagnostics['generated_candidates'], constructed)
    self.assertEqual(context.diagnostics['draft_pool_fingerprints'], fingerprints)
    self.assertEqual(context.diagnostics['reused_draft_pools'], 1)