"""Constrained, deterministically seeded autoregressive musical sampling.

The pooled model proposes joint note/rest, duration and pitch events. Complete
expression actions include their real resolution. Guitar and chord/function
validity precede sampling. No target access, hand-authored melody scorer or
legacy fallback. A bounded restart handles an exhausted/entirely silent sample.
"""
from __future__ import annotations
import collections
import dataclasses
import fractions
import hashlib
import math
import random

import generator.harmony
import generator.model
import generator.solo.expressions
import generator.solo.groups
import generator.solo.learned
import generator.solo.types
import generator.vocabulary

class SearchExhausted(ValueError):
  pass

@dataclasses.dataclass(frozen=True)
class State:
  events: tuple = ()
  notes: tuple = ()
  contracts: tuple = ()
  history: tuple = ()
  previous: int | None = None
  cost: float = 0.

def event_key(events): return tuple((a, b, -1 if p is None else p) for a, b, p in events)

class Decoder:
  def __init__(self, analysis, harmony, model, role, start, end, *, expressions_enabled=True, attempts=4, seed='', composition=None):
    self.composition = composition or generator.vocabulary.current()
    if type(attempts) is not int or not 1 <= attempts <= 16:
      raise ValueError('attempts must be between 1 and 16')
    self.analysis, self.harmony, self.model, self.role = analysis, harmony, model, role
    self.start, self.end = start, end
    self.attempts = attempts
    source = tuple((str(n.start), str(n.end), n.pitch.number()) for n in analysis.melody)
    self.seed = seed or hashlib.sha256(repr((source, role, generator.solo.learned.VERSION)).encode()).hexdigest()
    self.expressions_enabled = expressions_enabled
    self.low, self.high = harmony.profile.register
    limits = self.composition.data('roles')['difficulty']
    self.durations = tuple(d for d in model.durations if 0 < d <= 4 and d.denominator in (1, 2, 4) and
      (d.denominator != 4 or analysis.tempo <= limits['sixteenth_max_bpm']))
    self.allow_triplets = analysis.tempo <= limits['triplet_max_bpm']
    self._actions = {}
    self._context = {}
    self._pitches = {}
    self._colors = {}
    self._permissions = {}
    self.trace = []
    self.counts = collections.Counter()
  def context(self, time):
    if time not in self._context:
      harmonic, root, future = generator.solo.learned.chord_context(self.harmony, time)
      self._context[time] = (harmonic, root, generator.solo.learned.phase_at(self.analysis, time), generator.solo.learned.source_feature(self.analysis, time, root), future)
    return self._context[time]
  def pitch(self, number, time):
    span = self.harmony.at(time)
    identity = (number, span.start if span else None)
    if identity not in self._pitches:
      self._pitches[identity] = self._spell(number, time)
    return self._pitches[identity]
  def _spell(self, number, time):
    span = self.harmony.at(time)
    if span and span.chord != 'N.C.':
      chord = generator.solo.expressions.parsed_chord(span.chord)
      for degree, alter in generator.harmony.degrees(chord).items():
        p = generator.solo.expressions.degree_pitch(chord, degree, alter, number)
        if p.number() == number:
          return p
      for letter, alter in zip(generator.model.LETTERS, span.key['alterations']):
        accidental = generator.model.fraction(alter)
        if (generator.model.NATURAL[letter] + accidental) % 12 == number % 12:
          return generator.solo.types.Pitch(letter, accidental, int((number - generator.model.NATURAL[letter] - accidental) // 12) - 1)
    names = (('C', 0), ('D', -1), ('D', 0), ('E', -1), ('E', 0), ('F', 0), ('F', 1), ('G', 0), ('A', -1), ('A', 0), ('B', -1), ('B', 0))
    letter, alter = names[number % 12]
    return generator.solo.types.Pitch(letter.lower(), fractions.Fraction(alter), number // 12 - 1)
  def note(self, a, b, p, role='learned'):
    pitch = p if isinstance(p, generator.solo.types.Pitch) else self.pitch(p, a)
    return generator.solo.types.Note(a, b, pitch, role=role, event_id=f'{a}:{b}:{pitch.number()}')
  def color_bases(self, n):
    key = (n.start, n.end, n.pitch.number())
    if key in self._colors:
      return self._colors[key]
    bases = []
    regions = [(None, n)]
    regions.extend((s.start, dataclasses.replace(n, start=s.start)) for s in self.harmony.spans
      if n.start < s.start < n.end and s.start - n.start <= fractions.Fraction(self.composition.data('harmony')['anticipation_max']))
    for boundary, arrived in regions:
      exposed = [s for s, _, _ in self.harmony.segments(arrived) if n.pitch.number() % 12 not in s.roles.basic | s.roles.specified]
      if not exposed:
        continue
      for rule in self.composition.data('harmony')['colors']:
        if arrived.end - arrived.start > fractions.Fraction(rule['max_beats']):
          continue
        if any(s.chord == 'N.C.' or not generator.solo.expressions.eligible_color(rule, s, self.harmony) or
          generator.solo.expressions.degree_pitch(generator.solo.expressions.parsed_chord(s.chord), rule['degree'], rule['alter'], n.pitch.number()).number() != n.pitch.number()
          for s in exposed):
          continue
        bases.append((rule, boundary))
    self._colors[key] = tuple(bases)
    return self._colors[key]
  def color_contract(self, n, target, rule, boundary):
    span = self.harmony.at(boundary if boundary is not None else n.start)
    if not span:
      return None
    kind = 'anticipated_color' if boundary is not None else 'color'
    return generator.solo.types.ExpressionContract(n.event_id + ':' + rule['id'] + ':' + str(boundary), kind, rule['id'], (n,), target, span.chord, boundary)
  def permission(self, n):
    key = (n.start, n.end, n.pitch.number())
    if key not in self._permissions:
      self._permissions[key] = self._permission(n)
    return self._permissions[key]
  def _permission(self, n):
    """Find justified color/anticipation support, never authorize by label."""
    # Source spans are contiguous. These are precisely the ordinary-event
    # branches of Harmony.evaluate, without rebuilding an empty contract map.
    if self.harmony.at(n.start) and n.end <= self.harmony.spans[-1].end:
      if n.pitch.number() % 12 in self.harmony.members(n.start, n.end):
        return ()
      if n.end - n.start < 2 and n.pitch.number() % 12 in self.harmony.permitted(n.start, n.end):
        return ()
    if not self.expressions_enabled:
      return None
    span = self.harmony.at(n.start)
    if not span or span.chord == 'N.C.':
      return None
    for rule, boundary in self.color_bases(n):
      if rule.get('resolve'):
        continue
      c = self.color_contract(n, n, rule, boundary)
      if self.harmony.evaluate((n,), complete=True, expressions=(c,), infer=False).valid:
        return tuple([c])
    for arriving in self.harmony.spans:
      if n.start < arriving.start < n.end and arriving.start - n.start <= fractions.Fraction(self.composition.data('harmony')['anticipation_max']):
        c = generator.solo.types.ExpressionContract(n.event_id + ':anticipation', 'anticipation', 'early_chord_tone', (n,), n, arriving.chord, arriving.start)
        if self.harmony.evaluate((n,), complete=True, expressions=(c,), infer=False).valid:
          return tuple([c])
    return None
  def tails(self, n, target):
    """Finite resolution timing: direct, real space, or one valid intervening note."""
    gap = target.start - n.end
    if gap < 0:
      return
    if gap == 0:
      yield ((n.start, n.end, n.pitch.number()), (target.start, target.end, target.pitch.number())), (n, target)
      return
    yield ((n.start, n.end, n.pitch.number()), (n.end, target.start, None), (target.start, target.end, target.pitch.number())), (n, target)
    windows = {(n.end, target.start)} if gap <= 1 else set()
    if gap >= 1:
      windows |= {(n.end, n.end + fractions.Fraction(1, 2)), (target.start - fractions.Fraction(1, 2), target.start)}
    for a, b in sorted(windows):
      for number in range(max(self.low, n.pitch.number() - 3), min(self.high, n.pitch.number() + 3) + 1):
        mid = self.note(a, b, number)
        if number in (n.pitch.number(), target.pitch.number()) or not self.harmony.ordinary(mid):
          continue
        events = [(n.start, n.end, n.pitch.number())]
        if a > n.end:
          events.append((n.end, a, None))
        events.append((a, b, number))
        if b < target.start:
          events.append((b, target.start, None))
        events.append((target.start, target.end, target.pitch.number()))
        yield tuple(events), (n, mid, target)
  def resolutions(self, n):
    """Joint color/delayed-approach actions; the destination is never optional."""
    requests = []
    for rule, boundary in self.color_bases(n):
      if rule.get('resolve'):
        for gap in (fractions.Fraction(), fractions.Fraction(1, 2), fractions.Fraction(1)):
          if gap <= fractions.Fraction(self.composition.data('harmony')['resolution_gap']):
            requests.append(('color', rule, boundary, n.end + gap, (-2, -1, 1, 2)))
    delay = self.composition.data('lead_ins')['delayed_approach']
    if n.end - n.start <= fractions.Fraction(delay['max_note']) and n.pitch.number() % 12 not in self.harmony.permitted(n.start, n.end):
      for arrival in map(fractions.Fraction, delay['arrivals']):
        if n.end <= n.start + arrival <= n.start + fractions.Fraction(delay['deadline']):
          requests.append(('delayed', delay, None, n.start + arrival, delay['steps']))
    for kind, rule, boundary, arrival, steps in requests:
      for delta in steps:
        number = n.pitch.number() + delta
        if not self.low <= number <= self.high:
          continue
        for hold in (fractions.Fraction(1, 2), fractions.Fraction(1), fractions.Fraction(2)):
          if arrival + hold > self.end or number % 12 not in self.harmony.members(arrival, arrival + hold):
            continue
          target = self.note(arrival, arrival + hold, number, 'destination')
          c = self.color_contract(n, target, rule, boundary) if kind == 'color' else generator.solo.types.ExpressionContract(n.event_id + ':delayed', 'delayed_approach', rule['id'], (n,), target, self.harmony.at(arrival).chord)
          # Color and delayed-approach contracts permit this intervening
          # window. Validate the invariant endpoints once; tails() checks
          # every inserted note's ordinary full-duration permission itself.
          if not self.harmony.evaluate((n, target), complete=True, expressions=(c,), infer=False).valid:
            continue
          for events, ns in self.tails(n, target):
            if not generator.solo.groups.valid(tuple((a, b) for a, b, _ in events)):
              continue
            yield events, ns, (c,)
  def actions(self, time):
    if time in self._actions:
      return self._actions[time]
    result = []
    def add(events, notes, contracts=()): result.append((events, notes, contracts))
    for duration in self.durations:
      end = time + duration
      if end > self.end:
        continue
      add(((time, end, None),), ())
      for number in range(self.low, self.high + 1):
        n = self.note(time, end, number)
        support = self.permission(n)
        if support is not None:
          add(((time, end, number),), (n,), support)
        if self.expressions_enabled and not self.harmony.ordinary(n):
          for events, ns, contracts in self.resolutions(n):
            add(events, ns, contracts)
          rule = self.composition.data('lead_ins')['passing']
          if duration <= fractions.Fraction(rule['max_note']):
            for step in rule['steps']:
              p = number + step
              if not self.low <= p <= self.high:
                continue
              for hold in (fractions.Fraction(1, 2), fractions.Fraction(1), fractions.Fraction(2)):
                if end + hold > self.end:
                  continue
                target = self.note(end, end + hold, p)
                extra = self.permission(target)
                if extra is None or any(c.kind not in ('color', 'anticipated_color') for c in extra):
                  continue
                c = generator.solo.types.ExpressionContract(n.event_id + ':passing', 'passing', rule['id'], (n,), target, self.harmony.at(end).chord)
                add(((time, end, number), (end, end + hold, p)), (n, target), (c, *extra))
    if self.expressions_enabled:
      limits = self.composition.data('roles')['difficulty']
      for rule in self.composition.data('lead_ins')['patterns']:
        difficulty = rule.get('difficulty')
        if difficulty and self.analysis.tempo > limits[difficulty + '_max_bpm']:
          continue
        # Connections require a concrete preceding note, checked at expansion.
        durations = tuple(map(fractions.Fraction, rule['durations']))
        gaps = tuple(map(fractions.Fraction, rule.get('internal_gaps', ['0'] * (len(durations) - 1))))
        target_time = time + sum(durations, fractions.Fraction()) + sum(gaps, fractions.Fraction()) + fractions.Fraction(rule['gap'])
        span = self.harmony.at(target_time)
        if not span or span.chord == 'N.C.':
          continue
        holds = (fractions.Fraction(1, 3),) if difficulty == 'triplet' else (fractions.Fraction(1, 2), fractions.Fraction(1), fractions.Fraction(2))
        for duration in holds:
          end = target_time + duration
          if end > self.end:
            continue
          for number in range(self.low, self.high + 1):
            if number % 12 not in self.harmony.members(target_time, end):
              continue
            target = self.note(target_time, end, number, 'destination')
            cursor = time
            notes = []
            events = []
            for i, (delta, d) in enumerate(zip(rule['semitones'], durations)):
              n = self.note(cursor, cursor + d, generator.solo.expressions.chromatic(target.pitch, delta), rule['kind'])
              notes.append(n)
              events.append((n.start, n.end, n.pitch.number()))
              cursor = n.end
              gap = gaps[i] if i < len(gaps) else fractions.Fraction(rule['gap'])
              if gap:
                events.append((cursor, cursor + gap, None))
                cursor += gap
            if any(not self.low <= n.pitch.number() <= self.high for n in notes):
              continue
            if not generator.solo.groups.valid(tuple((a, b) for a, b, p in (*events, (target_time, end, number)))):
              continue
            contract = generator.solo.types.ExpressionContract(f'{time}:{number}:{rule["id"]}', rule['kind'], rule['id'], tuple(notes), target, span.chord)
            ns = (*notes, target)
            if rule['kind'] == 'connection':
              # Validation waits for the actual incoming event.
              add((*events, (target_time, end, number)), ns, (contract,))
            elif self.harmony.evaluate(ns, complete=True, expressions=(contract,), infer=False).valid:
              add((*events, (target_time, end, number)), ns, (contract,))
    # Complete ordinary triplet groups; do not leave a third-beat residual.
    if self.allow_triplets and time + 1 <= self.end:
      for number in range(self.low, self.high + 1):
        for direction in (-1, 0, 1):
          numbers = [number]
          cursor = time
          for _ in range(2):
            nxt = next((p for p in range(numbers[-1] + direction, numbers[-1] + direction * 4, direction or 1)
              if self.low <= p <= self.high and p % 12 in self.harmony.permitted(cursor + fractions.Fraction(1, 3), cursor + fractions.Fraction(2, 3))), None) if direction else number
            if nxt is None:
              break
            numbers.append(nxt)
            cursor += fractions.Fraction(1, 3)
          if len(numbers) != 3:
            continue
          ns = tuple(self.note(time + i * fractions.Fraction(1, 3), time + (i + 1) * fractions.Fraction(1, 3), p) for i, p in enumerate(numbers))
          if self.harmony.evaluate(ns, complete=True, infer=False).valid:
            add(tuple((n.start, n.end, n.pitch.number()) for n in ns), ns)
    self._actions[time] = tuple(result)
    self.counts['valid_structures'] += len(result)
    return self._actions[time]
  def extend(self, state, action):
    events, notes, contracts = action
    # Training tokenization merges contiguous silence. Never manufacture a
    # rest-to-rest transition absent from that event representation.
    if state.events and state.events[-1][2] is None and events[0][2] is None:
      return None
    if notes and state.previous is not None and abs(notes[0].pitch.number() - state.previous) > 12:
      return None
    if any(c.kind in ('connection', 'passing') and c.preparation is None for c in contracts):
      if not state.notes:
        return None
      contracts = tuple(dataclasses.replace(c, preparation=state.notes[-1]) if c.kind in ('connection', 'passing') and c.preparation is None else c for c in contracts)
      if self.harmony.ordinary(state.notes[-1]):
        _, error = generator.solo.expressions.validate((state.notes[-1], *notes), contracts, self.harmony, composition=self.composition)
      else:
        _, error = generator.solo.expressions.validate((*state.notes, *notes), (*state.contracts, *contracts), self.harmony, composition=self.composition)
      if error:
        return None
    history = state.history
    previous = state.previous
    cost = state.cost
    for a, b, p in events:
      harmonic, root, phase, source, _ = self.context(a)
      action = generator.solo.learned.token(b - a, p, previous, root)
      context = generator.solo.learned.contexts(harmonic, phase, source, history)
      cost += self.model.cost(action, context)
      if p is not None:
        if previous is None:
          cost += self.model.register_cost(p)
        previous = p
      history = (*history, action)[-4:]
    return State((*state.events, *events), (*state.notes, *notes), (*state.contracts, *contracts), history, previous, cost)
  def compose(self):
    for attempt in range(self.attempts):
      rng = random.Random(self.seed + ':' + str(attempt))
      state = State()
      time = self.start
      trace = []
      while time < self.end:
        actions = self.actions(time)
        options = [s for a in actions if (s := self.extend(state, a)) is not None]
        # Different operator names cannot multiply the probability of identical
        # sounded music. Prefer plain permission when no exception is needed.
        unique = {}
        for s in options:
          signature = event_key(s.events[len(state.events):])
          if signature not in unique or len(s.contracts) < len(unique[signature].contracts):
            unique[signature] = s
        options = sorted(unique.values(), key=lambda s: event_key(s.events[len(state.events):]))
        self.counts['scored_continuations'] += len(options)
        if not options:
          break
        low = min(s.cost - state.cost for s in options)
        weights = [math.exp(-(s.cost - state.cost - low)) for s in options]
        draw = rng.random() * sum(weights)
        selected = options[-1]
        for s, w in zip(options, weights):
          draw -= w
          if draw <= 0:
            selected = s
            break
        trace.append({'time': str(time), 'valid_joint_actions': len(options),
          'selected_events': [(str(a), str(b), p) for a, b, p in selected.events[len(state.events):]],
          'negative_log_likelihood': selected.cost - state.cost})
        state = selected
        time = state.events[-1][1]
      if time == self.end and state.notes:
        self.trace = trace
        self.completed = (state,)
        self.counts['attempts'] = attempt + 1
        return state
    raise SearchExhausted(f'{self.role}: no complete non-silent sample within {self.attempts} deterministic attempts; end={self.end}')