from __future__ import annotations

import dataclasses
import fractions
import functools
import heapq

import generator.fretboard
import generator.model
import generator.rhythm
import generator.scales.selection
import generator.solo.placement
import generator.solo.planning
import generator.solo.types

BEAM_WIDTH = 64
CELL_LIMIT = 32

class SearchExhausted(ValueError):
  pass

@dataclasses.dataclass(frozen=True)
class Cell:
  start: fractions.Fraction
  end: fractions.Fraction
  patterns: tuple[tuple[str, tuple[tuple[fractions.Fraction, fractions.Fraction], ...]], ...]
  motif: generator.solo.types.Motif | None = None

@dataclasses.dataclass(frozen=True)
class Candidate:
  notes: tuple[generator.solo.types.Note, ...]
  source: str
  treatment: str
  rank: tuple

def rhythm_patterns(library: generator.rhythm.Library, meter: str, length: fractions.Fraction, activity: fractions.Fraction):
  result = {}
  for snippet in library.for_meter(meter):
    for index, bar in enumerate(snippet.bars):
      spans = []
      for event in bar.events:
        if not event.sounding:
          continue
        if event.continuation:
          if spans and spans[-1][1] == event.onset:
            spans[-1] = (spans[-1][0], event.onset + event.duration)
          else:
            # A continuation without its attack is not an independent cell.
            spans = []
            break
        else:
          spans.append((event.onset, event.onset + event.duration))
      pattern = tuple(spans)
      if all(end <= length for _, end in pattern) and len(pattern) <= 4:
        result.setdefault(pattern, f'rhythm {snippet.id}:{index + 1}')
  ranked = sorted(result.items(), key=lambda item: (abs(len(item[0]) - length * activity), -sum(end - start for start, end in item[0]), item[1]))
  return tuple((identifier, pattern) for pattern, identifier in ranked[:8])

def cells(plan, measures, rhythms):
  result, start, index = [], fractions.Fraction(), 0
  boundaries = [fractions.Fraction()]
  for measure in measures:
    boundaries.append(boundaries[-1] + generator.model.fraction(measure['duration']))
  start = plan.start
  music_end = plan.end - plan.release
  arrival_fits = not plan.target or plan.motif.duration < music_end - start or any(plan.target.earliest <= start + n.start <= plan.target.latest for n in plan.motif.notes)
  if plan.motif.duration <= music_end - start and arrival_fits:
    motif = plan.motif
    spans = tuple((n.start, n.end) for n in motif.notes)
    patterns = [(motif.id, spans)]
    if plan.relationship == 'variation' and spans[-1][1] + fractions.Fraction(1, 2) <= motif.duration:
      patterns.insert(0, (motif.id + ':displaced', tuple((a + fractions.Fraction(1, 2), b + fractions.Fraction(1, 2)) for a, b in spans)))
    result.append(Cell(start, start + motif.duration, tuple(patterns), motif))
    start += motif.duration
  elif plan.motif.notes and plan.motif.duration > music_end - start:
    # Fragmentation is an explicit transformation, never a truncated quotation.
    available = (plan.target.earliest if plan.target else music_end) - start
    notes = tuple(n for n in plan.motif.notes if n.end <= available)
    if notes:
      motif = dataclasses.replace(plan.motif, notes=notes, duration=notes[-1].end)
      result.append(Cell(start, start + motif.duration, ((motif.id + ':fragment', tuple((n.start, n.end) for n in notes)),), motif))
      start += motif.duration
  while start < music_end:
    index = next(i for i in range(len(measures)) if boundaries[i] <= start < boundaries[i + 1])
    end = min(music_end, boundaries[index + 1])
    patterns = rhythm_patterns(rhythms, measures[index]['time_signature'], end - start, plan.activity)
    # These primitives are deliberate alternatives, independent of catalog counts.
    patterns += (('hold', ((fractions.Fraction(), end - start),)), ('space', ()))
    if plan.target and start <= plan.target.latest < end:
      onset = max(start, plan.target.earliest) - start
      patterns = (('arrival_and_release', ((onset, min(end - start, onset + 1)),)), *patterns)
    result.append(Cell(start, end, patterns))
    start = end
  result.append(Cell(music_end, plan.end, (('planned_release', ()),)))
  return tuple(result)

@functools.lru_cache(maxsize=2048)
def scale_options(scale, allowed, low, high):
  return tuple(generator.solo.types.scale_pitch(scale, index) for index in range(-21, 28) if low <= generator.solo.types.scale_pitch(scale, index).number() <= high and generator.solo.types.scale_pitch(scale, index).number() % 12 in allowed and generator.fretboard.positions_for_pitch(fractions.Fraction(generator.solo.types.scale_pitch(scale, index).number())))

