"""Phrase-wide rhythm construction and contracts for concrete melodic drafts.

No bar winners are assembled here. A rhythm is stated, transformed in place,
and timed to its destinations before the pitch search begins.
"""
from __future__ import annotations

import dataclasses
import fractions

import generator.model
import generator.solo.harmonic
import generator.solo.policy
import generator.solo.scoring
import generator.solo.sequences
import generator.solo.types
import generator.solo.references

F = fractions.Fraction

@dataclasses.dataclass(frozen=True)
class Slot:
  id: str
  start: F
  end: F
  group: int
  reference: int
  forced: generator.solo.types.Pitch | None = None
  protected: bool = False

@dataclasses.dataclass(frozen=True)
class Group:
  id: str
  operation: str
  start: F
  end: F
  source: str

@dataclasses.dataclass(frozen=True)
class RhythmDraft:
  id: str
  slots: tuple[Slot, ...]
  groups: tuple[Group, ...]
  anchors: tuple[generator.solo.types.Anchor, ...]
  reference: generator.solo.types.Motif | generator.solo.types.RealizedMotif
  relationship: str
  octave: int = 0

def catalog(library, meter):
  """Read complete snippets; ties across authored bars stay one attack."""
  result = {}
  for snippet in library.for_meter(meter):
    offset, spans = F(), []
    for bar in snippet.bars:
      for event in bar.events:
        if not event.sounding:
          continue
        start, end = offset + event.onset, offset + event.onset + event.duration
        if event.continuation:
          if spans and spans[-1][1] == start:
            spans[-1] = (spans[-1][0], end)
        else:
          spans.append((start, end))
      offset += bar.length
    if 2 <= len(spans) <= 12:
      result.setdefault(tuple(spans), (snippet.id, offset))
  return tuple((identifier, spans, length) for spans, (identifier, length) in result.items())

def templates(plan, reference, library, meter):
  source = tuple((n.start, n.end) for n in reference.notes)
  if plan.treatment in ('quote', 'exact_quote', 'rhythmic_reference'):
    return tuple([(reference.id, source, reference.duration)])
  # A compact identifying portion, including its authored internal rests.
  source = source[:6]
  result = []
  if source:
    length = min(reference.duration, source[-1][1] + F(1, 2))
    result.append((reference.id, source, length))
  choices = catalog(library, meter)
  ranked = generator.solo.sequences.diverse(choices,
    lambda item: (abs(len(item[1]) / float(item[2]) - float(plan.activity)) > 0.75, generator.solo.policy.stable_rank(plan.seed, item)),
    lambda item: (len(item[1]), bool(item[1][0][0]), item[2]), 4)
  result.extend(ranked)
  return tuple(result)

def reopen(slots, cursor):
  """An extension changes the previous provisional ending, not its core."""
  if not slots:
    return slots, cursor
  last = slots[-1]
  if last.protected or last.forced:
    return slots, cursor
  end = max(last.start + F(1, 2), last.end - F(1, 2))
  return (*slots[:-1], dataclasses.replace(last, end=end)), end

def develop_timing(pattern, period, unit, occupancy, connect):
  """Diminish/broaden one established rhythm, retaining order and rests.

  Ordinary diminution stops at eighth notes. Connected development may
  broaden existing sounds into part of their following rest, not insert a new
  independent pattern. These choices precede pitch/harmonic realization.
  """
  result, cursor, previous = [], F(), F()
  for start, end in pattern:
    cursor += (start - previous) * unit
    duration = max(F(1, 2), (end - start) * unit) if unit < 1 else (end - start) * unit
    result.append((cursor, cursor + duration))
    cursor += duration
    previous = end
  length = cursor + (period - previous) * unit
  if connect:
    needed = max(F(), occupancy * length - sum(b - a for a, b in result))
    for i, (a, b) in enumerate(result):
      following = result[i + 1][0] if i + 1 < len(result) else length
      # Occupancy is a soft objective, not a duration denominator. Use a half
      # beat or the existing rest endpoint; never manufacture e.g. 7/40 beats
      # solely to hit a percentage exactly.
      extra = min(F(1, 2), following - b)
      if abs(needed - extra) >= needed:
        extra = F()
      result[i] = (a, b + extra)
      needed -= extra
  return tuple(result), length

