"""Source-only sparse accompaniment, frozen once at each form position."""
from __future__ import annotations
import copy
import functools
import json
from fractions import Fraction as F

from generator import accompaniment, fretboard, guitar, harmony, model, parts, score, vocabulary
from generator.solo.policy import stable_rank


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)


@functools.lru_cache(maxsize=2048)
def voicings(key, bass):
  chord = json.loads(key)
  degrees = harmony.degrees(chord)
  root = int(model.NATURAL[chord['root']['letter']] + model.alteration(chord['root']['alter']))
  classes = {d: (root + harmony.degree_semitones(d, a)) % 12 for d, a in degrees.items()}
  required = {p for d, p in classes.items() if d in (2, 3, 4, 6, 7) or d == 5 and degrees[d] or d > 7}
  full = fretboard.chord_requirements(chord)
  required.update(full.basses)
  if bass and 1 in classes:
    required.add(classes[1])
  if len(required) < 2:
    required.update(classes[d] for d in (1, 5) if d in classes)
  allowed = set(full.allowed)
  requirements = fretboard.ChordRequirements(tuple(sorted(required)), tuple(sorted(allowed)), tuple(sorted(allowed)), full.basses)
  options = guitar.chord_positions(requirements, prune=False)
  size = max(2, len(required))
  options = [v for v in options if len(v) == size and fretboard.playable_attack((p, F(1)) for p in v)]
  if not options:
    raise ValueError(harmony.symbol(chord) + ': no legal sparse voicing preserving essential written degrees')
  # Retain distinct hand regions and inversions, independent of enumeration.
  best = {}
  for v in options:
    family = (min(p.fret for p in v), min(p.midi for p in v) % 12)
    rank = (guitar.compactness(v), stable_rank('comp-voicing', v))
    if family not in best or rank < best[family][0]:
      best[family] = rank, v
  return tuple(v for _, v in sorted(best.values()))


class Canonical:
  def __init__(self, document, part, seed):
    self.part, self.seed = part, seed
    self.policy = vocabulary.role(part + '_accompaniment')
    self.cache, self.previous, self.report = {}, (), []
    measures = score.source_staff(document)['measures']
    self.first_index = score.body_start_measure_index(measures)
    self.last_index = len(measures)-1
    carried = accompaniment.carried_chords(measures)
    for index in range(score.body_start_measure_index(measures), len(measures)):
      source = copy.deepcopy(measures[index])
      for h in source['harmonies']:
        h.pop('chord_optional', None)
      self(source, carried[index], index)

  def __call__(self, source, prior, index):
    key = (index, tuple((model.fraction(h['onset']), semantic_key(h)) for h in source['harmonies']), semantic_key(prior) if prior else '', source['duration'].__repr__())
    if key in self.cache:
      return copy.deepcopy(self.cache[key])
    length = model.fraction(source['duration'])
    boundaries = sorted({F(), length, *(model.fraction(h['onset']) for h in source['harmonies'])})
    attacks = []
    source_attacks = {model.fraction(n['onset']) for n in score.measure_notes(source) if n['pitches'] and not all('tie_from' in p for p in n['pitches'])}
    for start, end in zip(boundaries, boundaries[1:]):
      chord = accompaniment.active_chord(source, start, prior)
      if not chord:
        continue
      duration = end - start
      # Essential bass is a source-derived form responsibility, independent
      # of which other guitar happens to be playing or what it generated.
      bass = self.policy.get('bass') == 'primary' or bool(chord.get('bass')) or index in (self.first_index,self.last_index) or start > 0
      options = voicings(semantic_key(chord), bass)
      center = self.policy.get('center', 55)
      selected = min(options, key=lambda v: (
        .08 * abs(sum(p.midi for p in v) / len(v) - center)
        + (.15 * guitar.movement(self.previous, v) if self.previous else 0)
        + .1 * guitar.compactness(v), stable_rank(self.seed, (index, start, v))))
      self.previous = selected
      patterns = []
      for figure in vocabulary.data('rhythms')['comping']:
        events = tuple((start+F(a), min(end, start+F(a)+F(d))) for a, d in figure['attacks'] if F(a) < duration)
        # Space is available on held harmony, never an accidental empty winner.
        if not events and (not prior or not accompaniment.same_chord(chord, prior) or start):
          continue
        conflicts = sum(a in source_attacks for a, _ in events)
        target = 1 if duration <= 3 else 1.5
        cost = .2*abs(len(events)-target) + .08*conflicts
        if figure['id'] == 'sustain' and duration < 3:
          cost += .2
        patterns.append((cost, stable_rank(self.seed, (index, start, figure['id'])), events, figure['id']))
      # Source-specific variation within a broad acceptability band, not a loop.
      low = min(p[0] for p in patterns)
      _, _, events, label = min((p for p in patterns if p[0] <= low+.25), key=lambda p:p[1])
      for a, b in events:
        attacks.append((a, b, selected, chord))
      self.report.append({'source_index': index, 'onset': str(start), 'chord': harmony.symbol(chord), 'figure': label, 'tones': len(selected), 'bass': bass,
        'fuller_voicing_reason': 'explicit written degrees/bass' if len(selected)>3 else ''})
    notes, cursor = [], F()
    def emit(a, b, positions=(), chord=None):
      nonlocal cursor
      previous = None
      for duration, written in accompaniment.written_durations(b-a):
        pitches = fretboard.chord_pitches(chord, positions) if positions else []
        n = model.make_event(f'comp_{self.part}_{index}_{len(notes)}', a, duration, 'note' if pitches else 'rest', pitches, written)
        if previous:
          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 start, end, positions, chord in attacks:
      if cursor < start: emit(cursor, start)
      emit(start, end, positions, chord)
    if cursor < length: emit(cursor, length)
    result = parts.build_measure(source, 'canonical', 'rest', self.part, prior)
    result['voices'] = [{'id': 'voice_arranged', 'notes': notes}]
    for note in notes: note['voice'] = 'voice_arranged'
    self.cache[key] = copy.deepcopy(result)
    return result