from __future__ import annotations

import fractions
import itertools

import generator.accompaniment
import generator.model

def written_durations(length: fractions.Fraction):
  while length:
    written = generator.model.written_for(length)
    duration = length
    if 'ratio' in written and length.denominator & (length.denominator - 1) == 0:
      duration = next((size for size in generator.model.TYPE_LENGTHS.values() if size < length), length)
      written = generator.model.written_for(duration)
    yield duration, written
    length -= duration

def measures(document: dict):
  staff = document['staves'][0]
  if staff.get('arrangement_part_type') == 'acoustic':
    yield from (m for m in staff['measures'] if {v['id'] for v in m['voices']} == {'bass', 'upper'})

def prepare(document: dict) -> None:
  """Release each grip before the next attack; never require overlapping lanes."""
  for measure in measures(document):
    length = generator.model.fraction(measure['duration'])
    attacks = sorted((n for v in measure['voices'] for n in v['notes'] if n['pitches'] or '_harmony' in n), key=lambda n: generator.model.fraction(n['onset']))
    groups = [(time, list(notes)) for time, notes in itertools.groupby(attacks, key=lambda n: generator.model.fraction(n['onset']))]
    for index, (start, notes) in enumerate(groups):
      end = min(start + generator.model.fraction(n['duration']) for n in notes)
      end = min(end, groups[index + 1][0] if index + 1 < len(groups) else length)
      # A duration needing ties becomes one shorter attack followed by silence.
      duration, written = next(written_durations(end - start))
      for note in notes:
        if generator.model.fraction(note['duration']) != duration:
          note.update(duration=generator.model.rational(duration), written=written)
          note.pop('beams', None)
          for pitch in note['pitches']:
            pitch.update(duration=note['duration'], written=written)
    for voice in measure['voices']:
      notes, cursor = [], fractions.Fraction()
      for note in voice['notes']:
        start = generator.model.fraction(note['onset'])
        if start > cursor:
          notes.append(generator.model.make_event(f'{note["id"]}-release', cursor, start - cursor, 'rest') | {'voice': voice['id']})
        notes.append(note)
        cursor = start + generator.model.fraction(note['duration'])
      if cursor < length:
        notes.append(generator.model.make_event(f'{measure["id"]}-{voice["id"]}-release', cursor, length - cursor, 'rest') | {'voice': voice['id']})
      voice['notes'] = notes
      generator.accompaniment.normalize_beams(notes)

def combine(document: dict) -> None:
  """Write prepared bass/upper attacks as one voice with shared stems and rests."""
  for measure in measures(document):
    length = generator.model.fraction(measure['duration'])
    attacks = sorted((n for v in measure['voices'] for n in v['notes'] if n['pitches']), key=lambda n: generator.model.fraction(n['onset']))
    notes, cursor = [], fractions.Fraction()

    def rest(end):
      nonlocal cursor
      for duration, written in written_durations(end - cursor):
        notes.append(generator.model.make_event(f'{measure["id"]}-acoustic-{len(notes) + 1}', cursor, duration, 'rest', written=written) | {'voice': 'voice-arranged'})
        cursor += duration

    for start, group in itertools.groupby(attacks, key=lambda n: generator.model.fraction(n['onset'])):
      rest(start)
      group = list(group)
      pitches = sorted((p for n in group for p in n['pitches']), key=lambda p: (generator.model.pitch_number(p), p['tablature']['string']))
      pitches = [{k: v for k, v in p.items() if k != 'stem'} for p in pitches]
      duration = generator.model.fraction(group[0]['duration'])
      notes.append(generator.model.make_event(f'{measure["id"]}-acoustic-{len(notes) + 1}', start, duration, 'note', pitches, group[0]['written']) | {'voice': 'voice-arranged'})
      cursor = start + duration
    rest(length)
    if not attacks and length == generator.model.meter_length(measure['meter']):
      notes = [generator.model.make_event(f'{measure["id"]}-acoustic-1', fractions.Fraction(), length, 'rest', written={'type': 'whole', 'dots': 0}) | {'voice': 'voice-arranged', 'measure_rest': True}]
    measure['voices'] = [{'id': 'voice-arranged', 'notes': notes}]