from __future__ import annotations

import copy
import fractions
import pathlib

import generator.accompaniment
import generator.acoustic
import generator.arrange.form
import generator.electric
import generator.guitar
import generator.melody
import generator.model
import generator.musicxml
import generator.parts
import generator.rhythm
import generator.score

FINAL_SECTION_INDEX = 6

def arrange_document(document: dict, part_type: str, *, rhythms: generator.rhythm.Library) -> dict:
  document = generator.parts.prepare_document(document)
  plans = plan_arrangement(document, rhythms) if part_type != 'vocal' else {}
  return _arrange_document(document, part_type, plans)

def section_source(source: list[dict], final: bool) -> list[dict]:
  result = copy.deepcopy(source)
  for measure in result:
    measure['harmonies'] = [h for h in measure['harmonies'] if not final or not h.get('chord_optional')]
    for chord in measure['harmonies']:
      chord.pop('chord_optional', None)
  return result

def plan_arrangement(document: dict, rhythms: generator.rhythm.Library) -> dict[int, generator.accompaniment.Plan]:
  source = generator.score.source_staff(document)['measures']
  result = {}
  for section in generator.arrange.form.full_arrangement_sections(source, 'acoustic'):
    if section.role == 'slash':
      continue
    measures = section_source(source, section.section_index == FINAL_SECTION_INDEX)
    previous = generator.accompaniment.carried_chords(measures)
    start = section.source_measure_index
    result[section.section_index] = generator.accompaniment.plan_section(measures[start:start + section.measure_count], rhythms, initial=previous[start], electric=section.section_index in (1, 5, 6))
  return result

