"""Assemble source-defined role windows without composing or respelling."""
from __future__ import annotations
import copy
import fractions

import generator.accompaniment
import generator.arrange.form
import generator.arrange.pickup
import generator.composition_types
import generator.fretboard
import generator.model
import generator.parts
import generator.score
import generator.validation

def mask(document, allowed):
  """Apply source-defined role exclusions, without observing another part."""
  absolute = fractions.Fraction()
  for measure in document['staves'][0]['measures']:
    length = generator.model.fraction(measure['duration'])
    for voice in measure['voices']:
      selected = []
      for n in voice['notes']:
        start = absolute + generator.model.fraction(n['onset'])
        end = start + generator.model.fraction(n['duration'])
        intervals = [(max(start, a), min(end, b)) for a, b in allowed if a < end and b > start]
        for a, b in intervals:
          note = copy.deepcopy(n)
          if a != start:
            note['onset'] = generator.model.rational(a - absolute)
            for pitch in note['pitches']:
              pitch.pop('tie_from', None)
          generator.arrange.pickup.truncate(note, b - absolute)
          if a < b:
            selected.append(note)
      voice['notes'] = fill(selected, length, measure['id'] + '_' + voice['id'], voice['id'])
    absolute += length
  staff = document['staves'][0]
  staff['tuplets'] = generator.arrange.pickup.retain_links(staff['measures'], staff.get('tuplets', []))

def fill(notes, length, prefix, voice='voice_arranged'):
  notes = sorted(notes, key=lambda n: generator.model.fraction(n['onset']))
  result = []
  cursor = fractions.Fraction()

  def rests(end):
    nonlocal cursor
    for duration, written in generator.accompaniment.written_durations(end - cursor):
      result.append(generator.model.make_event(prefix + '_rest_' + str(len(result)), cursor, duration, 'rest', written=written) | {'voice': voice})
      cursor += duration

  for n in notes:
    start = generator.model.fraction(n['onset'])
    if start < cursor:
      raise ValueError(prefix + ': overlapping independently generated roles')
    rests(start)
    n['voice'] = voice
    result.append(n)
    cursor = start + generator.model.fraction(n['duration'])
  rests(length)
  return result

def improvisation(source, template, part):
  """Fill the source's solo allocation, including pickups, with beat slashes."""
  document = copy.deepcopy(template)
  staff = document['staves'][0]
  for field in ('tuplets', 'octave_shifts', 'scale_spans'):
    staff[field] = []
  for measure in staff['measures']:
    measure['id'] = part + '_improv_' + measure['id']
    measure['voices'] = [{'id': 'voice_arranged', 'notes': generator.parts.generated_voice(measure, 'slash')}]
  generator.score.reidentify(staff['measures'], [])
  plan = generator.arrange.form.FormPlan.read(source.document(), part)
  mask(document, [(a, b) for a, b, role in plan.windows if role == part + '_solo'])
  return document

def assemble(source, results, *, solos_enabled=True):
  """Union prevalidated role windows; no selection or pitch/fingering rewrite."""
  documents = {}
  for part in generator.parts.PART_TYPES:
    selected = [results[r]['document'] for r in generator.composition_types.ROLES
      if r.startswith(part + '_') and (solos_enabled or not r.endswith('_solo'))]
    if not solos_enabled and part != 'vocal':
      selected.append(improvisation(source, selected[0], part))
    document = copy.deepcopy(selected[0])
    staff = document['staves'][0]
    if part == 'vocal':
      documents[part] = document
      continue
    staff['tuplets'] = []
    staff['octave_shifts'] = []
    staff['scale_spans'] = []
    for other in selected:
      for field in ('tuplets', 'octave_shifts', 'scale_spans'):
        staff[field].extend(copy.deepcopy(other['staves'][0].get(field, [])))
    for i, m in enumerate(staff['measures']):
      # Tuplet rests belong to the source/role rhythm and must survive union.
      lanes = {}
      for d in selected:
        for voice in d['staves'][0]['measures'][i]['voices']:
          lane = voice['id'] if part == 'acoustic' and voice['id'] in ('upper', 'bass') else 'voice_arranged'
          sounding = [copy.deepcopy(n) for n in voice['notes'] if n['pitches'] or n.get('tuplet_ids') or n.get('notehead') == 'slash']
          if sounding:
            lanes.setdefault(lane, []).extend(sounding)
      lanes = lanes or {'voice_arranged': []}
      m['voices'] = [{'id': lane, 'notes': fill(notes, generator.model.fraction(m['duration']), m['id'] + '_' + lane, lane)} for lane, notes in lanes.items()]
      # Attribute changes belong to source melody pickups and are assembly metadata.
      for d in selected:
        other = d['staves'][0]['measures'][i]
        if other.get('attributes'):
          m['attributes'] = copy.deepcopy(other['attributes'])
    staff['tuplets'] = generator.arrange.pickup.retain_links(staff['measures'], staff['tuplets'])
    owners = {n['id']: v['id'] for m in staff['measures'] for v in m['voices'] for n in v['notes']}
    for span in staff['tuplets']:
      span['voice'] = owners[span['start']]
    generator.validation.validate_document(document)
    generator.fretboard.validate_document(document)
    documents[part] = document
  return documents