from __future__ import annotations

import dataclasses
import fractions
import itertools

import generator.accompaniment
import generator.model

@dataclasses.dataclass
class HarmonySpan:
  harmony: dict
  location: str
  notes: list[dict] = dataclasses.field(default_factory=list)
  iteration: int = 0

def harmony_spans(document: dict) -> list[HarmonySpan]:
  """Group attacks by continuous harmony within each arrangement iteration."""
  spans, span, active = [], None, None
  iteration = 0
  for measure in document['staves'][0]['measures']:
    if measure.get('section_label'):
      iteration += 1
      span = None
    if {v['id'] for v in measure['voices']} != {'bass', 'upper'}:
      span, active = None, None
      continue
    events = [(generator.model.fraction(h['onset']), 0, h) for h in measure['harmonies']]
    events.extend((generator.model.fraction(n['onset']), 1, n) for v in measure['voices'] for n in v['notes'] if n.get('_guitar_profile') == 'acoustic')
    for onset, kind, item in sorted(events, key=lambda event: event[:2]):
      harmony = item if kind == 0 else item['_harmony']
      if not generator.accompaniment.same_chord(active, harmony):
        span = None
      active = harmony
      if kind == 0:
        continue
      if span is None:
        span = HarmonySpan(harmony, f'measure {measure["number"]}, event {item["id"]}, onset {onset}', iteration=iteration)
        spans.append(span)
      span.notes.append(item)
  return spans

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(generator.accompaniment.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 generator.accompaniment.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}]