from __future__ import annotations

import dataclasses
import fractions
import itertools

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

@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 sounding_lanes(measures):
  """Recover complete string-lane events, preserving ties versus reattacks."""
  tones, starts, cursor = {}, {}, fractions.Fraction()
  for measure in 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:
      if tones[pitch['tie_from']].get('tie_to') != identifier:
        raise ValueError('broken acoustic sustain link')
      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')
      continuation = tones[target]
      if continuation.get('tie_from') != tail['id'] or starts[target] != starts[tail['id']] + generator.model.fraction(tail['duration']):
        raise ValueError('broken acoustic sustain link or timing')
      if generator.model.pitch_number(continuation) != generator.model.pitch_number(tail) or continuation.get('tablature') != tail.get('tablature'):
        raise ValueError('acoustic sustain changes pitch or position')
      seen.add(target)
      tail = continuation
    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')
  return groups

def validate_assignment(measures, events):
  expected = {}
  for e in events:
    lanes = ['bass', 'upper'] if e.lane == 'combined' else [e.lane]
    expected.setdefault((e.start, e.end), set()).update(lanes)
  if sounding_lanes(measures) != expected:
    raise ValueError('rendered acoustic attacks disagree with the selected bass/upper assignment')

def validate_comping(document: dict) -> None:
  """Check complete sounding events, not barline or notation fragments."""
  groups = sounding_lanes(document['staves'][0]['measures'])
  events = []
  for (start, end), lanes in groups.items():
    if end <= start or (end - start < 1 and end - start != fractions.Fraction(1, 2)):
      raise ValueError('acoustic comping requires quarter/eighth notes or authored sustains')
    events.append((start, end))
  events.sort()
  # A combined attack is one interval with both lanes. Separate attacks must
  # be sequential, including bass/upper alternation and cross-bar ties.
  for i, a in enumerate(events):
    for b in events[i + 1:]:
      if a[1] > b[0]:
        raise ValueError('acoustic sound must release before the next attack')
  spans = []
  for span in generator.harmonic_time.read(document['staves'][0]):
    if spans and generator.accompaniment.same_chord(spans[-1].chord, span.chord):
      spans[-1] = dataclasses.replace(spans[-1], end=span.end)
    else:
      spans.append(span)
  for span in spans:
    if span.chord is None:
      continue
    first = next((event for event in events if span.start <= event[0] < span.end), None)
    # A bounded anticipated bass can already be sounding into this harmony.
    anticipated = any(span.start - fractions.Fraction(1, 2) <= a < span.start < b and 'bass' in groups[a, b] for a, b in events)
    if first and not anticipated and not any('bass' in groups[e] for e in events if e[0] == first[0]):
      raise ValueError(f'measure {span.source_measure}, beat {span.beat}: first acoustic attack for a chord must include bass')