from __future__ import annotations

import copy
import fractions

import generator.accompaniment
import generator.acoustic
import generator.acoustic_voicing
import generator.arrange.form
import generator.electric
import generator.guitar
import generator.model
import generator.parts
import generator.score
import generator.solo.analysis
import generator.solo.composer
import generator.solo.notation
import generator.solo.report
import generator.solo.types

def regions(document):
  result, first, start = {}, 0, fractions.Fraction()
  source = generator.score.source_staff(document)['measures']
  for section in generator.arrange.form.full_arrangement_sections(source, 'acoustic'):
    length = sum((generator.model.fraction(m['duration']) for m in source[section.source_measure_index:section.source_measure_index + section.measure_count]), fractions.Fraction())
    result[section.section_index] = (first, section.measure_count, start, length)
    first += section.measure_count
    start += length
  return result

def annotate(document, solo, offset):
  spans = []
  for choice in solo.choices:
    if choice.scale is None:
      continue
    label = choice.scale.label()
    if spans and spans[-1]['label'] == label and generator.model.fraction(spans[-1]['end']) == offset + choice.span.start:
      spans[-1]['end'] = generator.model.rational(offset + choice.span.end)
    else:
      spans.append({'id': f'{solo.part}_scale_{len(spans) + 1}', 'start': generator.model.rational(offset + choice.span.start), 'end': generator.model.rational(offset + choice.span.end), 'label': label, 'matched': True})
  document['staves'][0]['scale_spans'] = spans

def reuse_head_grips(document, references, choices):
  previous = generator.accompaniment.carried_chords(references)
  target_previous = generator.accompaniment.carried_chords(document['staves'][0]['measures'])
  links, absolute = [], fractions.Fraction()
  for index, measure in enumerate(document['staves'][0]['measures']):
    for note in generator.score.measure_notes(measure):
      onset = generator.model.fraction(note['onset'])
      chord = generator.accompaniment.active_chord(measure, onset, target_previous[index])
      if not note['pitches'] or chord is None:
        continue
      allowed = generator.solo.types.permitted(choices, absolute + onset, absolute + onset + generator.model.fraction(note['duration']))
      matches = []
      for candidate in generator.score.measure_notes(references[index]):
        if len(candidate['pitches']) < 3:
          continue
        time = generator.model.fraction(candidate['onset'])
        head_chord = generator.accompaniment.active_chord(references[index], time, previous[index])
        if generator.accompaniment.same_chord(chord, head_chord) and {int(generator.model.pitch_number(p)) % 12 for p in candidate['pitches']} <= allowed:
          matches.append((abs(time - onset), candidate['id'], candidate))
      links.append(generator.electric.VoicingLink(note, tuple(item[2] for item in sorted(matches))))
    absolute += generator.model.fraction(measure['duration'])
  generator.electric.reuse_voicings(links)

def backing(baseline, foreground, part, rhythms, references, choices, initial=None):
  plan = generator.accompaniment.plan_section(foreground, rhythms, initial=initial, electric=part == 'electric', acoustic_companions=False)
  previous = generator.accompaniment.carried_chords(foreground, initial)
  measures = []
  for index, measure in enumerate(foreground):
    rhythm = (plan.acoustic_bass[index], plan.acoustic_upper[index]) if part == 'acoustic' else plan.electric[index]
    result = generator.parts.build_measure(measure, measure['id'], 'accompaniment', part, previous[index], measure.get('section_label', ''), rhythm)
    measures.append(result)
  generator.score.reidentify(measures, [])
  staff = baseline['staves'][0]
  document = generator.score.make_arranged_document(baseline, staff, staff_id=staff['id'], staff_name=staff['name'], role='guitar', arrangement_part_type=part, measures=measures)
  if part == 'acoustic':
    generator.acoustic.prepare(document)
    generator.acoustic_voicing.place(document)
    generator.acoustic.combine(document)
  else:
    generator.guitar.place_guitar(document)
    reuse_head_grips(document, references, choices)
  generator.model.validate_document(document)
  return document['staves'][0]['measures']

def apply(source, parts, *, rhythms, scales, sequences, solos=None):
  """Compose once per song, replacing only already-placed solo section windows."""
  analysis = generator.solo.analysis.analyze(source)
  sections = regions(source)
  source_measures = generator.score.source_staff(source)['measures']
  head_first, head_count, _, _ = sections[2]
  head = parts['electric']['staves'][0]['measures'][head_first:head_first + head_count]
  reports = []
  for section_index, part, other in ((3, 'acoustic', 'electric'), (4, 'electric', 'acoustic')):
    first, count, offset, length = sections[section_index]
    initial = generator.accompaniment.carried_chords(parts[part]['staves'][0]['measures'])[first]
    template = copy.deepcopy(parts[part]['staves'][0]['measures'][first:first + count])
    cutoff = length
    if part == 'electric' and generator.score.body_start_measure_index(source_measures):
      pickup = source_measures[0]
      sounding = [generator.model.fraction(n['onset']) for n in generator.score.measure_notes(pickup) if n['pitches']]
      if sounding:
        cutoff -= generator.model.fraction(pickup['duration']) - min(sounding)
    solo = generator.solo.composer.compose(analysis, template, part, cutoff, scales=scales, sequences=sequences, rhythms=rhythms, initial=initial)
    if solos is not None:
      solos.append(solo)
    foreground = generator.solo.notation.measures(template, solo.notes(), 'solo_' + part)
    generator.solo.notation.replace_region(parts[part], first, foreground)
    annotate(parts[part], solo, offset)
    # The other guitarist responds to the actual foreground, including vocal return.
    if cutoff < length:
      vocal = parts['vocal']['staves'][0]['measures'][first:first + count]
      time = fractions.Fraction()
      for measure, voice_measure in zip(foreground, vocal):
        returning = [copy.deepcopy(event) | {'voice': 'vocal_return'} for event in generator.score.measure_notes(voice_measure) if event['pitches'] and time + generator.model.fraction(event['onset']) >= cutoff]
        if returning:
          measure['voices'].append({'id': 'vocal_return', 'notes': returning})
        time += generator.model.fraction(measure['duration'])
    accompaniment = backing(parts[other], foreground, other, rhythms, head, solo.choices, initial)
    generator.solo.notation.replace_region(parts[other], first, accompaniment)
    reports.append(generator.solo.report.describe(solo, scales, first, offset))
  for document in parts.values():
    generator.model.validate_document(document)
  return '\n'.join(reports)