"""Complete paired lead ideas, realized by the existing phrase/output pipeline.

The bounded search branches on a whole timed hook or answer. Every pitch in
that group is measured from one fixed origin, never from a previously repaired
note. The original identity survives harmonic adaptation and passage changes.
"""
from __future__ import annotations

import dataclasses
import fractions
import itertools
import time

import generator.harmony
import generator.model
import generator.solo.drafts as drafts
import generator.solo.harmonic as harmonic
import generator.solo.policy as policy
import generator.solo.types as types

F = fractions.Fraction

def identities(references, incoming):
  reference = references[0]
  low, high = reference.features[0][1:]
  source = tuple(n for n in reference.notes if low <= n.start < high)
  rhythm = tuple(n.end - n.start for n in source)
  if len(rhythm) != 4:
    raise ValueError('lead reference requires a complete four-attack identifying figure')
  rhythms = (rhythm, (rhythm[1], rhythm[2], rhythm[0], rhythm[3]), (F(1), F(1, 2), F(3, 2), F(1)))
  shapes = ((0, 3, 2, 0), (0, -3, -2, 0), (0, 2, 1, -1), (0, 0, 2, 1))
  result = [types.LeadIdentity(f'{reference.id}:lead:{i}:{j}', durations, degrees, reference.id)
    for i, durations in enumerate(rhythms) for j, degrees in enumerate(shapes)]
  if incoming.lead_identity:
    result = [incoming.lead_identity, *(x for x in result if x != incoming.lead_identity)]
  return tuple(result[:policy.DEFAULT.lead_identities])

def timings(start, length, durations, finish, harmony):
  """Fit meter on an eighth-note grid; never compress into tiny tuplets."""
  cursor, result = start, []
  ideal = [duration * length / sum(durations) * 2 for duration in durations]
  ticks = [max(1, int(value)) for value in ideal]
  available = int(length * 2)
  if available < len(ticks) or length * 2 != available:
    return ()
  while sum(ticks) < available:
    index = max(range(len(ticks)), key=lambda i: (ideal[i] - ticks[i], -i))
    ticks[index] += 1
  while sum(ticks) > available:
    index = max((i for i in range(len(ticks)) if ticks[i] > 1), key=lambda i: (ticks[i] - ideal[i], -i))
    ticks[index] -= 1
  for duration in ticks:
    end = min(finish, cursor + F(duration, 2))
    if cursor < end:
      # An authored rest in the harmony is not filled with a melodic fallback.
      spans = [s for s in harmony.spans if s.start < end and s.end > cursor]
      if spans and all(s.scale for s in spans):
        result.append((cursor, end))
    cursor = end
  return tuple(result)