def ending_options(anchors, plan, choices):
  if not anchors or plan.target or plan.treatment in ('quote', 'exact_quote', 'rhythmic_reference'):
    return tuple([anchors])
  terminal = next((a for a in anchors if a.role == 'terminal'), None)
  if terminal is None or terminal.hold_until is None:
    return tuple([anchors])
  passage_end = plan.passage_end or plan.end
  active = [c for c in choices if c.scale and c.span.start < passage_end and c.span.end > plan.start]
  if not active:
    return tuple([anchors])
  # Missing harmony is mandatory space. Retiming an ordinary ending must not
  # move its destination into a trailing N.C. span.
  harmonic_end = min(passage_end, active[-1].span.end)
  result = []
  for release in (F(), F(1, 2), F(1)):
    finish = harmonic_end - release
    last = next((c for c in choices if c.span.start < finish <= c.span.end), None)
    for duration in (F(1, 2), F(1), F(2), F(3)):
      start = finish - duration
      if last is None or start < max(last.span.start, plan.passage_start or plan.start) or any(a.role != 'terminal' and a.latest + a.duration > start for a in anchors):
        continue
      changed = dataclasses.replace(terminal, earliest=start, latest=start, duration=duration, hold_until=finish)
      result.append(tuple(changed if a.id == terminal.id else a for a in anchors))
  return tuple(result) or tuple([anchors])

def ending_plan(plan, anchors):
  if plan.target or plan.treatment in ('quote', 'exact_quote', 'rhythmic_reference'):
    return plan
  terminal = next((a for a in anchors if a.role == 'terminal'), None)
  if terminal is None or terminal.hold_until is None:
    return plan
  finish = plan.passage_end or plan.end
  release = finish - terminal.hold_until
  # The candidate chooses where arrival begins. Earlier portions remain active.
  spans = []
  for span in plan.passages:
    low, high = span.start, min(span.end, terminal.earliest)
    if low < high:
      spans.append(dataclasses.replace(span, end=high, intent='active', occupancy=F(str(generator.solo.policy.DEFAULT.active_occupancy))))
  if not spans and plan.start < terminal.earliest:
    spans.append(generator.solo.types.ActivitySpan(plan.passage_start if plan.passage_start is not None else plan.start, terminal.earliest, 'active', plan.activity, F(str(generator.solo.policy.DEFAULT.active_occupancy))))
  arrival_length = terminal.hold_until - terminal.earliest
  desired = max(F(), plan.activity * (finish - (plan.passage_start if plan.passage_start is not None else plan.start)) - 1)
  expected = sum(s.activity * (s.end - s.start) for s in spans)
  spans = [dataclasses.replace(s, activity=min(F(2), s.activity * desired / expected) if expected else plan.activity) for s in spans]
  spans.append(generator.solo.types.ActivitySpan(terminal.earliest, terminal.hold_until, 'sustained_arrival', 1 / arrival_length, F(str(generator.solo.policy.DEFAULT.arrival_occupancy))))
  if release:
    spans.append(generator.solo.types.ActivitySpan(terminal.hold_until, finish, 'space', F(), F()))
  complete = plan.end == finish
  return dataclasses.replace(plan, release=release if complete else F(), passages=tuple(spans), ending=('breath' if release else 'sustain') if complete else 'connected')

