"""Source-only accompaniment: joint layered rhythms and connected grips."""
from __future__ import annotations
import copy
import dataclasses
import functools
import json
import fractions
import typing

import generator.accompaniment
import generator.comping_rhythm
import generator.chord_families
import generator.acoustic_voicing
import generator.arrange.form
import generator.fretboard
import generator.guitar
import generator.harmony
import generator.model
import generator.parts
import generator.score
import generator.vocabulary
import generator.solo.policy

def semantic_key(chord):
  return json.dumps({k: chord[k] for k in ('root', 'kind', 'degrees', 'bass', 'inversion') if k in chord}, sort_keys=True)

def requirements(chord, part):
  result = generator.fretboard.chord_requirements(chord, respect_bass=part != 'electric')
  if part == 'electric' or result.basses:
    return result
  degrees = generator.harmony.degrees(chord)
  plain = degrees in ({1: 0, 3: 0, 5: 0}, {1: 0, 3: -1, 5: 0})
  root = int(generator.model.NATURAL[chord['root']['letter']] + generator.model.alteration(chord['root']['alter']))
  basses = tuple(sorted({(root + generator.harmony.degree_semitones(d, degrees[d])) % 12 for d in ((1, 3, 5) if plain else (1, 5)) if d in degrees}))
  if not basses:
    raise ValueError('acoustic accompaniment requires a written bass, root or fifth')
  return dataclasses.replace(result, basses=basses)

@functools.lru_cache(maxsize=2048)
def voicings(key, part):
  chord = json.loads(key)
  required = requirements(chord, part)
  if part == 'electric':
    options = generator.guitar.chord_positions(required, prune=False, profile='electric')
  elif part == 'acoustic':
    degrees = generator.harmony.degrees(chord)
    plain = degrees in ({1: 0, 3: 0, 5: 0}, {1: 0, 3: -1, 5: 0})
    options = tuple(v for v in generator.acoustic_voicing.candidates(required, plain)
      if v.bass.fret <= generator.fretboard.MAX_SINGLE_FRET
      and generator.fretboard.has_fret_anchor(v[1:])
      and (not plain or set(required.allowed) <= {p.midi % 12 for p in v}))
  else:
    raise ValueError('Unknown accompaniment instrument: ' + part)
  options = tuple(v for v in options if generator.fretboard.playable_attack((p, fractions.Fraction(1)) for p in v))
  if not options:
    layout = 'bass on 5/6 plus upper strings 2/3/4' if part == 'acoustic' else 'melody chord families on strings 1–4'
    raise ValueError(f'{generator.harmony.symbol(chord)}: no legal {part} voicing ({layout}); required degrees/bass, no open strings, fret anchor 1–12 and at most {generator.fretboard.MAX_FRET_SPAN} consecutive frets remain binding')
  return tuple(sorted(options))

def movement(part, previous, current):
  if not previous or not current:
    return 0
  if part == 'acoustic':
    return generator.acoustic_voicing.movement(previous, current)
  return generator.guitar.movement(previous, current)

class VoicingCost(typing.NamedTuple):
  extended_grips: int = 0
  family: int = 0
  effort: int = 0
  missing_tones: int = 0
  register: fractions.Fraction = fractions.Fraction()
  def plus(self, other):
    return VoicingCost(*(a + b for a, b in zip(self, other)))
  def moved(self, distance):
    return self._replace(effort=self.effort + distance)

def local_cost(part, key, grip, center):
  required = requirements(json.loads(key), part)
  register = abs(fractions.Fraction(sum(p.midi for p in grip), len(grip)) - center)
  if part == 'electric':
    value = generator.guitar.voicing_cost(grip, 'electric', required.allowed)
    return VoicingCost(value.extended_grips, 3 * value.string_shortfall + value.family_penalty, value.effort, value.missing_tones, register)
  # Familiar degree orders on either bass string, then whole-form movement.
  family = generator.chord_families.acoustic(json.loads(key), grip)
  return VoicingCost(generator.fretboard.span_penalty(grip), int(not family['preferred']), 0, len(set(required.allowed) - {p.midi % 12 for p in grip}), register)

