from __future__ import annotations

import dataclasses
import fractions

import generator.fretboard
import generator.harmony
import generator.model
import generator.scales.selection
import generator.solo.harmonic
import generator.solo.analysis
import generator.solo.policy
import generator.solo.sequences
import generator.solo.types

def target_pitches(chord: dict, final: bool) -> frozenset[int]:
  degrees = generator.harmony.degrees(chord)
  explicit = {d['value'] for d in chord.get('degrees', []) if d['type'] != 'subtract'}
  preferred = {1, 3} if final else {3, 7}
  preferred |= explicit | {d for d, a in degrees.items() if d == 5 and a}
  if 3 not in degrees:
    preferred |= {2, 4}
  root = generator.scales.selection.root_number(chord)
  return frozenset(int(root + generator.harmony.degree_semitones(d, a)) % 12 for d, a in degrees.items() if d in preferred) or frozenset(generator.harmony.pitch_classes(chord))

def destination(choices, start, end, final=False, release=fractions.Fraction(1, 2), closure='closed'):
  last = next((c for c in reversed(choices) if c.span.start < end and c.span.end > start), None)
  if not last or not last.span.chord:
    return None
  latest = end - release - fractions.Fraction(1, 4)
  if last.span.start > latest:
    return None
  pitches = target_pitches(last.span.chord, final)
  if closure == 'open' and last.scale:
    root = generator.scales.selection.root_number(last.span.chord)
    degrees = generator.harmony.degrees(last.span.chord)
    additions = {int(root + generator.harmony.degree_semitones(d, degrees.get(d, 0))) % 12 for d in (5, 9, 11, 13) if d not in {v['value'] for v in last.span.chord.get('degrees', []) if v['type'] == 'subtract'}}
    pitches = frozenset(additions & last.scale.pitches) or pitches
  return generator.solo.types.TargetEvent(max(start, last.span.start, end - 2 - release), latest, pitches, generator.harmony.symbol(last.span.chord), 'open continuation' if closure == 'open' else 'phrase resolution' if final else 'defining chord-tone arrival')

def new_motif(sequences: generator.solo.sequences.Library, length: fractions.Fraction, *, choices=(), register=(55, 74), seed='', variant=0) -> generator.solo.types.Motif:
  if choices and any(c.scale for c in choices):
    scale = next(c.scale for c in choices if c.scale)
    arrival = next(c for c in reversed(choices) if c.scale)
    count = min(6, max(3, int(length * fractions.Fraction(3, 2))))
    shapes = (('rising', 4), ('falling', -4), ('arch', 1), ('valley', -1), ('leap_return', 2), ('neighbor', 0))
    shape, displacement = shapes[(int(generator.solo.policy.stable_rank(seed, 'invention')[:8], 16) + variant) % len(shapes)]
    paths = sequences.retrieve(count, displacement, (displacement > 0) - (displacement < 0))
    candidates = []
    target = target_pitches(arrival.span.chord, False)
    for path in paths:
      for start in range(-14, 15):
        pitches = tuple(generator.solo.types.scale_pitch(scale, start + offset) for offset in path.offsets())
        if not all(register[0] <= pitch.number() <= register[1] and generator.fretboard.single_positions_for_pitch(fractions.Fraction(pitch.number())) for pitch in pitches):
          continue
        score = (generator.solo.sequences.family(path.steps) != shape, pitches[-1].number() % 12 not in target, abs(path.offsets()[-1] - displacement), abs(pitches[0].number() - (register[0] + 5)), generator.solo.policy.stable_rank(seed, (variant, path.steps, start)))
        candidates.append((score, pitches, path))
    if candidates:
      _, pitches, path = min(candidates, key=lambda c: c[0])
      unit = min(fractions.Fraction(1, 2), length / count)
      notes = tuple(generator.solo.types.Note(i * unit, (i + 1) * unit, pitch) for i, pitch in enumerate(pitches))
      identifier = generator.solo.policy.stable_rank(seed, (variant, path.steps))[:12]
      return generator.solo.types.Motif('invented_' + identifier, notes, length, (), 'sequences ' + ' + '.join(path.sources))
  paths = sequences.connecting(3) or sequences.connecting(2) or sequences.connecting(4)
  if not paths:
    raise ValueError('solo composition requires a two-, three-, or four-note sequence')
  path = min(paths, key=lambda s: (sum(a * b >= 0 for a, b in zip(s.steps, s.steps[1:])), s.rank()))
  count = len(path.steps) + 1
  unit = min(fractions.Fraction(1, 2), length / (count + 1))
  notes = []
  for i, offset in enumerate(path.offsets()):
    octave, degree = divmod(28 + offset, 7)
    pitch = generator.solo.types.Pitch(generator.model.LETTERS[degree], fractions.Fraction(), octave)
    notes.append(generator.solo.types.Note(unit * i, unit * (i + 1), pitch))
  return generator.solo.types.Motif('new_' + path.id.replace('/', '_').removesuffix('.xml'), tuple(notes), length, (), 'sequence ' + path.id)