def align_anchors(slots, anchors, plan):
  """Choose arrival timing with the rhythm, before assigning any free pitches."""
  slots = list(slots)
  exact = plan.treatment in ('quote', 'exact_quote', 'rhythmic_reference')
  protected_end = plan.start + plan.motif.duration if exact else plan.start
  for anchor in anchors:
    if not plan.start <= anchor.earliest < plan.end - plan.release:
      continue
    possible = [s for s in slots if anchor.earliest <= s.start <= anchor.latest]
    fit = [s for s in possible if s.end - s.start >= anchor.duration]
    selected = (fit or possible)
    slot = selected[0] if selected else None
    time = slot.start if slot else anchor.earliest
    end = anchor.hold_until or time + anchor.duration
    end = min(end, plan.end - plan.release)
    if time < protected_end:
      if slot is None or slot.end < end or (slot.forced and slot.forced != anchor.pitch):
        return None
      slots[slots.index(slot)] = dataclasses.replace(slot, forced=anchor.pitch)
      continue
    group = slot.group if slot else next((s.group for s in reversed(slots) if s.start <= time), 0)
    event_id = slot.id if slot else anchor.id
    replacement = Slot(event_id, time, end, group, slot.reference if slot else 0, anchor.pitch)
    before = [dataclasses.replace(s, end=min(s.end, time)) for s in slots if s.start < time]
    slots = [s for s in before if s.end > s.start] + [replacement] + [s for s in slots if s.start >= end and s.start != time]
  return tuple(slots)

def construct(plan, reference, relationship, template, anchors, variant, choices, octave=0):
  identifier, pattern, period = template
  if not pattern or period <= 0:
    return None
  finish = plan.end - plan.release
  slots, groups, cursor = (), [], plan.start
  protected = plan.treatment in ('quote', 'exact_quote', 'rhythmic_reference')
  while cursor < finish:
    index = len(groups)
    # Establish the idea before departing. Choose from the remaining musical
    # span instead of cycling a fixed list of transformations indefinitely.
    remaining = finish - cursor
    operation = relationship if index == 0 else 'repeat' if index == 1 and variant == 0 else 'answer' if remaining <= period else 'extend' if variant == 1 and index == 1 else 'recall'
    activity = next((p for p in plan.passages if p.start <= cursor < p.end), None)
    rate = F(len(pattern), 1) / period
    unit = F(2) if activity and activity.activity < F(3, 4) else F(1, 2) if activity and activity.activity > rate + F(1, 4) and index else F(1)
    local = pattern
    length = period
    if index and operation == 'extend':
      slots, cursor = reopen(slots, cursor)
      # Continue a recognizable opening beyond the old cadence; this group
      # can span either side of a barline without a new catalog selection.
      tail = pattern[:2]
      local = (*pattern, *((period + a, period + b) for a, b in tail))
      length += tail[-1][1]
    elif index and operation == 'answer':
      local = pattern[:max(2, len(pattern) - 1)]
      length = local[-1][1]
    elif index and operation == 'recall':
      operation = 'repeat_with_harmonic_adaptation'
    elif index and operation == 'displace':
      local = tuple((a + F(1, 2), b + F(1, 2)) for a, b in pattern)
      length += F(1, 2)
    if index and unit == 2:
      operation = 'fragment'
      local = local[:3]
      length = local[-1][1]
    if index and not protected:
      connected = operation not in ('repeat', 'repeat_with_harmonic_adaptation')
      local, length = develop_timing(local, length, unit, activity.occupancy if activity else F(3, 4), connected)
      if unit < 1:
        operation = 'diminish' if operation != 'extend' else 'extend_diminished'
      elif unit > 1:
        operation = 'broaden_ending'
      unit = F(1)
    if activity and activity.intent == 'space':
      cursor = min(activity.end, finish)
      continue
    if cursor + unit * length > finish:
      operation = 'close' if index else operation
    group_end = min(finish, cursor + unit * length)
    if group_end <= cursor:
      break
    group = Group(f'{plan.id}_development_{index + 1}', operation, cursor, group_end, identifier)
    groups.append(group)
    new = []
    for j, (a, b) in enumerate(local):
      start, end = cursor + a * unit, min(group_end, cursor + b * unit)
      if start >= end:
        continue
      quote = protected and index == 0
      pitch = None
      if quote and plan.treatment in ('quote', 'exact_quote'):
        pitch = dataclasses.replace(reference.notes[j].pitch, octave=reference.notes[j].pitch.octave + octave)
      new.append(Slot(f'{group.id}_event_{j + 1}', start, end, index, j % len(pattern), pitch, quote))
    slots = (*slots, *new)
    cursor = group_end
  if protected and plan.start + reference.duration > finish:
    return None
  aligned = align_anchors(slots, anchors, plan)
  if aligned is None:
    return None
  # Missing harmony is explicit space. Only ordinary material may fragment.
  filtered = []
  for slot in aligned:
    pieces = [(slot.start, slot.end)]
    for choice in choices:
      if choice.scale is None:
        pieces = [(a, b) for low, high in pieces for a, b in ((low, min(high, choice.span.start)), (max(low, choice.span.end), high)) if b > a]
    if slot.protected and pieces != [(slot.start, slot.end)]:
      return None
    filtered.extend(dataclasses.replace(slot, id=slot.id + (f'_fragment_{i}' if len(pieces) > 1 else ''), start=a, end=b) for i, (a, b) in enumerate(pieces))
  # The finalized rhythm is the contract, including changes made by extension
  # and anchor alignment. It is never rebased when a cell or bar ends.
  identity = generator.solo.policy.stable_rank(plan.seed, (identifier, variant, anchors, relationship, octave))[:12]
  return RhythmDraft(identity, tuple(filtered), tuple(groups), anchors, reference, relationship, octave)

