from __future__ import annotations

import dataclasses
import fractions
import functools

import generator.model
import generator.rhythm

ATTACKS_PER_BEAT = {'bass': fractions.Fraction(1, 2), 'upper': fractions.Fraction(3, 4), 'electric': fractions.Fraction(1, 2)}

@dataclasses.dataclass(frozen=True)
class Context:
  measure: dict
  previous: dict | None
  offset: fractions.Fraction
  intervals: tuple[tuple[fractions.Fraction, fractions.Fraction], ...]
  attacks: frozenset[fractions.Fraction]
  occupied: tuple[tuple[fractions.Fraction, fractions.Fraction], ...]

@dataclasses.dataclass(frozen=True)
class Plan:
  acoustic_bass: tuple[generator.rhythm.Bar, ...]
  acoustic_upper: tuple[generator.rhythm.Bar, ...]
  electric: tuple[generator.rhythm.Bar, ...]

def active_chord(measure: dict, onset: fractions.Fraction, previous: dict | None) -> dict | None:
  active = previous
  for chord in measure['harmonies']:
    if generator.model.fraction(chord['onset']) > onset:
      break
    active = chord
  return active if active and active['kind'] != 'none' else None

def carried_chords(measures: list[dict], initial: dict | None = None) -> list[dict | None]:
  result = []
  for measure in measures:
    result.append(initial)
    if measure['harmonies']:
      initial = measure['harmonies'][-1]
  return result

def context(measure: dict, previous: dict | None) -> Context:
  length = generator.model.fraction(measure['duration'])
  nominal = generator.model.meter_length(measure['meter'])
  offset = nominal - length if measure.get('implicit') and length < nominal else fractions.Fraction()
  boundaries = sorted({fractions.Fraction(), length} | {generator.model.fraction(h['onset']) for h in measure['harmonies']})
  intervals = tuple((a, b) for a, b in zip(boundaries, boundaries[1:]) if b > a and active_chord(measure, a, previous))
  attacks, occupied = set(), []
  for voice in measure['voices']:
    for note in voice['notes']:
      pitches = [p for p in note['pitches'] if not p.get('cue')]
      if not pitches or note.get('cue'):
        continue
      start = generator.model.fraction(note['onset'])
      end = start + generator.model.fraction(note['duration'])
      if not all('tie_from' in p for p in pitches):
        attacks.add(start)
      occupied.append((start, end))
  return Context(measure, previous, offset, intervals, frozenset(attacks), tuple(occupied))

@functools.lru_cache(maxsize=4096)
def sounding_rhythm(bar, offset, intervals):
  attacks = tuple(t - offset for t in bar.attacks() if any(a <= t - offset < b for a, b in intervals))
  spans = tuple((max(event.onset - offset, a), min(event.onset + event.duration - offset, b)) for event in bar.events if event.sounding and not event.continuation and event.onset + event.duration > offset for a, b in intervals if a <= max(event.onset - offset, 0) < b)
  return attacks, spans

def sounding_attacks(bar: generator.rhythm.Bar, item: Context) -> tuple[fractions.Fraction, ...]:
  return sounding_rhythm(bar, item.offset, item.intervals)[0]

def sounding_spans(bar: generator.rhythm.Bar, item: Context) -> tuple[tuple[fractions.Fraction, fractions.Fraction], ...]:
  return sounding_rhythm(bar, item.offset, item.intervals)[1]

@functools.lru_cache(maxsize=16384)
def measure_rank(bar, offset, intervals, melody_attacks, melody_spans, companions, role):
  attacks, spans = sounding_rhythm(bar, offset, intervals)
  missing_bars = int(bool(intervals) and not spans and role != 'electric')
  uncovered = sum(not any(start < b and end > a for start, end in spans) for a, b in intervals)
  duration = sum((b - a for a, b in intervals), fractions.Fraction())
  density = abs(len(attacks) - duration * ATTACKS_PER_BEAT[role])
  collisions = sum((t in melody_attacks) + (t in companions) for t in attacks)
  occupied = sum(any(a <= t < b for a, b in melody_spans) for t in attacks)
  pulse = -sum(2 if t == 0 else 1 if t.denominator == 1 else 0 for t in attacks)
  if role == 'electric':
    return missing_bars, density, collisions, occupied, uncovered, pulse
  return missing_bars, uncovered, density, pulse if role == 'bass' else collisions, collisions if role == 'bass' else pulse, occupied

