"""One learned accompaniment per instrument, reused at authored form positions."""
from __future__ import annotations

import copy
from fractions import Fraction as F

import generator.accompaniment
import generator.acoustic
import generator.acoustic_voicing
import generator.electric
import generator.guitar
import generator.fretboard
import generator.model
import generator.parts
import generator.score
from generator.arrange import form
import generator.arrange.pickup


def harmonic_key(measure, previous):
  def key(chord):
    return repr({k: chord.get(k) for k in ('root', 'kind', 'degrees', 'bass', 'inversion')}) if chord else ''
  active = generator.accompaniment.active_chord(measure, F(), previous)
  return (key(active), tuple((generator.model.fraction(h['onset']), key(h)) for h in measure['harmonies']
    if generator.model.fraction(h['onset']) > 0), measure['duration'], measure['time_signature'])


def voice(source, document, measures, part, initial=None, source_offset=0):
  staff = document['staves'][0]
  result = generator.score.make_arranged_document(source, staff, staff_id=staff['id'],
    staff_name=staff['name'], role='guitar', arrangement_part_type=part, measures=measures)
  if part == 'acoustic':
    generator.acoustic.prepare(result)
    generator.acoustic_voicing.place(result)
    generator.acoustic.combine(result)
  else:
    generator.guitar.place_guitar(result)
  generator.model.validate_document(result)
  generator.fretboard.validate_document(result)
  return result['staves'][0]['measures']


def install(source, document, plans):
  """Freeze complete canonical grips; copies are never placed a second time."""
  from generator.arrange.generate import section_source
  part = document['staves'][0]['arrangement_part_type']
  if part == 'vocal': return
  authored = generator.score.source_staff(source)['measures']
  sections = form.full_arrangement_sections(authored, part)
  if 1 not in plans or not any(s.section_index == 1 for s in sections):
    return []  # A standalone section already is its own canonical part.
  body = generator.score.body_start_measure_index(authored)
  body_measures = copy.deepcopy(authored[body:])
  for m in body_measures:
    for h in m['harmonies']: h.pop('chord_optional', None)
  previous = generator.accompaniment.carried_chords(authored)
  plan = plans[1]

  def build(measures, prior, offset=0):
    carried = generator.accompaniment.carried_chords(measures, prior)
    result = []
    for i, m in enumerate(measures):
      j = offset + i
      rhythm = (plan.acoustic_bass[j], plan.acoustic_upper[j]) if part == 'acoustic' else plan.electric[j]
      result.append(generator.parts.build_measure(m, f'canonical_{part}_{j}', 'accompaniment', part, carried[i], rhythm=rhythm))
    return result

  canonical = voice(source, document, build(body_measures, previous[body]), part, previous[body])
  canonical_previous = generator.accompaniment.carried_chords(body_measures, previous[body])
  keys = [harmonic_key(m, p) for m, p in zip(body_measures, canonical_previous)]
  cache, exceptions = {}, []
  destination = 0
  for si, section in enumerate(sections):
    if section.role != 'accompaniment':
      destination += section.measure_count
      continue
    effective = section_source(authored, section, sections[si + 1] if si + 1 < len(sections) else None)
    carried = generator.accompaniment.carried_chords(effective)
    for j in range(section.measure_count):
      source_index = section.source_measure_index + j
      index = source_index - body
      old = document['staves'][0]['measures'][destination + j]
      key = harmonic_key(effective[source_index], carried[source_index])
      if key == keys[index]:
        template = canonical[index]
      else:
        # Changed pickups/ending harmony reuse the rhythm and only revoice the
        # affected measure. These are chart differences, never solo responses.
        identity = (index, repr(key))
        if identity not in cache:
          cache[identity] = voice(source, document,
            build([effective[source_index]], carried[source_index], index), part, carried[source_index], index)[0]
        template = cache[identity]
        exceptions.append({'measure': old['number'], 'source_measure': authored[source_index]['number'], 'reason': 'changed harmony or arrangement boundary'})
      replacement = copy.deepcopy(old)
      replacement['voices'] = copy.deepcopy(template['voices'])
      replacement['id'] = 'canonical_' + old['id']
      generator.score.reidentify([replacement], [])
      replacement['id'] = old['id']
      following = sections[si + 1] if si + 1 < len(sections) else None
      if j == section.measure_count-1 and following and following.role == 'melody' and generator.arrange.pickup.enters(following, authored):
        sounding = [generator.model.fraction(n['onset']) for n in generator.score.measure_notes(authored[0]) if n['pitches']]
        if sounding:
          cutoff = generator.arrange.pickup.offset(old, authored[0]) + min(sounding)
          prefix = [n for n in generator.score.measure_notes(replacement) if generator.model.fraction(n['onset']) < cutoff]
          for n in prefix: generator.arrange.pickup.truncate(n, cutoff)
          # Preserve original event IDs: incoming melody ties and tuplets refer
          # to these exact pickup events, not to a freshly voiced copy.
          tail = copy.deepcopy([n for n in generator.score.measure_notes(old) if generator.model.fraction(n['onset']) >= cutoff])
          identifier = old['voices'][0]['id']
          for n in prefix: n['voice'] = identifier
          replacement['voices'] = [{'id': identifier, 'notes': prefix + tail}]
          exceptions.append({'measure': old['number'], 'source_measure': authored[source_index]['number'], 'reason': 'protected melody pickup'})
      document['staves'][0]['measures'][destination + j] = replacement
    destination += section.measure_count
  # Build-only diagnostics live outside the published musical schema.
  return exceptions


def audit(source, documents):
  """Compare music, not generated IDs, at every reused form position."""
  authored = generator.score.source_staff(source)['measures']
  result = {}
  for part, document in documents.items():
    if part == 'vocal': continue
    sections = form.full_arrangement_sections(authored, part)
    measures = document['staves'][0]['measures']
    carried = generator.accompaniment.carried_chords(measures)
    seen, exceptions, offset, reused = {}, [], 0, 0
    for si, section in enumerate(sections):
      following = sections[si + 1] if si + 1 < len(sections) else None
      for j in range(section.measure_count):
        if section.role != 'accompaniment': continue
        m = measures[offset+j]
        if j == section.measure_count-1 and following and following.role == 'melody' and generator.arrange.pickup.enters(following, authored):
          exceptions.append({'measure': m['number'], 'reason': 'protected melody pickup'})
          continue
        key = (section.source_measure_index+j, repr(harmonic_key(m, carried[offset+j])))
        signature = tuple((generator.model.fraction(n['onset']), generator.model.fraction(n['duration']),
          tuple((generator.model.pitch_number(p),p['tablature']['string'],p['tablature']['fret']) for p in n['pitches']))
          for n in generator.score.measure_notes(m))
        if key in seen:
          if signature != seen[key]:
            raise ValueError(f'{part}, arrangement measure {m["number"]}: canonical accompaniment changed')
          reused += 1
        else:
          if any(k[0] == key[0] for k in seen):
            exceptions.append({'measure': m['number'], 'reason': 'changed effective harmony or duration'})
          seen[key] = signature
      offset += section.measure_count
    result[part] = {'reused_bars': reused, 'distinct_form_bars': len(seen), 'exceptions': exceptions}
  return result