def valid_notes(notes, choices):
  return all(n.pitch.number() % 12 in generator.solo.types.permitted(choices, n.start, n.end) for n in notes)

def options(choices, start, end, register):
  scale = next((c.scale for c in choices if c.span.start <= start < c.span.end), None)
  allowed = generator.solo.types.permitted(choices, start, end)
  return scale_options(scale, allowed, max(41, register[0] - 7), min(84, register[1] + 5)) if scale and allowed else ()

def target_ok(notes, target):
  return target is None or any(target.earliest <= n.start <= target.latest and n.pitch.number() % 12 in target.pitches for n in notes)

def thematic_paths(cell, spans, choices, plan):
  states = [(0, ())]
  for index, (start, end) in enumerate(spans):
    following = []
    for cost, path in states:
      available = options(choices, start, end, plan.register)
      changing_ending = plan.relationship == 'answer' and index == len(spans) - 1
      if changing_ending:
        chord = next((c.span.chord for c in choices if c.span.start <= start < c.span.end), None)
        if chord:
          available = tuple(p for p in available if p.number() % 12 in generator.solo.planning.target_pitches(chord, False))
      for pitch in available:
        reference = cell.motif.notes[index]
        expected = cell.motif.contour()[index - 1] if index else 0
        motion = pitch.number() - path[-1].pitch.number() if path else 0
        # Repeated notes are identity, not missing sequence transitions.
        mismatch = int(bool(index) and ((expected == 0) != (motion == 0) or expected * motion < 0))
        note = generator.solo.types.Note(start, end, pitch, reference.sources, 'motif')
        following.append((cost + (0 if changing_ending else 8 * mismatch + abs(motion - expected)), (*path, note)))
    states = sorted(following, key=lambda item: (item[0], tuple(n.pitch for n in item[1])))[:BEAM_WIDTH]
    if not states:
      break
  return states

def adapt_path(pitches, spans, choices, register):
  """Connect actual pitches across collections without carrying a stale degree index."""
  states = [(0, ())]
  for index, ((start, end), intended) in enumerate(zip(spans, pitches)):
    available = options(choices, start, end, register)
    following = []
    for cost, path in states:
      desired = path[-1].pitch.number() + intended.number() - pitches[index - 1].number() if path else intended.number()
      for pitch in sorted(available, key=lambda p: (abs(p.number() - desired), p))[:2]:
        following.append((cost + abs(pitch.number() - desired), (*path, generator.solo.types.Note(start, end, pitch, role='adapted_path'))))
    states = sorted(following, key=lambda item: (item[0], tuple(n.pitch for n in item[1])))[:8]
  return [notes for _, notes in states]

def sequence_paths(spans, choices, plan, sequences):
  if not spans:
    return [((), 'space')]
  first_scale = next((c.scale for c in choices if c.span.start <= spans[0][0] < c.span.end), None)
  if first_scale is None:
    return []
  paths = []
  allowed = tuple(generator.solo.types.permitted(choices, start, end) for start, end in spans)
  changing = any(classes != allowed[0] for classes in allowed)
  sequence_options = sequences.connecting(len(spans))[:16]
  offsets = [(s.offsets(), 'sequence ' + s.id) for s in sequence_options]
  # Single notes and explicit repetitions do not pretend to be catalog paths.
  offsets.append(((0,) * len(spans), 'single' if len(spans) == 1 else 'repeated_note'))
  for steps, identifier in offsets:
    for target_index in range(-14, 21):
      start_index = target_index - steps[-1]
      pitches = tuple(generator.solo.types.scale_pitch(first_scale, start_index + step) for step in steps)
      if not all(max(41, plan.register[0] - 7) <= p.number() <= min(84, plan.register[1] + 5) for p in pitches):
        continue
      if all(pitch.number() % 12 in classes for pitch, classes in zip(pitches, allowed)):
        notes = tuple(generator.solo.types.Note(start, end, pitch, role='path') for (start, end), pitch in zip(spans, pitches))
        paths.append((notes, identifier))
      elif changing and not paths and steps == offsets[0][0]:
        paths.extend((notes, 'harmonic_adaptation of ' + identifier) for notes in adapt_path(pitches, spans, choices, plan.register))
  return paths