def rank(bar: generator.rhythm.Bar, contexts: list[tuple[Context, tuple[fractions.Fraction, ...]]], role: str) -> tuple:
  result = (0, 0, 0, 0, 0, 0)
  for item, companions in contexts:
    result = add_ranks(result, measure_rank(bar, item.offset, item.intervals, item.attacks, item.occupied, companions, role))
  return result

def add_ranks(left: tuple, right: tuple) -> tuple:
  return tuple(a + b for a, b in zip(left, right))

def cycle(snippets: tuple[generator.rhythm.Snippet, ...], contexts: list[Context], role: str, companions: list[tuple[fractions.Fraction, ...]] | None = None) -> tuple[generator.rhythm.Bar, ...]:
  bars = {}
  for snippet in snippets:
    for index, bar in enumerate(snippet.bars):
      bars.setdefault(bar, f'{snippet.id}:{index}')
  phases = [[(item, companions[i] if companions else ()) for i, item in enumerate(contexts) if i % 2 == phase] for phase in range(2)]
  scores = [{bar: rank(bar, phase, role) for bar in bars} for phase in phases]
  candidates = [(add_ranks(scores[0][s.bars[0]], scores[1][s.bars[1]]), s.id, s.bars) for s in snippets if len(s.bars) == 2]
  single = [s for s in snippets if len(s.bars) == 1]
  if single:
    pair = tuple(min(single, key=lambda s: (scores[phase][s.bars[0]], s.id)) for phase in range(2))
    candidates.append((add_ranks(scores[0][pair[0].bars[0]], scores[1][pair[1].bars[0]]), '|'.join(s.id for s in pair), tuple(s.bars[0] for s in pair)))
  score, _, pair = min(candidates, key=lambda candidate: candidate[:2])
  if score[0]:
    failing = next(item.measure for i, item in enumerate(contexts) if item.intervals and not sounding_spans(pair[i % 2], item))
    raise ValueError(f'measure {failing["number"]}: no rhythm motif provides acoustic accompaniment in every harmonically active bar')
  fourth = [(item, companions[i] if companions else ()) for i, item in enumerate(contexts) if i % 4 == 3]
  second = pair[1]
  alternatives = [bar for bar in bars if sum(bar.beat(n) != second.beat(n) for n in range(int(second.length))) == 1]
  variation = min(alternatives, key=lambda bar: (rank(bar, fourth, role), bars[bar])) if alternatives and fourth else second
  if rank(variation, fourth, role)[0]:
    variation = second
  result = (pair[0], second, pair[0], variation)
  if role == 'electric' and any(item.intervals for item in contexts) and not any(sounding_spans(result[i % 4], item) for i, item in enumerate(contexts)):
    raise ValueError(f'measure {contexts[0].measure["number"]}: rhythm library cannot provide an audible electric response')
  return result

def plan_section(measures: list[dict], library: generator.rhythm.Library, *, initial: dict | None = None, electric: bool = True, acoustic_companions: bool = True) -> Plan:
  previous = carried_chords(measures, initial)
  basses, uppers, responses = [], [], []
  start = 0
  while start < len(measures):
    meter = measures[start]['time_signature']
    snippets = library.for_meter(meter)
    end = start + 1
    while end < len(measures) and measures[end]['time_signature'] == meter and not measures[end].get('section_label'):
      end += 1
    segment = measures[start:end]
    for measure in segment:
      if any('meter' in change and generator.model.fraction(change['onset']) != 0 for change in measure.get('attributes', [])):
        raise ValueError(f'measure {measure["number"]}: accompaniment does not support mid-measure meter changes')
    contexts = [context(m, previous[start + i]) for i, m in enumerate(segment)]
    bass = cycle(snippets, contexts, 'bass')
    upper = cycle(snippets, contexts, 'upper')
    companions = [tuple(sorted(set(sounding_attacks(bass[i % 4], item)) | set(sounding_attacks(upper[i % 4], item)))) for i, item in enumerate(contexts)]
    response = cycle(snippets, contexts, 'electric', companions if acoustic_companions else None) if electric else upper
    basses.extend(bass[i % 4] for i in range(len(segment)))
    uppers.extend(upper[i % 4] for i in range(len(segment)))
    responses.extend(response[i % 4] for i in range(len(segment)))
    start = end
  return Plan(tuple(basses), tuple(uppers), tuple(responses))