def choose_motifs(analysis, sequences, length):
  eligible = [m for m in analysis.motifs if m.duration <= 2 * length]
  head = next((m for m in eligible if len(m.notes) >= 2), eligible[0] if eligible else new_motif(sequences, length))
  others = [m for m in eligible if m.id != head.id]
  contrast = min(others, key=lambda m: (m.contour() == head.contour(), abs(len(m.notes) - 3), m.duration, m.id)) if others else head
  return head, contrast, new_motif(sequences, length)

def activity_passages(start, end, release, activity, shape=0):
  music_end = end - release
  boundaries = sorted({start, music_end} | {start + fractions.Fraction(round(float((end - start) * i / 4) * 2), 2) for i in (1, 2, 3) if start < start + fractions.Fraction(round(float((end - start) * i / 4) * 2), 2) < music_end})
  weights = ((1, 1.6, 1.1), (0.8, 1.7, 1.2), (1.5, 1.4, 0.8))[shape % 3]
  spans = list(zip(boundaries, boundaries[1:]))
  arrival_rate = fractions.Fraction(3, 8)
  desired = max(0, activity * (end - start) - arrival_rate * (spans[-1][1] - spans[-1][0]))
  denominator = sum(float(b - a) * weights[min(i, 2)] for i, (a, b) in enumerate(spans[:-1]))
  result = []
  for i, (a, b) in enumerate(spans):
    arrival = i == len(spans) - 1
    rate = arrival_rate if arrival else fractions.Fraction(str(min(2, float(desired) * weights[min(i, 2)] / max(denominator, 0.001)))).limit_denominator(1000)
    occupancy = generator.solo.policy.DEFAULT.arrival_occupancy if arrival else generator.solo.policy.DEFAULT.active_occupancy
    result.append(generator.solo.types.ActivitySpan(a, b, 'sustained_arrival' if arrival else 'active', rate, fractions.Fraction(str(occupancy))))
  if release:
    result.append(generator.solo.types.ActivitySpan(music_end, end, 'space', fractions.Fraction(), fractions.Fraction()))
  return tuple(result)

def decision(seed, name):
  return int(generator.solo.policy.stable_rank(seed, name)[:8], 16)

def same_harmony(left, right):
  return bool(left and right and left.span.chord and right.span.chord and generator.harmony.symbol(left.span.chord) == generator.harmony.symbol(right.span.chord))

def resolutions(choices):
  distinct = []
  for choice in choices:
    if not distinct or not same_harmony(distinct[-1], choice):
      distinct.append(choice)
  return tuple(right for left, right in zip(distinct, distinct[1:]) if left.span.chord and right.span.chord and generator.scales.selection.dominant(left.span.chord) and (generator.scales.selection.root_number(left.span.chord) - generator.scales.selection.root_number(right.span.chord)) % 12 == 7)