def candidates(cell, plan, choices, sequences):
  result = []
  target = plan.target if plan.target and cell.start <= plan.target.latest < cell.end else None
  for identifier, pattern in cell.patterns:
    spans = tuple((cell.start + start, cell.start + end) for start, end in pattern)
    active = tuple(i for i, (start, end) in enumerate(spans) if not any(c.scale is None and c.span.start < end and c.span.end > start for c in choices))
    local_cell = cell
    if len(active) != len(spans):
      spans = tuple(spans[i] for i in active)
      identifier += ':harmonic_fragment'
      if cell.motif:
        local_cell = dataclasses.replace(cell, motif=dataclasses.replace(cell.motif, notes=tuple(cell.motif.notes[i] for i in active)))
    paths = []
    if local_cell.motif and spans:
      if plan.treatment == 'quote' and ':fragment' not in identifier and ':harmonic_fragment' not in identifier and ':displaced' not in identifier:
        for octaves in (0, 1, -1, 2, -2):
          notes = tuple(generator.solo.types.Note(start, end, dataclasses.replace(source.pitch, octave=source.pitch.octave + octaves), source.sources, 'quote') for (start, end), source in zip(spans, local_cell.motif.notes))
          if valid_notes(notes, choices):
            paths.append((notes, f'quote octave {octaves:+}', 0))
      for mismatch, notes in thematic_paths(local_cell, spans, choices, plan):
        treatment = 'fragment' if 'fragment' in identifier else 'rhythmic_displacement' if ':displaced' in identifier else 'changed_ending' if plan.relationship == 'answer' else 'harmonic_adaptation'
        paths.append((notes, treatment, 1 + mismatch))
    else:
      paths = [(notes, name, 0) for notes, name in sequence_paths(spans, choices, plan, sequences)]
    for notes, treatment, identity in paths:
      if not target_ok(notes, target):
        continue
      notes = tuple(dataclasses.replace(n, role='target' if target and target.earliest <= n.start <= target.latest and n.pitch.number() % 12 in target.pitches else n.role) for n in notes)
      register = 16 * sum(max(0, plan.register[0] - n.pitch.number(), n.pitch.number() - plan.register[1]) for n in notes)
      if notes:
        register += abs(max(n.pitch.number() for n in notes) - sum(plan.register) // 2)
      density = abs(len(notes) - (cell.end - cell.start) * plan.activity)
      contour = sum(max(0, abs(b.pitch.number() - a.pitch.number()) - 5) for a, b in zip(notes, notes[1:]))
      if not cell.motif and len({n.pitch.number() for n in notes}) == 1:
        contour += max(0, len(notes) - 2)
      rank = (identity, register, density if not cell.motif else 0, contour)
      result.append(Candidate(notes, identifier, treatment, rank))
  # A no-harmony passage is intentional silence, not a failed pitch search.
  if not any(c.span.chord and c.span.start < cell.end and c.span.end > cell.start for c in choices):
    result = [Candidate((), 'no_harmony', 'space', (0, 0, 0, 0))]
  return sorted(result, key=lambda c: (c.rank, c.source, c.treatment, tuple(n.pitch for n in c.notes)))

def compose(plan, measures, choices, sequences, rhythms, tempo, previous=None):
  states = [((0, 0, 0, 0, 0, 0, 0), (), previous)]
  for cell in cells(plan, measures, rhythms):
    pool = candidates(cell, plan, choices, sequences)
    following, extensions = [], {}
    for cost, decisions, last in states:

      def extension_rank(candidate):
        jump = abs(candidate.notes[0].pitch.number() - last.pitch.number()) if candidate.notes and last else 0
        return candidate.rank, jump, candidate.source, candidate.treatment, tuple(n.pitch for n in candidate.notes)

      key = last.pitch.number() if last else None
      if key not in extensions:
        extensions[key] = heapq.nsmallest(CELL_LIMIT, pool, key=extension_rank)
      for candidate in extensions[key]:
        placed = generator.solo.placement.place(candidate.notes, tempo, last)
        if placed is None:
          continue
        comfort, notes = placed
        jump = abs(notes[0].pitch.number() - last.pitch.number()) if notes and last else 0
        monotony = int(bool(notes and decisions and decisions[-1].notes) and tuple(n.pitch for n in notes) == tuple(n.pitch for n in decisions[-1].notes) and cell.motif is None)
        local = (*candidate.rank, jump, comfort, monotony)
        total = tuple(a + b for a, b in zip(cost, local))
        decision = generator.solo.types.Decision(cell.start, cell.end, candidate.source, candidate.treatment, notes)
        following.append((total, (*decisions, decision), notes[-1] if notes else last))
    if not following:
      raise SearchExhausted(f'{plan.id}, beats {cell.start + 1}-{cell.end + 1}: bounded solo search exhausted; no feasible phrase extension among {BEAM_WIDTH} states / {CELL_LIMIT} extensions (not proof that the passage is unplayable)')
    following.sort(key=lambda item: (item[0], tuple((d.source, d.treatment, tuple((n.pitch, n.position) for n in d.notes)) for d in item[1])))
    # Retain distinct endings for lookahead instead of 64 copies of one path.
    unique = {}
    for item in following:
      last = item[2]
      key = (last.pitch, last.position, last.end) if last else None
      unique.setdefault(key, item)
    states = list(unique.values())[:BEAM_WIDTH]
  _, decisions, _ = states[0]
  return generator.solo.types.Phrase(plan, decisions)