def same_chord(left: dict | None, right: dict | None) -> bool:
  return left is not None and right is not None and all(left.get(k) == right.get(k) for k in ('root', 'kind', 'degrees', 'bass', 'inversion'))

def normalize_beams(events: list[dict]) -> None:
  for level in sorted({beam['level'] for event in events for beam in event.get('beams', [])}):
    groups, group = [], []
    for event in events:
      beam = next((b for b in event.get('beams', []) if b['level'] == level), None)
      if not beam or beam['type'] in ('begin', 'forward hook', 'backward hook'):
        groups.append(group)
        group = []
      if beam:
        group.append((event, beam))
        if beam['type'] in ('end', 'forward hook', 'backward hook'):
          groups.append(group)
          group = []
    groups.append(group)
    for group in groups:
      for index, (event, beam) in enumerate(group):
        if len(group) > 1:
          beam['type'] = 'begin' if index == 0 else 'end' if index == len(group) - 1 else 'continue'
        elif level > 1 and any(b['level'] == level - 1 for b in event['beams']):
          beam['type'] = 'backward hook' if beam['type'] in ('end', 'backward hook') else 'forward hook'
        else:
          event['beams'].remove(beam)
  for event in events:
    if not event.get('beams'):
      event.pop('beams', None)

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 realize(measure: dict, bar: generator.rhythm.Bar, previous: dict | None, *, lane: str = '') -> list[dict]:
  item = context(measure, previous)
  length = generator.model.fraction(measure['duration'])
  events = []
  for source in bar.events:
    start, end = max(fractions.Fraction(), source.onset - item.offset), min(length, source.onset + source.duration - item.offset)
    if start >= end:
      continue
    chord = active_chord(measure, start, previous) if source.sounding else None
    continuation = source.continuation and source.onset - item.offset == start
    if continuation:
      chord = None
    sounding_end = end
    for change in measure['harmonies']:
      time = generator.model.fraction(change['onset'])
      if chord and start < time < sounding_end and not same_chord(chord, active_chord(measure, time, previous)):
        sounding_end = time
    unchanged = end - start == source.duration and sounding_end == end
    duration = sounding_end - start
    written = {'type': source.written_type, 'dots': source.dots} if unchanged else generator.model.written_for(duration)
    if not unchanged and (lane == 'electric' or not chord):
      duration, written = next(written_durations(duration))
    event = generator.model.make_event(f'{measure["id"]}_rhythm_{lane}_{len(events)}', start, duration, 'note' if chord else 'rest', written=written)
    if chord:
      event['_harmony'] = chord
      if lane:
        event.update(_guitar_profile='electric' if lane == 'electric' else 'acoustic', _guitar_lane=lane)
    if source.beams and unchanged and chord:
      event['beams'] = [{'level': level, 'type': kind} for level, kind in source.beams]
    events.append(event)
    cursor = start + duration
    for duration, written in written_durations(end - cursor):
      events.append(generator.model.make_event(f'{measure["id"]}_rhythm_{lane}_{len(events)}', cursor, duration, 'rest', written=written))
      cursor += duration
  normalize_beams(events)
  return events

def connect_continuations(document: dict) -> bool:
  previous = {}
  changed = False
  for measure in document['staves'][0]['measures']:
    if measure.get('section_label'):
      previous.clear()
    for voice in measure['voices']:
      for note in voice['notes']:
        continuation = note.pop('_rhythm_continuation', False)
        source_note = previous.get(voice['id'])
        if continuation and source_note and source_note['pitches'] and note['pitches']:
          old = {(p['tablature']['string'], p['tablature']['fret']): p for p in source_note['pitches']}
          for tone in note['pitches']:
            source = old[(tone['tablature']['string'], tone['tablature']['fret'])]
            tone['tie_from'] = source['id']
            source['tie_to'] = tone['id']
          changed = True
        elif continuation:
          note.update(kind='rest', pitches=[])
          note.pop('beams', None)
          changed = True
        previous[voice['id']] = note
  return changed