def anchor_backbones(identifier, start, end, release, register, choices, bar, seed):
  active = [c for c in choices if c.scale and c.span.start < end - release and c.span.end > start]
  if not active:
    return tuple([()])
  finish = min(end - release, active[-1].span.end)
  terminal_duration = min(fractions.Fraction(2), (finish - start) / 2)
  terminal_start = max(start, active[-1].span.start, finish - terminal_duration)
  terminal_duration = finish - terminal_start
  opening = max(start, active[0].span.start)
  moments = [(opening, 'opening')]
  meaningful = list(resolutions(active)) + [b for a, b in zip(active, active[1:]) if a.scale.pitches != b.scale.pitches]
  for choice in sorted(meaningful, key=lambda c: c.span.start):
    time = choice.span.start
    if time >= moments[-1][0] + bar * fractions.Fraction(3, 2) and time + bar < terminal_start and len(moments) < 4:
      moments.append((time, 'intermediate'))
  if terminal_start <= opening + fractions.Fraction(1, 2):
    moments = []
  moments.append((terminal_start, 'terminal'))
  states = [(0.0, ())]
  for index, (time, role) in enumerate(moments):
    choice = next(c for c in active if c.span.start <= time < c.span.end)
    chord_roles = generator.solo.harmonic.roles(choice.span.chord)
    eligible = (chord_roles.basic | chord_roles.specified) & choice.scale.pitches
    duration = terminal_duration if role == 'terminal' else fractions.Fraction(1, 2)
    latest = time if role == 'terminal' else min(time + fractions.Fraction(3, 2), choice.span.end - duration, terminal_start - duration)
    pitches = [generator.solo.types.scale_pitch(choice.scale, degree) for degree in range(-21, 28)]
    pitches = [p for p in pitches if register[0] <= p.number() <= register[1] and p.number() % 12 in eligible and generator.fretboard.single_positions_for_pitch(fractions.Fraction(p.number()))]
    desired = register[0] + (5 if index == 0 else 8 if role == 'terminal' else 11)
    following = []
    for cost, anchors in states:
      for pitch in pitches:
        jump = abs(pitch.number() - anchors[-1].pitch.number()) if anchors else 0
        direction = 1 if role != 'terminal' else -1
        wrong_way = bool(anchors and (pitch.number() - anchors[-1].pitch.number()) * direction < 0)
        preferred = target_pitches(choice.span.chord, role == 'terminal')
        score = cost + jump / 12 + 0.3 * wrong_way + abs(pitch.number() - desired) / 18 + 0.3 * (pitch.number() % 12 not in chord_roles.basic) + 0.2 * (role == 'terminal' and pitch.number() % 12 not in preferred)
        anchor = generator.solo.types.Anchor(f'{identifier}_anchor_{index + 1}', time, latest, pitch, frozenset(eligible), register, duration, role, generator.harmony.symbol(choice.span.chord), False, finish if role == 'terminal' else None)
        following.append((score, (*anchors, anchor)))
    states = list(generator.solo.sequences.diverse(following, lambda item: (item[0], generator.solo.policy.stable_rank(seed, tuple(a.pitch for a in item[1]))), lambda item: (item[1][0].pitch.number(), item[1][-1].pitch.number()), generator.solo.policy.DEFAULT.alternatives))
  return tuple(anchors for _, anchors in states)