def rhythms(plan, incoming, measures, choices, library, batch=0, reference_library=()):
  continuing = bool(plan.passage_id and incoming.passage_id == plan.passage_id)
  anchors = (incoming.selected_anchors,) if continuing and incoming.selected_anchors else plan.anchor_options or tuple([()])
  references = [(plan.motif, 'statement')]
  if incoming.home and plan.treatment not in ('quote', 'exact_quote', 'rhythmic_reference'):
    references = [(incoming.home, 'continue' if incoming.ending == 'connected' or continuing else 'answer')]
    if not continuing:
      references.append((incoming.home, 'recall'))
      if plan.motif.id != incoming.home.parent:
        references.append((plan.motif, 'contrast'))
      if incoming.themes and incoming.themes[0].id != incoming.home.id:
        references.append((incoming.themes[0], 'recall'))
  if plan.treatment not in ('quote', 'exact_quote', 'rhythmic_reference') and not continuing:
    references.extend((motif, 'contrast' if incoming.home else 'statement') for example in reference_library for motif in example.windows())
  elapsed = F()
  meter = measures[0]['time_signature']
  for measure in measures:
    end = elapsed + generator.model.fraction(measure['duration'])
    if elapsed <= plan.start < end:
      meter = measure['time_signature']
      break
    elapsed = end
  result = []
  for reference, relationship in references:
    available_templates = templates(plan, reference, library, meter)
    if isinstance(reference, generator.solo.types.RealizedMotif) or reference.origin == 'accepted reference profile':
      available_templates = tuple([(reference.id, tuple((n.start, n.end) for n in reference.notes), reference.duration)])
    for template in available_templates:
      for backbone in anchors:
        endings = (backbone,) if continuing else ending_options(backbone, plan, choices)
        for ending in endings:
          for variant in range(2):
            for octave in ((0, 1, -1, 2, -2) if plan.treatment in ('quote', 'exact_quote') else tuple([0])):
              rhythm = construct(ending_plan(plan, ending), reference, relationship, template, ending, variant, choices, octave)
              if rhythm:
                result.append(rhythm)

  def rank(rhythm):
    expected = float(plan.activity * (plan.end - plan.start))
    # Broad rhythm feasibility only. Pitch/rhythm relationships are evaluated
    # after realization; one attack cannot eliminate an otherwise useful tune.
    broad = abs(len(rhythm.slots) - expected) > max(4, expected / 2)
    return broad, generator.solo.policy.stable_rank(plan.seed, rhythm.id)

  def family(rhythm):
    pitches = tuple(a.pitch.number() for a in rhythm.anchors)
    shape = generator.solo.sequences.family(tuple(b - a for a, b in zip(pitches, pitches[1:])))
    return rhythm.relationship, rhythm.groups[0].source, shape

  limit = generator.solo.policy.DEFAULT.draft_frontier
  return generator.solo.sequences.diverse(result, rank, family, limit * (batch + 1))[limit * batch:]