def select_grips(layers, costs, part, seed, *, left=(), right=(), close=False):
  """Exact small-state dynamic programming; never prune by hand center.

  Electric effort includes its melody grip cost and hand-center travel.
  Acoustic effort follows all four voices, including a bass-string change.
  Fixing the first grip permits an exact repeat-join cost without a beam.
  """
  if not layers:
    return (), VoicingCost()
  # Every accompaniment transition is feasible (travel is a cost). A wider
  # option cannot beat a primary grip at the same span, regardless of travel.
  # This exact dominance reduction avoids growing the cubic repeat search.
  retained = []
  for i, layer in enumerate(layers):
    tier = min(costs[i][v].extended_grips for v in layer)
    retained.append(tuple(v for v in layer if costs[i][v].extended_grips == tier))
  layers = tuple(retained)
  ranks = [{v: generator.solo.policy.stable_rank(seed, (i, v)) for v in layer} for i, layer in enumerate(layers)]
  winners = []
  for first in layers[0] if close else [None]:
    states = {left: (VoicingCost(), (), ())}
    for i, layer in enumerate(layers):
      following = {}
      for grip in [first] if i == 0 and first is not None else layer:
        local = costs[i][grip]
        choices = []
        for previous, (cost, ties, path) in states.items():
          total = cost.plus(local).moved(movement(part, previous, grip))
          choices.append((total, (*ties, ranks[i][grip]), (*path, grip)))
        following[grip] = min(choices)
      states = following
    for grip, (cost, ties, path) in states.items():
      travel = movement(part, grip, first if close else right)
      winners.append((cost.moved(travel), ties, path))
  cost, _, path = min(winners)
  return path, cost

@dataclasses.dataclass(frozen=True)
class HarmonyChoice:
  index: int
  start: fractions.Fraction
  end: fractions.Fraction
  chord: dict
  events: tuple
  figure: dict
  options: tuple
  @property
  def key(self):
    return (self.index, self.start, self.end, semantic_key(self.chord), self.events)