def group_options(spans, degrees, register, harmony, previous, identifier, *, home=(), terminal=False):
  """Enumerate whole paired figures, including non-nearest melodic intervals."""
  from generator.solo.realization import options
  from generator.solo.planning import target_pitches
  if not spans:
    return ((0., ()),)
  cache = getattr(harmony, 'lead_groups', None)
  if cache is None:
    harmony.lead_groups = cache = {}
  key = (spans, degrees, register, previous.pitch if previous else None, identifier, tuple(n.pitch for n in home), terminal)
  if key in cache:
    return cache[key]
  domains = []
  for index, (start, end) in enumerate(spans):
    available = options(harmony.choices, start, end, register)
    if end - start >= policy.DEFAULT.long_note:
      members = harmony.members(start, end)
      available = tuple(p for p in available if p.number() % 12 in members)
    if terminal and index == len(spans) - 1:
      chord = next(c.span.chord for c in harmony.choices if c.span.start <= start < c.span.end)
      eligible = target_pitches(chord, True) & harmony.members(start, end)
      available = tuple(p for p in available if p.number() % 12 in eligible)
    if not available:
      return ()
    domains.append(available)
  collection = harmony.at(spans[0][0]).scale
  result = []
  for origin in domains[0]:
    index = types.scale_index(collection, origin.number())
    desired = tuple(origin.number() + n.pitch.number() - home[0].pitch.number() for n in home) if home else tuple(types.scale_pitch(collection, index + degree).number() for degree in degrees)
    desired = desired[:len(spans)]
    if len(desired) < len(spans):
      continue
    narrowed = [(origin,)] + [tuple(sorted(domain, key=lambda p: (abs(p.number() - target), p))[:2]) for domain, target in zip(domains[1:], desired[1:])]
    for pitches in itertools.product(*narrowed):
      actual = tuple(b.number() - a.number() for a, b in zip(pitches, pitches[1:]))
      intended = tuple(b - a for a, b in zip(desired, desired[1:]))
      # A return must supply evidence of the complete heard figure, not just
      # its first interval. Harmonic adaptation may alter interval quality.
      if home and any((a > 0) - (a < 0) != (b > 0) - (b < 0) or abs(a - b) > 2 for a, b in zip(actual, intended)):
        continue
      shape = sum(abs(a - b) / 12 + .5 * (((a > 0) - (a < 0)) != ((b > 0) - (b < 0))) for a, b in zip(actual, intended)) / max(1, len(actual))
      if not home and shape > .45:
        continue
      notes = tuple(types.Note(start, end, pitch, (identifier,), 'motif', event_id=f'{identifier}:{start}') for (start, end), pitch in zip(spans, pitches))
      stopped = sum(end - start >= 1 and p.number() % 12 not in harmony.at(start).roles.basic for (start, end), p in zip(spans, pitches)) / len(notes)
      connection = min(1, abs(origin.number() - previous.pitch.number()) / 12) if previous else 0
      result.append((shape + .3 * stopped + .15 * connection, notes))
  result.sort(key=lambda v: (v[0], tuple(n.pitch for n in v[1])))
  selected = tuple(result[:policy.DEFAULT.lead_realizations])
  cache[key] = selected
  return selected

def make_phrase(plan, notes, groups, identity, incoming, harmony, tempo, fit):
  from generator.solo import realization
  if not notes:
    return None
  first = next((g for g in groups if g[2] in ('statement', 'recall', 'answer_theme')), groups[0])
  home_notes = tuple(n for n in notes if first[0] <= n.start < first[1])
  home = realization.make_identity(first[3], identity.id, home_notes, first[0], first[1], first[2], harmony)
  transformations, features, pairs = [], [], []
  parent = incoming.development.id if incoming.development else identity.id
  for start, end, operation, group_id in groups:
    actual = tuple(n for n in notes if start <= n.start < end)
    if not actual:
      continue
    directions = tuple((b.pitch.number() > a.pitch.number()) - (b.pitch.number() < a.pitch.number()) for a, b in zip(actual, actual[1:]))
    transformations.append(types.Transformation(group_id, parent, home.id, operation, start, end, tuple(n.event_id for n in actual), tuple((n.start - start, n.end - start) for n in actual), directions, (), identity.source))
    parent = group_id
    if operation in ('statement', 'recall', 'answer_theme', 'repeat_with_harmonic_adaptation') and len(actual) == len(identity.degrees):
      features.append(types.IdentifyingFeature(group_id + ':hook', tuple(n.event_id for n in actual), tuple((n.start - actual[0].start, n.end - n.start) for n in actual), tuple(b.pitch.number() - a.pitch.number() for a, b in zip(actual, actual[1:])), identity.id))
      if operation == 'repeat_with_harmonic_adaptation':
        pairs.extend((a.event_id, b.event_id) for a, b in zip(home_notes, actual))
  anchors = []
  for note, role in ((notes[0], 'opening'), *((n, 'intermediate') for n in notes[1:-1] if any(n.start == g[0] and g[2] == 'repeat_with_harmonic_adaptation' for g in groups)), (notes[-1], 'terminal')):
    span = harmony.at(note.start)
    if note.pitch.number() % 12 not in span.roles.basic:
      if role != 'terminal':
        continue
      return None
    anchors.append(types.Anchor(note.event_id + ':anchor', note.start, note.start, note.pitch, span.roles.basic, plan.register, note.end - note.start, role, span.chord, hold_until=note.end))
  last = notes[-1]
  passages = (types.ActivitySpan(plan.start, last.end, 'active', plan.activity, F('0.925')),)
  if last.end < plan.end:
    passages += (types.ActivitySpan(last.end, plan.end, 'space', F(), F()),)
  motif = types.Motif(home.id, home.notes, home.duration, tuple(s.chord for s in harmony.spans if s.start < first[1] and s.end > first[0]), 'invented paired lead: ' + identity.id)
  actual_plan = dataclasses.replace(plan, motif=motif, passages=passages, release=plan.end - last.end, landmarks=(), anchor_options=(tuple(anchors),), ending='breath' if last.end < plan.end else 'sustain', closure='closed')
  draft = types.MelodicDraft(plan.id + ':' + identity.id + ':' + policy.stable_rank(plan.seed, tuple(n.pitch for n in notes))[:10], plan.start, plan.end, notes, tuple(transformations), tuple(a.id.removesuffix(':anchor') for a in anchors), (), tuple(anchors), home, first[2], tuple(g[0] for g in groups[1:]), tuple(features), identity.source, motif, identity, tuple(pairs))
  phrase = realization.checked_phrase(actual_plan, draft, notes, (), incoming, harmony, tempo)
  if phrase:
    state = dataclasses.replace(phrase.state, lead_identity=identity)
    # Different complete musical constructions compete before any ornament.
    return dataclasses.replace(phrase, state=state, score=phrase.score + fit, initial=draft, refined=draft)