def canonical_context(notes, incoming, plan):
  prefix = incoming.notes if incoming.passage_id == plan.passage_id and plan.passage_id else ((incoming.last,) if incoming.last else ())
  return generator.solo.types.sounding_notes((*prefix, *notes))

def local_sounding(notes):
  if notes and notes[0].continuation:
    notes = (dataclasses.replace(notes[0], continuation=False), *notes[1:])
  return generator.solo.types.sounding_notes(notes)

def contract_valid(draft, notes):
  by_id = {n.event_id: n for n in local_sounding(notes)}
  original = {n.event_id: n for n in draft.notes}
  for identity in (*draft.structural, *draft.protected):
    before, after = original[identity], by_id.get(identity)
    if after is None or (before.start, before.end, before.pitch) != (after.start, after.end, after.pitch):
      return False, f'{identity}: protected draft event changed'
  for transform in draft.transformations:
    if not transform.preserves:
      continue
    core = [by_id.get(identity) for identity in transform.events[:min(3, len(transform.events))]]
    if not all(core):
      return False, f'{transform.id}: identifying events missing'
    if 'rhythm' in transform.preserves and tuple((n.start - transform.start, n.end - transform.start) for n in core) != transform.rhythm[:len(core)]:
      return False, f'{transform.id}: identifying rhythm changed'
    directions = tuple((b.pitch.number() > a.pitch.number()) - (b.pitch.number() < a.pitch.number()) for a, b in zip(core, core[1:]))
    if 'opening_contour' in transform.preserves and directions != transform.directions[:len(directions)]:
      return False, f'{transform.id}: identifying contour changed'
  for feature in draft.features:
    if generator.solo.references.feature_error(feature, local_sounding(notes)):
      return False, f'{feature.id}: selected identifying relationship changed'
  return True, ''

def relationship_features(draft, notes):
  """Inspect relationships once, on canonical events, independent of cells."""
  notes = local_sounding(notes)
  by_id = {n.event_id: n for n in notes}
  preservation = sum(identity not in by_id for identity in draft.structural) / max(1, len(draft.structural))
  direction, development = [], []
  for transform in draft.transformations:
    group = [by_id[i] for i in transform.events if i in by_id]
    actual = tuple((b.pitch.number() > a.pitch.number()) - (b.pitch.number() < a.pitch.number()) for a, b in zip(group, group[1:]))
    direction.extend(a != b for a, b in zip(actual, transform.directions))
    reference = draft.reference_motif if transform.operation in ('statement', 'contrast') else draft.home
    if reference and transform.operation not in ('close', 'fragment', 'broaden_ending'):
      development.append(generator.solo.references.relationship(group, reference))
  joins = []
  for time in draft.joins:
    left = next((n for n in reversed(notes) if n.start < time), None)
    right = next((n for n in notes if n.start >= time), None)
    if left and right:
      original_left = next((n for n in reversed(draft.notes) if n.start < time), None)
      original_right = next((n for n in draft.notes if n.start >= time), None)
      planned_gap = max(F(), original_right.start - original_left.end) if original_left and original_right else F()
      gap = max(F(), right.start - left.end - planned_gap)
      joins.append(min(1, float(gap) / 2) + min(1, max(0, abs(right.pitch.number() - left.pitch.number()) - 4) / 12))
  terminal = next((a for a in draft.anchors if a.role == 'terminal' and a.earliest < draft.end), None)
  integration = 0.0
  if terminal:
    arrival = generator.solo.harmonic.anchor_note(terminal, notes)
    before = next((n for n in reversed(notes) if arrival and n.start < arrival.start), None)
    integration = min(1, max(0, abs(arrival.pitch.number() - before.pitch.number()) - 4) / 12) if arrival and before else 0
  return {'draft_preservation': preservation, 'development': sum(development) / max(1, len(development)), 'direction': sum(direction) / max(1, len(direction)), 'join_coherence': sum(joins) / max(1, len(joins)), 'ending_integration': integration}