def _arrange_document(document: dict, part_type: str, plans: dict[int, generator.accompaniment.Plan]) -> dict:
  if part_type not in generator.parts.PART_TYPES:
    raise ValueError(f'unknown arrangement part: {part_type}')
  staff = generator.score.source_staff(document)
  source = staff['measures']
  sections = generator.arrange.form.full_arrangement_sections(source, part_type)
  output, tuplets, octave_shifts, voicing_origins = [], [], [], []
  absolute = fractions.Fraction()
  source_starts = []
  source_time = fractions.Fraction()
  for m in source:
    source_starts.append(source_time)
    source_time += generator.model.fraction(m['duration'])
  for section in sections:
    section_measures = section_source(source, section.section_index == FINAL_SECTION_INDEX)
    previous = generator.accompaniment.carried_chords(section_measures)
    segment = []
    for si in range(section.source_measure_index, section.source_measure_index + section.measure_count):
      m = section_measures[si]
      rhythm = None
      if section.role == 'accompaniment':
        plan = plans[section.section_index]
        rhythm = (plan.acoustic_bass[len(segment)], plan.acoustic_upper[len(segment)]) if part_type == 'acoustic' else plan.electric[len(segment)]
      result = generator.parts.build_measure(m, f'measure-{len(output) + len(segment) + 1}', section.role, part_type, previous[si], section.label if not segment else '', rhythm, melody_octave=True)
      segment.append(result)
      if part_type == 'electric' and section.role in ('melody', 'accompaniment'):
        voicing_origins.append(generator.electric.MeasureOrigin(source_starts[si], section.role, result))
    identity_segment = segment
    segment_length = sum((generator.model.fraction(m['duration']) for m in segment), fractions.Fraction())
    if section.role in ('vocal', 'melody'):
      start = source_starts[section.source_measure_index]
      for shift in staff.get('octave_shifts', []):
        low, high = max(start, generator.model.fraction(shift['start'])), min(start + segment_length, generator.model.fraction(shift['end']))
        if low < high:
          octave_shifts.append({**shift, 'id': f'section-{section.section_index}-{shift["id"]}', 'start': generator.model.rational(absolute + low - start), 'end': generator.model.rational(absolute + high - start)})
    if part_type == 'vocal' and section.section_index in (1, 5) and section.source_measure_index == 1 and source[0].get('implicit') and output:
      pickup = generator.parts.build_measure(source[0], source[0]['id'], 'vocal', part_type, None)
      destination = output[-1]
      pad = generator.model.fraction(destination['duration']) - generator.model.fraction(pickup['duration'])
      if pad < 0:
        raise ValueError('pickup is longer than the preceding arrangement measure')
      voices = generator.score.copy_source_voices(pickup, strip_lyrics=False)
      for voice in voices:
        for n in voice['notes']:
          n['onset'] = generator.model.rational(generator.model.fraction(n['onset']) + pad)
        if pad:
          voice['notes'].insert(0, generator.model.make_event(f'pickup-rest-{voice["id"]}', fractions.Fraction(), pad, 'rest') | {'voice': voice['id']})
      destination['voices'] = voices
      identity_segment = [destination, *segment]
      if pickup['harmonies']:
        destination['harmonies'] = [h for h in destination['harmonies'] if generator.model.fraction(h['onset']) < pad]
        for h in pickup['harmonies']:
          h['onset'] = generator.model.rational(generator.model.fraction(h['onset']) + pad)
          destination['harmonies'].append(h)
      state = {k: copy.deepcopy(pickup[k]) for k in ('key', 'clef', 'meter', 'transpose')}
      if pad:
        destination['attributes'].append({'onset': generator.model.rational(pad), **state})
      else:
        destination.update(state)
        destination['key_signature'] = generator.model.key_label(state['key'])
        destination['time_signature'] = generator.model.meter_label(state['meter'])
      for change in pickup.get('attributes', []):
        destination['attributes'].append({**change, 'onset': generator.model.rational(generator.model.fraction(change['onset']) + pad)})
      for shift in staff.get('octave_shifts', []):
        low, high = generator.model.fraction(shift['start']), min(generator.model.fraction(pickup['duration']), generator.model.fraction(shift['end']))
        if low < high:
          offset = absolute - generator.model.fraction(pickup['duration'])
          octave_shifts.append({**shift, 'id': f'pickup-{section.section_index}-{shift["id"]}', 'start': generator.model.rational(offset + low), 'end': generator.model.rational(offset + high)})
    tuplets.extend(generator.score.reidentify(identity_segment, staff.get('tuplets', [])))
    output.extend(segment)
    absolute += segment_length
  result = generator.score.make_arranged_document(document, staff, staff_id='staff-arranged-' + part_type, staff_name=part_type.capitalize(), role='vocal' if part_type == 'vocal' else 'guitar', arrangement_part_type=part_type, measures=output, tuplets=tuplets, octave_shifts=octave_shifts)
  result['metadata'].setdefault('tempo', 124)
  if part_type != 'vocal':
    generator.acoustic.prepare(result)
    generator.melody.harmonize(result)
    voicing_links = generator.electric.link_voicings(source, voicing_origins) if part_type == 'electric' else []
    generator.guitar.place_guitar(result)
    if generator.accompaniment.connect_continuations(result):
      generator.guitar.place_guitar(result)
    generator.electric.reuse_voicings(voicing_links)
    generator.acoustic.combine(result)
  return result

def generate_performance_scores(source_dir: pathlib.Path, output_dir: pathlib.Path, *, rhythms: generator.rhythm.Library) -> int:
  count = 0
  for path in generator.parts.source_paths(source_dir):
    try:
      document = generator.parts.prepare_document(generator.musicxml.load_musicxml(path))
      plans = plan_arrangement(document, rhythms)
    except ValueError as error:
      raise ValueError(f'{path}: {error}') from error
    for part in generator.parts.PART_TYPES:
      try:
        generator.score.write_json(output_dir / f'{generator.score.slugify(path.stem)}-{part}.json', _arrange_document(document, part, plans))
      except ValueError as error:
        raise ValueError(f'{path}, part {part}: {error}') from error
      count += 1
  return count