def candidates(plan, measures, choices, incoming, tempo, context):
  from generator.solo import realization, sequences
  harmony_key = ('lead_harmony', policy.stable_rank('harmony', choices))
  if harmony_key not in context.resources:
    context.resources[harmony_key] = harmonic.Harmony(choices)
  harmony = context.resources[harmony_key]
  active = [s for s in harmony.spans if s.scale and s.start < plan.end and s.end > plan.start]
  if not active:
    silent_plan = dataclasses.replace(plan, passages=(types.ActivitySpan(plan.start, plan.end, 'space', F(), F()),), release=plan.end-plan.start, landmarks=(), anchor_options=())
    draft = types.MelodicDraft(plan.id + ':space', plan.start, plan.end, (), (), (), (), (), None, 'space')
    phrase = realization.checked_phrase(silent_plan, draft, (), (), incoming, harmony, tempo)
    return (dataclasses.replace(phrase, initial=draft, refined=draft),) if phrase else ()
  finish = min(plan.end, active[-1].end)
  boundaries, cursor = [], F()
  for measure in measures:
    length = generator.model.fraction(measure['duration'])
    if plan.start <= cursor < finish:
      boundaries.append((cursor, min(cursor + length, finish)))
    cursor += length
  result = []
  for identity in identities(context.references, incoming):
    states = [(0., (), (), ())]
    for index, (start, end) in enumerate(boundaries):
      final = index == len(boundaries) - 1
      hook = index in (0, 2) and not final
      returning = hook and index == 2
      operation = ('recall' if incoming.lead_identity == identity else 'answer_theme' if incoming.lead_identity else 'statement') if index == 0 else 'repeat_with_harmonic_adaptation' if returning else 'close' if final else 'extend' if index >= 4 else 'answer'
      # Paired answers retain the hook's direction of departure. Ending
      # alternatives supply timing and pitches together, not inserted holds.
      specs = [(identity.durations, identity.degrees)] if hook else [((F(1), F(1, 2), F(1, 2), F(1), F(1)), (0, 1, 2, 1, -1)), ((F(2), F(1, 2), F(1, 2), F(1)), (0, 1, 2, 1))]
      if index >= 4 and not final:
        reference = context.references[0]
        profiles = [d for d in reference.development if d[1] == index * 4 and d[3]]
        if profiles:
          _, low, high, degrees = profiles[0]
          source = tuple(n for n in reference.notes if low <= n.start < high)
          specs = [(tuple(n.end - n.start for n in source), degrees)]
      if final:
        specs = [((F(1), F(1, 2), F(1, 2), F(2)), (0, -1, -2, -1)), ((F(1, 2), F(1, 2), F(1), F(1)), (0, 2, 1, 0))]
        if end - start < 2:
          specs = [((end - start,), (0,))]
        elif end - start < 4:
          specs = [((F(1, 2), end - start - F(1, 2)), (0, -1)), ((F(1), end - start - 1), (0, 1))]
      following = []
      for cost, notes, groups, home in states:
        for variant, (durations, degrees) in enumerate(specs):
          release = F(1) if final and variant == 1 and end - start >= 4 else F()
          spans = timings(start, end - start - release, durations, end - release, harmony)
          identifier = f'{plan.id}:{identity.id}:{index}:{variant}'
          previous = notes[-1] if notes else incoming.last
          heard = home if returning else incoming.home.notes if index == 0 and incoming.home and incoming.lead_identity == identity else ()
          for local_cost, actual in group_options(spans, degrees, plan.register, harmony, previous, identifier, home=heard, terminal=final):
            if not actual:
              following.append((cost, notes, groups, home))
              continue
            following.append((cost + local_cost, (*notes, *actual), (*groups, (start, end, operation, identifier)), actual if index == 0 else home))
      states = sorted(following, key=lambda v: (v[0], tuple(n.pitch for n in v[1])))[:policy.DEFAULT.lead_realizations]
      context.diagnostics['evaluated_extensions'] += len(following)
      context.diagnostics['peak_states'] = max(context.diagnostics['peak_states'], len(states))
      if not states:
        context.diagnostics.setdefault('lead_failures', []).append({'phrase': plan.id, 'identity': identity.id, 'start': str(start), 'end': str(end), 'operation': operation, 'reason': 'no complete figure satisfies the harmonic, duration and return constraints within the bounded candidate pool'})
        break
    for cost, notes, groups, home in states:
      if not notes:
        continue
      phrase = make_phrase(plan, notes, groups, identity, incoming, harmony, tempo, cost / max(1, len(groups)))
      if phrase:
        # A heard identity is a real option across passages, not obligatory
        # perpetual development of whichever fragment happened most recently.
        if incoming.lead_identity and identity != incoming.lead_identity:
          phrase = dataclasses.replace(phrase, score=phrase.score + .15)
        result.append(phrase)
  context.diagnostics['generated_candidates'] += len(result)
  return sequences.diverse(result, lambda p: (p.score, policy.stable_rank(plan.seed, p.draft.id)), lambda p: p.draft.lead_identity.id, policy.DEFAULT.phrase_alternatives)

