from __future__ import annotations

import copy
import fractions
import pathlib

import generator.accompaniment
import generator.acoustic
import generator.acoustic_voicing
import generator.arrange.form
import generator.arrange.pickup
import generator.electric
import generator.guitar
import generator.melody
import generator.model
import generator.musicxml
import generator.parts
import generator.rhythm
import generator.scales.library
import generator.scales.mapping
import generator.scales.timeline
import generator.score

FINAL_SECTION_INDEX = 6

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

def section_source(source: list[dict], section, next_section=None) -> list[dict]:
  result = copy.deepcopy(source)
  for measure in result:
    measure['harmonies'] = [h for h in measure['harmonies'] if section.section_index != FINAL_SECTION_INDEX or not h.get('chord_optional')]
    for chord in measure['harmonies']:
      chord.pop('chord_optional', None)
  if generator.arrange.pickup.enters(next_section, source):
    destination = result[section.source_measure_index + section.measure_count - 1]
    generator.arrange.pickup.harmonies(destination, result[0])
  return result

def plan_arrangement(document: dict, rhythms: generator.rhythm.Library) -> dict[int, generator.accompaniment.Plan]:
  source = generator.score.source_staff(document)['measures']
  result = {}
  sections = generator.arrange.form.full_arrangement_sections(source, 'acoustic')
  for index, section in enumerate(sections):
    if section.role == 'slash':
      continue
    measures = section_source(source, section, sections[index + 1] if index + 1 < len(sections) else None)
    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], *, scales: generator.scales.library.Library | None = None, reports: list[str] | None = None, source_name: str = '') -> 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 = [], [], [], []
  solos = []
  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 index, section in enumerate(sections):
    section_measures = section_source(source, section, sections[index + 1] if index + 1 < len(sections) else None)
    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)
      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))
    segment_length = sum((generator.model.fraction(m['duration']) for m in segment), fractions.Fraction())
    if section.role == 'slash':
      solos.append(generator.scales.timeline.Solo(absolute, absolute + segment_length, section.label))
    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)})
    pickup = None
    if section.role in ('vocal', 'melody') and generator.arrange.pickup.enters(section, source) and output:
      pickup = generator.parts.build_measure(source[0], f'pickup_{section.section_index}', section.role, part_type, None)
    tuplets.extend(generator.score.reidentify(([pickup] if pickup else []) + segment, staff.get('tuplets', [])))
    if pickup:
      destination = output[-1]
      pad = generator.arrange.pickup.offset(destination, pickup)
      onset = generator.arrange.pickup.splice(destination, pickup)
      if part_type == 'electric':
        voicing_origins = [generator.electric.MeasureOrigin(o.onset, o.role, o.measure, o.start, onset if o.measure is destination else o.end) for o in voicing_origins]
        voicing_origins.append(generator.electric.MeasureOrigin(-pad, 'melody', destination, onset))
      for shift in staff.get('octave_shifts', []):
        low, high = max(onset - pad, 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)})
    output.extend(segment)
    absolute += segment_length
  tuplets = generator.arrange.pickup.retain_links(output, tuplets)
  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 == 'acoustic':
    generator.acoustic.prepare(result)
    generator.acoustic_voicing.place(result)
    generator.acoustic.combine(result)
  elif part_type == 'electric':
    generator.melody.harmonize(result)
    voicing_links = generator.electric.link_voicings(source, voicing_origins)
    generator.guitar.place_guitar(result)
    if generator.accompaniment.connect_continuations(result):
      generator.guitar.place_guitar(result)
    generator.electric.reuse_voicings(voicing_links)
  if scales and solos:
    report = generator.scales.mapping.annotate(result, solos, scales, source_name or document['metadata']['title'])
    if reports is not None:
      reports.append(report)
  return result

def generate_performance_scores(source_dir: pathlib.Path, output_dir: pathlib.Path, *, rhythms: generator.rhythm.Library, scales: generator.scales.library.Library | None = None) -> 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
    reports = []
    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, scales=scales, reports=reports, source_name=path.name))
      except ValueError as error:
        raise ValueError(f'{path}, part {part}: {error}') from error
      count += 1
    if reports:
      report_path = output_dir.parent / 'analysis' / f'{generator.score.slugify(path.stem)}.txt'
      report_path.parent.mkdir(parents=True, exist_ok=True)
      report_path.write_text(document['metadata']['title'] + '\nSource: ' + path.name + '\n\n' + '\n'.join(reports), encoding='utf-8')
  return count