def plans(analysis: generator.solo.analysis.Analysis, measures: list[dict], choices, end: fractions.Fraction, part: str, sequences: generator.solo.sequences.Library) -> tuple[generator.solo.types.PhrasePlan, ...]:
  boundaries = [fractions.Fraction()]
  for measure in measures:
    boundaries.append(min(end, boundaries[-1] + generator.model.fraction(measure['duration'])))
  head, contrast, _ = choose_motifs(analysis, sequences, boundaries[1])
  seed = generator.solo.policy.stable_rank(generator.solo.analysis.identity(analysis, measures, part), generator.solo.policy.SEED_NAMESPACE)
  activity = fractions.Fraction(str(generator.solo.policy.DEFAULT.acoustic_activity if part == 'acoustic' else generator.solo.policy.DEFAULT.electric_activity))
  result, index, passage_index, inventions = [], 0, 0, {}
  cadences = resolutions(choices)
  while index < len(measures) and boundaries[index] < end:
    start = boundaries[index]
    stops = sorted({min(index + size, len(measures)) for size in (4, 6, 8)})

    def destination_rank(stop):
      finish = boundaries[stop]
      last = next((c for c in reversed(choices) if c.scale and c.span.start < finish and c.span.end > start), None)
      resolved = bool(last and any(start < c.span.start < finish and same_harmony(c, last) for c in cadences))
      dominant = bool(last and generator.scales.selection.dominant(last.span.chord))
      return (0 if resolved else 2 if dominant else 1), abs(stop - index - 4), finish

    stop = min(stops, key=destination_rank)
    finish = boundaries[stop]
    if finish <= start:
      break
    passage_index += 1
    passage_id = f'{part}_passage_{passage_index}'
    passage_seed = generator.solo.policy.stable_rank(seed, (str(start), str(finish)))
    low, high = (55, 74) if part == 'acoustic' else (59, 79)
    phase = (start + finish) / (2 * end)
    register = (low + (5 if fractions.Fraction(1, 2) <= phase < fractions.Fraction(3, 4) else 0), high if phase >= fractions.Fraction(1, 2) else high - 4)
    register = generator.fretboard.single_register(register)
    bar = boundaries[index + 1] - start
    release = fractions.Fraction(1, 2) if decision(passage_seed, 'breathing') % 3 == 0 else fractions.Fraction()
    release = min(release, (finish - start) / 8)
    envelope = activity_passages(start, finish, release, activity, decision(passage_seed, 'activity') % 3)
    anchors = anchor_backbones(passage_id, start, finish, release, register, choices, bar, passage_seed)
    # Prefer four-bar phrasing; extend when a sustained harmony crosses that seam.
    middle = min(index + 4, stop)
    before = next((c for c in reversed(choices) if c.span.start < boundaries[middle]), None)
    after = next((c for c in choices if c.span.start <= boundaries[middle] < c.span.end), None)
    extended = stop - index > 4 and same_harmony(before, after)
    phrase_stops = [stop] if extended or middle == stop else [middle, stop]
    local = tuple(c for c in choices if c.span.start < boundaries[phrase_stops[0]] and c.span.end > start)
    slot = decision(passage_seed, 'motif') % 3
    if part == 'acoustic':
      slot = 0 if passage_index <= 2 else slot
      if slot == 2 and 0 not in inventions:
        inventions[0] = new_motif(sequences, bar, choices=local, register=register, seed=passage_seed)
      motif = head if slot == 0 else contrast if slot == 1 else inventions[0]
    else:
      slot = 0 if passage_index == 1 else slot
      if slot < 2 and slot not in inventions:
        inventions[slot] = new_motif(sequences, bar, choices=local, register=register, seed=passage_seed, variant=slot)
      motif = head if slot == 2 else inventions[slot]
    gestures, cursor = [], start
    while cursor < finish - release:
      gesture_end = min(cursor + bar, finish - release)
      closing = gesture_end == finish - release
      changed = any(c.span.start == cursor and c.scale and prior.scale and c.scale.pitches != prior.scale.pitches for prior, c in zip(choices, choices[1:]))
      relationship = 'statement' if not gestures else 'close' if closing else 'repeat_with_harmonic_adaptation' if changed else 'answer' if cursor >= start + (finish - start) / 2 else 'extend'
      treatment = 'opening_reference' if not gestures and motif.origin == 'melody' else 'invented' if not gestures else 'rhythmic_reference' if relationship == 'repeat_with_harmonic_adaptation' else 'answer' if relationship in ('answer', 'close') else 'contour_reference'
      gestures.append(generator.solo.types.Gesture(f'{passage_id}_gesture_{len(gestures) + 1}', cursor, gesture_end, relationship, treatment))
      cursor = gesture_end
    phrase_start = start
    for phrase_stop in phrase_stops:
      phrase_end = boundaries[phrase_stop]
      terminal = phrase_end == finish
      treatment = 'opening_reference' if motif.origin == 'melody' else 'invented'
      local_gestures = tuple(g for g in gestures if g.start < phrase_end and g.end > phrase_start)
      landmarks = tuple(generator.solo.types.Landmark(a.earliest, a.latest, min(v[i].pitch.number() for v in anchors), max(v[i].pitch.number() for v in anchors), 0, a.role) for i, a in enumerate(anchors[0])) if anchors and anchors[0] else ()
      result.append(generator.solo.types.PhrasePlan(f'{part}_phrase_{len(result) + 1}', phrase_start, phrase_end, motif, 'statement' if phrase_start == start else 'continue', treatment, None, register, activity, release if terminal else fractions.Fraction(), envelope, landmarks, 'breath' if terminal and release else 'sustain' if terminal else 'connected', 'closed' if terminal else 'open', passage_seed, passage_id, start, finish, anchors, local_gestures))
      phrase_start = phrase_end
    index = stop
  return tuple(result)