def approach_pitch(target, number):
  """Spell a semitone neighbor using the adjacent diatonic letter."""
  letters = generator.model.LETTERS
  step = 1 if number > target.number() else -1
  index = letters.index(target.letter) + step
  octave, degree = divmod(index, 7)
  candidates = []
  for letter, octave in ((letters[degree], target.octave + octave), (target.letter, target.octave)):
    natural = generator.model.NATURAL[letter] + 12 * (octave + 1)
    candidates.append(types.Pitch(letter, F(number - natural), octave))
  # A natural above A-flat is clearer than B-double-flat; an ordinary
  # leading tone still uses its adjacent letter (D-sharp into E).
  return min(candidates, key=lambda p: abs(p.alter))

def chromatic(phrase, incoming, harmony, tempo):
  from generator.solo import realization
  notes = realization.emitted_notes(phrase)
  protected = {*phrase.draft.protected, *phrase.draft.structural, *(e for f in phrase.draft.features for e in f.events)}
  budget = min(policy.DEFAULT.chromatic_max_per_eight_bars, max(1, int((phrase.plan.end - phrase.plan.start) / 8)))
  proposals = []
  for index, (note, target) in enumerate(zip(notes, notes[1:])):
    if note.event_id in protected or note.end - note.start != F(1, 2) or note.start % 1 != F(1, 2) or note.end != target.start or len(harmony.segments(note)) != 1:
      continue
    span, destination = harmony.at(note.start), harmony.at(target.start)
    if not span or not span.scale or not destination or target.pitch.number() % 12 not in destination.roles.basic:
      continue
    for number in (target.pitch.number() - 1, target.pitch.number() + 1):
      if number % 12 in span.scale.pitches or not phrase.plan.register[0] <= number <= phrase.plan.register[1]:
        continue
      pitch = approach_pitch(target.pitch, number)
      treatment = types.HarmonicTreatment(note.event_id + ':chromatic', 'chromatic_approach', number, note.start, note.end, target.pitch.number(), target.start, F(), note.event_id, target.event_id)
      smooth = abs(number - (notes[index - 1].pitch.number() if index else note.pitch.number()))
      proposals.append((smooth, abs(number - note.pitch.number()), note.start, index, pitch, treatment))
  chosen_times, used_routes = [], set()
  previous_times = [t.start for t in incoming.treatments if t.kind == 'chromatic_approach']
  original_score = phrase.score
  for _, _, _, index, pitch, treatment in sorted(proposals, key=lambda p: p[:4]):
    if len(chosen_times) >= budget or any(abs(treatment.start - t) < 8 for t in (*previous_times, *chosen_times)):
      continue
    route_ids = {r.id for r in phrase.routes if any(n.event_id == treatment.event for n in r.notes)}
    if route_ids & used_routes:
      continue
    original = realization.emitted_notes(phrase)
    changed = tuple(dataclasses.replace(n, pitch=pitch, position=None, role='chromatic_approach') if n.event_id == treatment.event else n for n in original)
    contracts = (*phrase.state.treatments, treatment)
    if not harmony.evaluate(changed, contracts, complete=True).valid:
      continue
    draft = realization.rebase_draft(phrase.draft, changed, harmony)
    edit = types.DraftEdit('chromatic_approach', treatment.start, treatment.deadline, (treatment.event,), (treatment.event,), 'semitone resolution to ' + treatment.resolution_event)
    checked = realization.checked_phrase(phrase.plan, draft, changed, contracts, incoming, harmony, tempo, (*phrase.edits, edit))
    if checked:
      routes = tuple(dataclasses.replace(r, notes=tuple(dataclasses.replace(n, pitch=pitch, position=None, role='chromatic_approach') if n.event_id == treatment.event else n for n in r.notes), relationship=r.relationship + ('; chromatic approach' if r.id in route_ids else '')) for r in phrase.routes)
      phrase = dataclasses.replace(checked, score=original_score, initial=phrase.initial, refined=phrase.refined, routes=routes, state=dataclasses.replace(checked.state, lead_identity=phrase.state.lead_identity))
      chosen_times.append(treatment.start)
      used_routes.update(route_ids)
  return phrase

