from __future__ import annotations

import bisect
import copy
import fractions

import generator.accompaniment
import generator.model

def harmonize(document: dict) -> None:
  """Split long melody runs at harmony changes before choosing guitar positions."""
  staff = document['staves'][0]
  tones, events, changes = {}, [], []
  absolute = fractions.Fraction()
  previous = None
  for measure in staff['measures']:
    for chord in measure['harmonies']:
      value = chord if chord['kind'] != 'none' else None
      if value != previous and not generator.accompaniment.same_chord(value, previous):
        time = absolute + generator.model.fraction(chord['onset'])
        if changes and changes[-1][0] == time:
          changes[-1] = (time, value)
        else:
          changes.append((time, value))
      previous = value
    for voice in measure['voices']:
      for note in voice['notes']:
        start = absolute + generator.model.fraction(note['onset'])
        events.append((note, start))
        for tone in note['pitches']:
          tones[tone['id']] = tone
    absolute += generator.model.fraction(measure['duration'])
  eligible = set()
  visited = set()
  for note, _ in events:
    if not note.get('_melody') or note.get('grace') or note.get('cue') or len(note['pitches']) != 1 or note['pitches'][0].get('cue'):
      continue
    tone = note['pitches'][0]
    if tone['id'] in visited:
      continue
    while tone.get('tie_from'):
      tone = tones[tone['tie_from']]
    chain = []
    while True:
      chain.append(tone)
      if not tone.get('tie_to'):
        break
      tone = tones[tone['tie_to']]
    visited.update(p['id'] for p in chain)
    if sum((generator.model.fraction(p['duration']) for p in chain), fractions.Fraction()) > 1:
      eligible.update(p['id'] for p in chain)
  times = [time for time, _ in changes]
  reattacks = {time for time, chord in changes if chord}
  replacements, tone_pieces = {}, {}
  for note, start in events:
    melodic = note.pop('_melody', False)
    if not melodic or not note['pitches'] or note['pitches'][0]['id'] not in eligible:
      continue
    end = start + generator.model.fraction(note['duration'])
    cuts = [start, *times[bisect.bisect_right(times, start):bisect.bisect_left(times, end)], end]
    pieces = []
    for index, (a, b) in enumerate(zip(cuts, cuts[1:])):
      fragment = copy.deepcopy(note)
      fragment['id'] = note['id'] if index == 0 else f'{note["id"]}-h{index}'
      fragment['onset'] = generator.model.rational(generator.model.fraction(note['onset']) + a - start)
      fragment['duration'] = generator.model.rational(b - a)
      if len(cuts) > 2:
        fragment['written'] = generator.model.written_for(b - a)
        fragment.pop('beams', None)
      for tone in fragment['pitches']:
        original = tone['id']
        tone['id'] = original if index == 0 else f'{original}-h{index}'
        tone['duration'], tone['written'] = fragment['duration'], copy.deepcopy(fragment['written'])
        tone.pop('tie_from', None)
        tone.pop('tie_to', None)
        tone_pieces.setdefault(original, []).append((tone, a))
      if index:
        fragment.pop('lyrics', None)
      current = bisect.bisect_right(times, a) - 1
      chord = changes[current][1] if current >= 0 else None
      if chord:
        fragment.update(_harmony=chord, _guitar_profile='electric', _harmonize=True)
      pieces.append(fragment)
    replacements[note['id']] = pieces
  for identifier, pieces in tone_pieces.items():
    original = tones[identifier]
    links = list(zip(pieces, pieces[1:]))
    if source := original.get('tie_from'):
      previous_piece = tone_pieces[source][-1] if source in tone_pieces else (tones[source], None)
      links.insert(0, (previous_piece, pieces[0]))
      previous_piece[0].pop('tie_to', None)
    if target := original.get('tie_to'):
      if target not in tone_pieces:
        pieces[-1][0]['tie_to'] = target
        tones[target]['tie_from'] = pieces[-1][0]['id']
    for (first, _), (second, start) in links:
      if start not in reattacks:
        first['tie_to'], second['tie_from'] = second['id'], first['id']
  by_tone = {p['id']: note for original, _ in events for note in replacements.get(original['id'], [original]) for p in note['pitches']}
  for measure in staff['measures']:
    for voice in measure['voices']:
      voice['notes'] = [piece for note in voice['notes'] for piece in replacements.get(note['id'], [note])]
      for note in voice['notes']:
        if note.get('_harmonize') and (source := note['pitches'][0].get('tie_from')):
          earlier = by_tone[source]
          if generator.accompaniment.same_chord(earlier.get('_harmony'), note['_harmony']):
            note['_rhythm_continuation'] = True
      generator.accompaniment.normalize_beams(voice['notes'])
  for tuplet in staff.get('tuplets', []):
    if tuplet['end'] in replacements:
      tuplet['end'] = replacements[tuplet['end']][-1]['id']