class Canonical:
  def __init__(self, document, part, seed):
    self.part, self.seed = part, seed
    self.policy = generator.vocabulary.role(part + '_accompaniment')
    self.cache, self._reports, self._rhythm_reports = {}, {}, {}
    measures = generator.score.source_staff(document)['measures']
    self.melody = generator.comping_rhythm.melody_contexts(document)
    first_index = generator.score.body_start_measure_index(measures)
    sections = generator.arrange.form.full_arrangement_sections(measures, part)
    self.closed = any(a.role == b.role == 'accompaniment' and a.source_measure_index + a.measure_count == len(measures)
      and b.source_measure_index == first_index for a, b in zip(sections, sections[1:]))
    carried = generator.accompaniment.carried_chords(measures)
    inputs = []
    for index in range(first_index, len(measures)):
      source = copy.deepcopy(measures[index])
      for chord in source['harmonies']:
        chord.pop('chord_optional', None)
      inputs.append((source, carried[index], index))
    layers = tuple(generator.comping_rhythm.candidates(m, prior, part, self.melody[i], seed + '/' + str(i)) for m, prior, i in inputs)
    rhythms = generator.comping_rhythm.select(layers, seed)
    self.rhythms = {i: r for (_, _, i), r in zip(inputs, rhythms)}
    prepared = [(m, prior, i, self.prepare(m, prior, i, r)) for (m, prior, i), r in zip(inputs, rhythms)]
    active = tuple(s for _, _, _, spans in prepared for s in spans if s.events)
    path, self.path_cost = self.select(active, close=self.closed)
    self.base = dict(zip((s.key for s in active), path))
    self.order = tuple((s.index, s.start, v) for s, v in zip(active, path))
    previous = path[-1] if self.closed and path else ()
    for source, prior, index, spans in prepared:
      selected = tuple(self.base[s.key] for s in spans if s.events)
      self.cache[self.key(source, prior, index)] = self.render(source, prior, index, spans, selected, previous, False)
      if selected:
        previous = selected[-1]
    if generator.vocabulary.accompaniment(part)['connect_holds']:
      self.connect_holds(inputs)
  @property
  def report(self):
    return [row for key in sorted(self._reports, key=lambda k: (k[0], repr(k))) for row in self._reports[key]]
  @property
  def rhythm_report(self):
    return [self._rhythm_reports[key] for key in sorted(self._rhythm_reports, key=lambda k: (k[0], repr(k)))]
  @staticmethod
  def key(source, prior, index):
    return (index, tuple((a, b, semantic_key(c) if c else '') for a, b, c in generator.comping_rhythm.spans(source, prior)),
      source['duration'].__repr__())
  def prepare(self, source, prior, index, rhythm):
    spans = []
    for start, end, chord in generator.comping_rhythm.spans(source, prior):
      if not chord:
        continue
      events = tuple(e for e in rhythm.events if start <= e.start < end)
      try:
        options = voicings(semantic_key(chord), self.part)
      except ValueError as error:
        raise ValueError(f'source measure {source["number"]}, beat {start + 1}: {error}') from error
      spans.append(HarmonyChoice(index, start, end, chord, events, dict(rhythm.families), options))
    self._rhythm_reports[self.key(source, prior, index)] = {
      'source_index': index, 'families': dict(rhythm.families), 'source_melody': [(str(a), str(b)) for a, b in self.melody[index]],
      'contributions': dict(rhythm.contributions), 'selection_cost': rhythm.cost, 'exceptions': rhythm.exceptions,
      'layers': {lane: {'attacks': sum(e.lane == lane for e in rhythm.events),
        'events': [{'onset': str(e.start), 'duration': str(e.end - e.start)} for e in rhythm.events if e.lane == lane], 'tied_continuations': 0} for lane, _ in rhythm.families},
      'overlap_beats': str(sum((max(0, min(a.end, b.end) - max(a.start, b.start)) for a in rhythm.events for b in rhythm.events if a.lane == 'bass' and b.lane == 'upper'), fractions.Fraction()))}
    return tuple(spans)
  def select(self, spans, *, left=(), right=(), close=False, fixed=None):
    layers = tuple((fixed[s.key],) if fixed and s.key in fixed else s.options for s in spans)
    costs = tuple({v: local_cost(self.part, semantic_key(s.chord), v, self.policy.get('center', 55)) for v in layer}
      for s, layer in zip(spans, layers))
    return select_grips(layers, costs, self.part, self.seed, left=left, right=right, close=close)
  def __call__(self, source, prior, index):
    key = self.key(source, prior, index)
    if key not in self.cache:
      # Effective pickup/turnaround harmony may differ from the main form.
      # Solve that measure between fixed canonical neighbors. Request order
      # and previously generated arrangement sections cannot change it.
      choices = generator.comping_rhythm.candidates(source, prior, self.part, self.melody[index], self.seed + '/' + str(index))
      rhythm = generator.comping_rhythm.select((choices,), self.seed, self.rhythms.get(index - 1), self.rhythms.get(index + 1))[0]
      spans = self.prepare(source, prior, index, rhythm)
      before = [v for i, _, v in self.order if i < index]
      after = [v for i, _, v in self.order if i > index]
      left = before[-1] if before else self.order[-1][2] if self.closed and self.order else ()
      right = after[0] if after else self.order[0][2] if self.closed and self.order else ()
      selected, _ = self.select(tuple(s for s in spans if s.events), left=left, right=right, fixed=self.base)
      self.cache[key] = self.render(source, prior, index, spans, selected, left, True)
    return copy.deepcopy(self.cache[key])
  def render(self, source, prior, index, spans, selected, previous, variant):
    length = generator.model.fraction(source['duration'])
    attacks, rows = [], []
    choices = iter(selected)
    for span in spans:
      grip = next(choices) if span.events else ()
      bass = grip[0] if grip and self.part == 'acoustic' else None
      degrees = generator.harmony.degrees(span.chord)
      root = int(generator.model.NATURAL[span.chord['root']['letter']] + generator.model.alteration(span.chord['root']['alter']))
      rows.append({'source_index': index, 'onset': str(span.start), 'chord': generator.harmony.symbol(span.chord),
        'figure': span.figure, 'tones': len(grip), 'bass': bool(bass), 'variant': variant,
        'family': (generator.chord_families.acoustic(span.chord, grip) if self.part == 'acoustic' else generator.chord_families.electric(grip)) if grip else None,
        'family_coverage': generator.chord_families.coverage(span.chord, span.options, self.part),
        'fret_span': generator.fretboard.fret_span(grip),
        'span_group': ('fallback' if generator.fretboard.span_penalty(grip) else 'primary') if grip else None,
        'span_candidates': generator.chord_families.span_coverage(span.options),
        'positions': [{'string': p.string_number, 'fret': p.fret, 'sounding_midi': p.midi} for p in grip],
        'bass_degree': next((d for d, a in degrees.items() if bass and (root + generator.harmony.degree_semitones(d, a)) % 12 == bass.midi % 12), None),
        'bass_string': bass.string_number if bass else None,
        'movement': movement(self.part, previous, grip), 'voicing_candidates': len(span.options)})
      for event in span.events:
        positions = (grip[0],) if event.lane == 'bass' else grip[1:] if event.lane == 'upper' else grip
        attacks.append((event.lane, event.start, event.end, positions, span.chord))
      if grip:
        previous = grip
    self._reports[self.key(source, prior, index)] = rows
    voices = []
    notes, cursor, lane = [], fractions.Fraction(), ''
    def emit(a, b, positions=(), chord=None):
      nonlocal cursor
      previous = None
      for duration, written in generator.accompaniment.written_durations(b - a):
        pitches = generator.fretboard.chord_pitches(chord, positions) if positions else []
        n = generator.model.make_event(f'comp_{self.part}_{index}_{lane}_{len(notes)}', a, duration, 'note' if pitches else 'rest', pitches, written)
        n['voice'] = lane
        if self.part == 'acoustic':
          n['stem'] = 'down' if lane == 'bass' else 'up'
        if previous and pitches:
          for old, new in zip(previous['pitches'], n['pitches']):
            old['tie_to'], new['tie_from'] = new['id'], old['id']
        notes.append(n)
        previous = n
        a += duration
      cursor = b
    for lane in ['upper', 'bass'] if self.part == 'acoustic' else ['electric']:
      notes, cursor = [], fractions.Fraction()
      for _, start, end, positions, chord in sorted((e for e in attacks if e[0] == lane), key=lambda e: e[1]):
        if cursor < start:
          emit(cursor, start)
        emit(start, end, positions, chord)
      if cursor < length:
        emit(cursor, length)
      voices.append({'id': lane, 'notes': notes})
    result = generator.parts.build_measure(source, 'canonical', 'rest', self.part, prior)
    result['voices'] = voices
    return result
  def connect_holds(self, inputs):
    """Only explicit sustained neighbors can share an attack across a barline.

    These are canonical internal joins, never links across role boundaries.
    The direct reattack remains when a grip or required attack count differs.
    """
    for left, right in zip(inputs, inputs[1:]):
      before = self.cache[self.key(*left)]
      after = self.cache[self.key(*right)]
      previous_chord = generator.accompaniment.active_chord(left[0], generator.model.fraction(left[0]['duration']), left[1])
      next_chord = generator.accompaniment.active_chord(right[0], fractions.Fraction(), right[1])
      if not generator.accompaniment.same_chord(previous_chord, next_chord):
        continue
      for a, b in zip(before['voices'], after['voices']):
        if a['id'] == 'bass':
          continue
        old, new = a['notes'][-1], b['notes'][0]
        minimum = 1 if a['id'] == 'upper' else 0
        attacks = sum(bool(n['pitches']) and not any('tie_from' in p for p in n['pitches']) for n in b['notes'])
        if not old['pitches'] or not new['pitches'] or attacks <= minimum or min(generator.model.fraction(n['duration']) for n in (old, new)) < 1:
          continue
        if [p['tablature'] for p in old['pitches']] != [p['tablature'] for p in new['pitches']]:
          continue
        for p, q in zip(old['pitches'], new['pitches']):
          p['tie_to'], q['tie_from'] = q['id'], p['id']
        row = self._rhythm_reports[self.key(*right)]['layers'][a['id']]
        row['attacks'] -= 1
        row['tied_continuations'] += 1