def alternatives(plan, measures, choices, sequences, tempo, previous, context):
  from generator.solo import realization
  incoming = previous if isinstance(previous, types.ContinuationState) else types.ContinuationState('', last=previous)
  context.resources[id(sequences)] = sequences
  key = ('paired-lead', plan, incoming, policy.stable_rank('harmony', choices), tuple(generator.model.fraction(m['duration']) for m in measures), tempo, id(sequences), tuple(r.version for r in context.references), policy.describe()['fingerprint'], context.stage)
  if key in context.results:
    context.diagnostics['cache_hits'] += 1
    return context.results[key]
  started = time.monotonic()
  context.diagnostics['searches'] += 1
  complete = candidates(plan, measures, choices, incoming, tempo, context)
  if not complete:
    raise realization.SearchExhausted(f'{plan.id}: no complete paired lead reaches the terminal harmony within {policy.DEFAULT.lead_identities} identities / {policy.DEFAULT.lead_realizations} realizations; not proof of incompatibility')
  harmony = harmonic.Harmony(choices)
  result = []
  for phrase in complete:
    score, identity = phrase.score, phrase.state.lead_identity
    if context.stage != 'initial':
      # The complete plain tune is selected by the same score in all stages.
      # Ornament quality cannot rescue a weak underlying idea.
      phrase = realization.elaborate(phrase, incoming, harmony, sequences, tempo, context)
      diatonic = realization.rebase_draft(phrase.draft, realization.emitted_notes(phrase), harmony)
      phrase = dataclasses.replace(phrase, refined=diatonic, score=score, state=dataclasses.replace(phrase.state, lead_identity=identity))
      if context.stage == 'elaborated':
        phrase = chromatic(phrase, incoming, harmony, tempo)
    result.append(dataclasses.replace(phrase, score=score))
  context.diagnostics['construction_seconds'] = context.diagnostics.get('construction_seconds', 0) + time.monotonic() - started
  context.diagnostics.setdefault('complete_ideas', {})[policy.stable_rank('incoming', (plan.id, incoming))[:16]] = [{'id': p.draft.id, 'identity': p.draft.lead_identity.id if p.draft.lead_identity else 'space', 'score': p.score, 'attacks': len(p.initial.notes), 'return_pairs': len(p.draft.correspondence), 'edits': [e.kind for e in p.edits]} for p in result]
  context.results[key] = tuple(result)
  return tuple(result)