from __future__ import annotations

import dataclasses
import fractions
import itertools

import generator.accompaniment
import generator.model
import generator.score

@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:
  """One rhythmic voice; split held pitches into ties rather than release them."""
  originals = {p['id']: p for m in document['staves'][0]['measures'] for v in m['voices'] for n in v['notes'] for p in n['pitches']}
  edges = [(p['id'], p['tie_to']) for p in originals.values() if p.get('tie_to')]
  fragments = {}
  for measure in measures(document):
    length = generator.model.fraction(measure['duration'])
    sounding = [(generator.model.fraction(n['onset']), generator.model.fraction(n['onset']) + generator.model.fraction(p['duration']), p)
      for v in measure['voices'] for n in v['notes'] for p in n['pitches']]
    boundaries = sorted({fractions.Fraction(), length, *(t for a, b, _ in sounding for t in (a, b))})
    notes = []
    for a, b in zip(boundaries, boundaries[1:]):
      active = sorted((p for start, end, p in sounding if start <= a < end), key=lambda p: (generator.model.pitch_number(p), p['tablature']['string']))
      for duration, written in generator.accompaniment.written_durations(b - a):
        pitches = [{k: v for k, v in p.items() if k not in ('tie_from', 'tie_to', 'stem')} for p in active]
        note = generator.model.make_event(f'{measure["id"]}_acoustic_{len(notes) + 1}', a, duration, 'note' if pitches else 'rest', pitches, written)
        note['voice'] = 'voice_arranged'
        for old, new in zip(active, note['pitches']):
          fragments.setdefault(old['id'], []).append(new)
        notes.append(note)
        a += duration
    if not sounding 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}]
  for pieces in fragments.values():
    for a, b in zip(pieces, pieces[1:]):
      a['tie_to'], b['tie_from'] = b['id'], a['id']
  for start, end in edges:
    if start in fragments or end in fragments:
      a = fragments[start][-1] if start in fragments else originals[start]
      b = fragments[end][0] if end in fragments else originals[end]
      a['tie_to'], b['tie_from'] = b['id'], a['id']

def validate_comping(document: dict) -> None:
  """Check complete sounding events, not barline or notation fragments."""
  tones, starts, cursor = {}, {}, fractions.Fraction()
  for measure in document['staves'][0]['measures']:
    for note in generator.score.measure_notes(measure):
      for pitch in note['pitches']:
        tones[pitch['id']] = pitch
        starts[pitch['id']] = cursor + generator.model.fraction(note['onset'])
    cursor += generator.model.fraction(measure['duration'])
  groups = {}
  for identifier, pitch in tones.items():
    if pitch.get('tie_from') in tones:
      continue
    start, tail = starts[identifier], pitch
    seen = {identifier}
    while tail.get('tie_to') in tones:
      target = tail['tie_to']
      if target in seen:
        raise ValueError('cyclic acoustic sustain')
      seen.add(target)
      tail = tones[target]
    end = starts[tail['id']] + generator.model.fraction(tail['duration'])
    string = pitch.get('tablature', {}).get('string')
    if string is None:
      raise ValueError('acoustic comping requires placed notes')
    groups.setdefault((start, end), set()).add('bass' if string >= 5 else 'upper')
  events = []
  for (start, end), lanes in groups.items():
    if len(lanes) == 2:
      if end - start <= 1:
        raise ValueError('combined acoustic sustain must exceed one beat')
    elif end - start not in (fractions.Fraction(1, 2), 1):
      raise ValueError('separate acoustic comping requires quarter/eighth notes')
    events.append((start, end))
  events.sort()
  if any(a[1] > b[0] for a, b in zip(events, events[1:])):
    raise ValueError('acoustic bass and upper must release together or remain separate')
  cursor = fractions.Fraction()
  for measure in document['staves'][0]['measures']:
    length = generator.model.fraction(measure['duration'])
    first = next((event for event in events if cursor <= event[0] < cursor + length), None)
    if first and 'bass' not in groups[first]:
      raise ValueError(f'measure {measure["number"]}: first acoustic attack must include bass')
    cursor += length