from __future__ import annotations

import dataclasses
import fractions
import functools
import collections

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

BEAM_WIDTH = generator.solo.policy.DEFAULT.beam_width
CELL_LIMIT = generator.solo.policy.DEFAULT.extensions

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
  fidelity: float = 0

def rhythm_patterns(library: generator.rhythm.Library, meter: str, length: fractions.Fraction, activity: fractions.Fraction, offset=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:
          spans.append((event.onset, event.onset + event.duration))
      # Slice exact authored rhythm; later attacks outside the window do not
      # disqualify a useful prefix. A tied continuation never becomes an attack.
      pattern = tuple((start - offset, min(end - offset, length)) for start, end in spans if offset <= start < offset + length)
      result.setdefault(pattern, f'rhythm {snippet.id}:{index + 1}@{offset}+{length}')

  def rank(item):
    pattern, identifier = item
    return abs(len(pattern) - length * activity), abs(sum(end - start for start, end in pattern) - length * fractions.Fraction(3, 4)), identifier

  def family(item):
    pattern = item[0]
    return len(pattern), bool(pattern and pattern[0][0]), any(a.denominator > 1 for a, _ in pattern)

  selected = generator.solo.sequences.diverse(result.items(), rank, family, 12)
  return tuple((identifier, pattern) for pattern, identifier in selected)

def cells(plan, measures, rhythms):
  result, boundaries = [], [fractions.Fraction()]
  for measure in measures:
    boundaries.append(boundaries[-1] + generator.model.fraction(measure['duration']))
  start, music_end = plan.start, plan.end - plan.release
  protected = plan.treatment in ('quote', 'exact_quote', 'rhythmic_reference')
  if protected:
    motif = plan.motif
    if not motif.notes or motif.duration > music_end - start:
      raise SearchExhausted(f'{plan.id}: protected motif rhythm does not fit the phrase')
    spans = tuple((n.start, n.end) for n in motif.notes)
    result.append(Cell(start, start + motif.duration, ((motif.id, spans),), motif))
    start += motif.duration
  while start < music_end:
    index = next(i for i in range(len(measures)) if boundaries[i] <= start < boundaries[i + 1])
    passage = next((p for p in plan.passages if p.start <= start < p.end), None)
    end = min(music_end, boundaries[index + 1], passage.end if passage else music_end)
    activity = passage.activity if passage else plan.activity
    patterns = rhythm_patterns(rhythms, measures[index]['time_signature'], end - start, activity, start - boundaries[index])
    patterns += (('hold', ((fractions.Fraction(), end - start),)), ('space', ()))
    if plan.target and start <= plan.target.latest and plan.target.earliest < end:
      onset = max(start, plan.target.earliest) - start
      patterns = (('arrival_and_release', ((onset, end - start),)), *patterns)
    motif = plan.motif if not result else None
    result.append(Cell(start, end, patterns, motif))
    start = end
  if plan.release:
    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.0, ())]
  for index, (start, end) in enumerate(spans):
    following = []
    available = options(choices, start, end, plan.register)
    for _, path in states:
      for pitch in available:
        reference = cell.motif.notes[min(len(cell.motif.notes) - 1, index * len(cell.motif.notes) // max(1, len(spans)))]
        note = generator.solo.types.Note(start, end, pitch, reference.sources, 'motif')
        notes = (*path, note)
        cost = generator.solo.scoring.motif_fidelity(notes, cell.motif, plan.treatment)
        direction = plan.landmarks[1].direction if len(plan.landmarks) > 1 else 0
        motion = notes[-1].pitch.number() - notes[0].pitch.number()
        cost += 0.1 * (direction and motion * direction < 0)
        cost += sum(generator.solo.scoring.corridor(n.pitch.number(), *plan.register) for n in notes) / (12 * len(notes))
        following.append((cost, notes))
    states = list(generator.solo.sequences.diverse(following, lambda item: (item[0], tuple(n.pitch for n in item[1])), lambda item: (item[1][-1].pitch.number(), generator.solo.sequences.family(tuple(b.pitch.number() - a.pitch.number() for a, b in zip(item[1], item[1][1:])))), BEAM_WIDTH))
    if not states:
      break
  return [(generator.solo.scoring.motif_fidelity(notes, cell.motif, plan.treatment), notes) for _, notes in 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 []
  allowed = tuple(generator.solo.types.permitted(choices, start, end) for start, end in spans)
  if not all(allowed):
    return []
  available = options(choices, *spans[0], plan.register)
  if len(spans) == 1:
    return [((generator.solo.types.Note(*spans[0], pitch, role='path'),), 'single') for pitch in available]
  landmark = next((p for p in plan.landmarks if p.latest >= spans[0][0]), None)
  center = (landmark.low + landmark.high) // 2 if landmark else sum(plan.register) // 2
  direction = landmark.direction if landmark else 0
  starts = sorted(available, key=lambda p: (generator.solo.scoring.corridor(p.number(), *plan.register), abs(p.number() - center), p))[:8]
  paths, identities = [], set()
  for first in starts:
    degree = next(i for i in range(-21, 28) if generator.solo.types.scale_pitch(first_scale, i) == first)
    destination = center
    if plan.target and spans[-1][0] >= plan.target.earliest:
      compatible = [p for p in options(choices, *spans[-1], plan.register) if p.number() % 12 in plan.target.pitches]
      if compatible:
        destination = min(compatible, key=lambda p: (abs(p.number() - center), p)).number()
    retrieved = sequences.retrieve(len(spans), round((destination - first.number()) / 2), direction)
    routes = [(p.offsets(), 'sequences ' + ' + '.join(p.sources)) for p in retrieved]
    routes.append(((0,) * len(spans), 'repeated_note'))
    for offsets, identifier in routes:
      pitches = tuple(generator.solo.types.scale_pitch(first_scale, degree + step) for step in offsets)
      if not all(max(41, plan.register[0] - 7) <= p.number() <= min(84, plan.register[1] + 5) for p in pitches):
        continue
      if all(p.number() % 12 in classes for p, classes in zip(pitches, allowed)):
        realized = [tuple(generator.solo.types.Note(start, end, pitch, role='path') for (start, end), pitch in zip(spans, pitches))]
      elif any(classes != allowed[0] for classes in allowed):
        realized = adapt_path(pitches, spans, choices, plan.register)
        identifier = 'harmonic_adaptation of ' + identifier
      else:
        continue
      for notes in realized:
        identity = tuple(n.pitch for n in notes)
        if identity not in identities:
          identities.add(identity)
          paths.append((notes, identifier))
  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
  local_plan = dataclasses.replace(plan, start=cell.start, end=cell.end, release=fractions.Fraction(), passages=tuple(dataclasses.replace(p, start=max(cell.start, p.start), end=min(cell.end, p.end)) for p in plan.passages if p.start < cell.end and p.end > cell.start))
  evaluation = generator.solo.scoring.Evaluation(local_plan, choices)
  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):
      if cell.motif and plan.treatment in ('quote', 'exact_quote'):
        continue
      spans = tuple(spans[i] for i in active)
      identifier += ':harmonic_fragment'
      if cell.motif:
        references = tuple(cell.motif.notes[min(len(cell.motif.notes) - 1, i * len(cell.motif.notes) // len(pattern))] for i in active)
        local_cell = dataclasses.replace(cell, motif=dataclasses.replace(cell.motif, notes=references))
    paths = []
    if local_cell.motif and spans:
      if plan.treatment in ('quote', 'exact_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))
      if plan.treatment not in ('quote', 'exact_quote'):
        for mismatch, notes in thematic_paths(local_cell, spans, choices, plan):
          treatment = 'fragment' if 'fragment' in identifier else plan.treatment
          paths.append((notes, treatment, mismatch))
    else:
      paths = [(notes, name, 0) for notes, name in sequence_paths(spans, choices, plan, sequences)]
    for notes, treatment, identity in paths:
      notes = tuple(dataclasses.replace(n, role='target' if n.role not in ('motif', 'quote') and target and target.earliest <= n.start <= target.latest and n.pitch.number() % 12 in target.pitches else n.role) for n in notes)
      features = evaluation.features(notes)
      fidelity = identity if cell.motif else 0
      features['fidelity'] = fidelity
      rank = (generator.solo.scoring.value(features),)
      result.append(Candidate(notes, identifier, treatment, rank, fidelity))
  # A no-harmony passage is intentional silence, not a failed pitch search.
  if not (cell.motif and plan.treatment in ('quote', 'exact_quote')) and 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,))]
  return sorted(result, key=lambda c: (c.rank, c.source, c.treatment, tuple(n.pitch for n in c.notes)))

@dataclasses.dataclass
class SearchContext:
  pools: dict = dataclasses.field(default_factory=dict)
  results: dict = dataclasses.field(default_factory=dict)
  evaluations: dict = dataclasses.field(default_factory=collections.OrderedDict)
  diagnostics: dict = dataclasses.field(default_factory=lambda: {'searches': 0, 'cache_hits': 0, 'generated_candidates': 0, 'evaluated_extensions': 0, 'peak_states': 0, 'candidate_families': {}, 'surviving_families': {}})

def note_family(notes):
  return generator.solo.sequences.family(tuple(b.pitch.number() - a.pitch.number() for a, b in zip(notes, notes[1:]))) if notes else 'space'

def state_identity(state, plan):
  _, decisions, last, _, notes = state
  ending = (last.pitch, last.position, last.start, last.end) if last else None
  times = tuple((float(n.start), float(n.end)) for n in notes)
  windows = tuple((float(p.start), float(p.end)) for p in generator.solo.scoring.passages(plan))
  activity = tuple((sum(start <= a < end for a, _ in times), sum(max(0, min(b, end) - max(a, start)) for a, b in times)) for start, end in windows)
  rhythm = tuple((n.start, n.end) for n in notes[-3:])
  progress = tuple(any(p.earliest <= n.start <= p.latest and p.low <= n.pitch.number() <= p.high for n in notes) for p in plan.landmarks)
  return ending, activity, rhythm, progress, tuple(b.pitch.number() - a.pitch.number() for a, b in zip(notes[-4:], notes[-4:][1:])), tuple(d.treatment for d in decisions)

def alternatives(plan, measures, choices, sequences, rhythms, tempo, previous=None, context=None):
  context = context if context is not None else SearchContext()
  key = (plan, previous)
  if key in context.results:
    context.diagnostics['cache_hits'] += 1
    return context.results[key]
  context.diagnostics['searches'] += 1
  if plan not in context.pools:
    context.pools[plan] = tuple((cell, candidates(cell, plan, choices, sequences)) for cell in cells(plan, measures, rhythms))
    for _, pool in context.pools[plan]:
      context.diagnostics['generated_candidates'] += len(pool)
      for candidate in pool:
        family = note_family(candidate.notes)
        counts = context.diagnostics['candidate_families']
        counts[family] = counts.get(family, 0) + 1
  states = [(0.0, (), previous, 0, ())]
  for cell, pool in context.pools[plan]:
    evaluation = generator.solo.scoring.Evaluation(plan, choices, cell.end)
    following, extensions = [], {}
    for _, decisions, last, travel, sounded in states:
      target_met = target_ok(sounded, plan.target)

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

      extension_key = (last.pitch if last else None, target_met)
      if extension_key not in extensions:
        eligible = [c for c in pool if not plan.target or plan.target.latest >= cell.end and cell.end < plan.end or target_met or target_ok(c.notes, plan.target)]
        extensions[extension_key] = generator.solo.sequences.diverse(eligible, extension_rank, lambda c: (len(c.notes), note_family(c.notes), c.treatment if cell.motif else ''), CELL_LIMIT)
      for candidate in extensions[extension_key]:
        variants = [candidate.notes]
        passage = next((p for p in plan.passages if p.start <= cell.start < p.end), None)
        if cell.motif is None and last and last.end == cell.start and len(candidate.notes) == 1 and candidate.notes[0].start == cell.start and candidate.notes[0].pitch == last.pitch and passage and passage.intent == 'sustained_arrival':
          variants.append((dataclasses.replace(candidate.notes[0], continuation=True, role='sustain'),))
        for variant in variants:
          context.diagnostics['evaluated_extensions'] += 1
          placed = generator.solo.placement.place(variant, tempo, last)
          if placed is None:
            continue
          comfort, notes = placed
          decision = generator.solo.types.Decision(cell.start, cell.end, candidate.source, candidate.treatment, notes, candidate.fidelity)
          combined = (*decisions, decision)
          canonical = generator.solo.types.sounding_notes((*sounded, *notes))
          if plan.target and (plan.target.latest < cell.end or cell.end == plan.end) and not target_ok(canonical, plan.target):
            continue
          musical_key = (plan, cell.end, tuple((n.start, n.end, n.pitch, n.role) for n in canonical))
          if musical_key not in context.evaluations:
            context.evaluations[musical_key] = evaluation.features(canonical) if evaluation.feasible(canonical) else None
            if len(context.evaluations) > 32768:
              context.evaluations.popitem(last=False)
          else:
            context.evaluations.move_to_end(musical_key)
          evaluated = context.evaluations[musical_key]
          if evaluated is None:
            continue
          features = dict(evaluated)
          features['connection'] = min(1, abs(canonical[0].pitch.number() - previous.pitch.number()) / 12) if canonical and previous else 0
          features['movement'] = min(1, (travel + comfort) / max(1, len(canonical)) / 2000)
          total = generator.solo.scoring.value(features)
          following.append((total, combined, canonical[-1] if canonical else last, travel + comfort, canonical))
    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)')

    def rank(state):
      return state[0], tuple((n.pitch, n.position, n.start, n.end) for n in state[4])

    unique = {}
    for state in sorted(following, key=rank):
      unique.setdefault(state_identity(state, plan), state)
    states = list(generator.solo.sequences.diverse(unique.values(), rank, lambda s: (note_family(s[4]), len(s[4]), s[4][-1].pitch.number() if s[4] else None), BEAM_WIDTH))
    states.sort(key=rank)
    context.diagnostics['peak_states'] = max(context.diagnostics['peak_states'], len(states))
  selected = generator.solo.sequences.diverse(states, lambda s: (s[0], tuple(n.pitch for n in s[4])), lambda s: (note_family(s[4]), len(s[4]), s[4][-1].pitch if s[4] else None), generator.solo.policy.DEFAULT.alternatives)
  for state in selected:
    family = note_family(state[4])
    counts = context.diagnostics['surviving_families']
    counts[family] = counts.get(family, 0) + 1
  result = tuple(generator.solo.types.Phrase(plan, decisions, score) for score, decisions, _, _, _ in selected)
  context.results[key] = result
  return result

def compose(plan, measures, choices, sequences, rhythms, tempo, previous=None):
  return alternatives(plan, measures, choices, sequences, rhythms, tempo, previous)[0]