from __future__ import annotations

import dataclasses
import fractions

import generator.fretboard
import generator.harmony
import generator.model
import generator.scales.selection
import generator.scales.timeline
import generator.solo.harmonic
import generator.solo.analysis
import generator.solo.planning
import generator.solo.policy
import generator.solo.realization
import generator.solo.types

def compose(analysis, measures, part, end, *, scales, sequences, rhythms, initial=None):
  timeline = generator.scales.timeline.read({'measures': measures}, initial=initial)
  region = generator.scales.timeline.Solo(fractions.Fraction(), end, part.capitalize() + ' Solo')
  choices = generator.scales.selection.select(timeline, region, scales)
  for choice in choices:
    if choice.span.chord and choice.scale is None:
      raise ValueError(f'{part} solo, source measure {choice.span.source_measure}, beat {choice.span.beat}, {generator.harmony.symbol(choice.span.chord)}: No exact match within the permitted catalog')
  plans = generator.solo.planning.plans(analysis, measures, choices, end, part, sequences)
  retained = tuple((plan.start + n.start, plan.start + n.end, {n.pitch.number() % 12}) for plan in plans if plan.treatment in ('quote', 'exact_quote') for n in plan.motif.notes)
  if retained and not all(pc <= generator.solo.types.permitted(choices, start, finish) for start, finish, pc in retained):
    choices = generator.scales.selection.select(timeline, region, scales, retained=retained)
    if not all(pc <= generator.solo.types.permitted(choices, start, finish) for start, finish, pc in retained):
      raise ValueError(f'{part} solo: protected quotation has no exact compatible scale collection')
    choices = tuple(dataclasses.replace(c, reason=c.reason + '; retained quotation') if any(a < c.span.end and b > c.span.start for a, b, _ in retained) else c for c in choices)
  context = generator.solo.realization.SearchContext()

  def realize(plan, previous):
    original = plan
    for attempt in range(generator.solo.policy.DEFAULT.replan_attempts + 1):
      try:
        options = generator.solo.realization.alternatives(plan, measures, choices, sequences, rhythms, analysis.tempo, previous, context)
        if attempt:
          context.results[(original, previous)] = options
        return options
      except generator.solo.realization.SearchExhausted as error:
        if attempt == generator.solo.policy.DEFAULT.replan_attempts or plan.treatment in ('quote', 'exact_quote'):
          raise
        context.diagnostics.setdefault('replanning', []).append({'phrase': plan.id, 'attempt': attempt + 1, 'reason': str(error), 'adjustment': 'contour reference' if attempt == 0 else 'opening reference'})
        # Replan only ordinary development contracts; harmony, anchor pitches,
        # required holds, timing, activity and protected quotations stay fixed.
        treatment = 'contour_reference' if attempt == 0 else 'opening_reference'
        gestures = tuple(dataclasses.replace(g, treatment=treatment) if g.relationship not in ('statement', 'close') else g for g in plan.gestures)
        plan = dataclasses.replace(plan, gestures=gestures)

  phrases = choose_phrases(plans, realize)
  context.diagnostics['policy'] = generator.solo.policy.describe()['fingerprint']
  context.diagnostics['musical_identity'] = generator.solo.analysis.identity(analysis, measures, part)
  passage_ids = tuple(dict.fromkeys(p.passage_id for p in plans if p.passage_id))
  passages = tuple(generator.solo.types.MelodicPassagePlan(identifier, min(p.start for p in plans if p.passage_id == identifier), max(p.end for p in plans if p.passage_id == identifier), tuple(p.id for p in plans if p.passage_id == identifier), next(p.seed for p in plans if p.passage_id == identifier)) for identifier in passage_ids)
  result = generator.solo.types.Solo(part, choices, tuple(phrases), analysis.tempo, end, context.diagnostics, passages)
  validate(result)
  return result

def phrase_notes(phrase):
  notes = tuple(n for d in phrase.decisions for n in d.notes)
  if notes and notes[0].continuation:
    notes = (dataclasses.replace(notes[0], continuation=False), *notes[1:])
  return generator.solo.types.sounding_notes(notes)

def transition_cost(first, second):
  before, after = phrase_notes(first), phrase_notes(second)
  if not before or not after:
    return 0.0
  gap = max(0, after[0].start - before[-1].end - first.plan.release)
  jump = abs(after[0].pitch.number() - before[-1].pitch.number())
  # The next phrase's local score already includes physical movement. Here we
  # evaluate pacing and conversational contrast at the join, once per pair.
  before_shape = tuple(b.pitch.number() - a.pitch.number() for a, b in zip(before, before[1:]))
  after_shape = tuple(b.pitch.number() - a.pitch.number() for a, b in zip(after, after[1:]))
  repeated_answer = second.plan.relationship == 'answer' and before_shape == after_shape
  return min(1, float(gap) / 2) + 0.5 * min(1, max(0, jump - 5) / 12) + 0.5 * repeated_answer

def choose_phrases(plans, realize):
  """Commit only after the next phrase has been evaluated from each ending."""
  result, previous = [], None
  for index, plan in enumerate(plans):
    options = realize(plan, previous)
    ranked, failure = [], None
    for option in options:
      notes = phrase_notes(option)
      last = notes[-1] if notes else previous
      incoming = transition_cost(result[-1], option) if result else 0
      if index + 1 == len(plans):
        ranked.append((option.score + incoming, option))
        continue
      try:
        following = realize(plans[index + 1], option.state or last)
      except generator.solo.realization.SearchExhausted as error:
        failure = error
        continue
      length = float(option.plan.end - option.plan.start)
      next_length = float(plans[index + 1].end - plans[index + 1].start)
      shared_passage = bool(plan.passage_id and plan.passage_id == plans[index + 1].passage_id)
      cost = min((future.score if shared_passage else (option.score * length + future.score * next_length) / (length + next_length)) + transition_cost(option, future) for future in following)
      ranked.append((cost + incoming, option))
    if not ranked:
      raise failure or generator.solo.realization.SearchExhausted(f'{plan.id}: no feasible phrase transition')
    _, chosen = min(ranked, key=lambda item: (item[0], tuple((n.pitch, n.start, n.end, n.position) for n in phrase_notes(item[1]))))
    result.append(chosen)
    notes = phrase_notes(chosen)
    if chosen.state:
      previous = chosen.state
    elif notes:
      previous = notes[-1]
  return tuple(result)

def validate(solo):
  previous = None
  for note in solo.notes():
    if note.start < 0 or note.end > solo.end or note.end <= note.start or previous and previous.end > note.start:
      raise ValueError(f'{solo.part} solo: overlapping or out-of-bounds note at {note.start}')
    if note.pitch.number() % 12 not in generator.solo.types.permitted(solo.choices, note.start, note.end):
      raise ValueError(f'{solo.part} solo: held pitch {note.pitch.label()} violates a scale span at {note.start}-{note.end}')
    if note.position not in generator.fretboard.positions_for_pitch(fractions.Fraction(note.pitch.number())):
      raise ValueError(f'{solo.part} solo: invalid guitar placement at {note.start}')
    previous = note
  for phrase in solo.phrases:
    notes = phrase_notes(phrase)
    if not generator.solo.realization.target_ok(notes, phrase.plan.target):
      raise ValueError(f'{phrase.plan.id}: planned harmonic arrival was not realized')
  harmony = generator.solo.harmonic.Harmony(solo.choices)
  contracts = tuple(dict.fromkeys(t for phrase in solo.phrases if phrase.state for t in phrase.state.treatments))
  checked = harmony.evaluate(solo.notes(), contracts, until=solo.end, complete=True)
  if not checked.valid:
    raise ValueError(f'{solo.part} solo: {checked.error}')
  for phrase in solo.phrases:
    if phrase.state:
      complete = phrase.plan.end == (phrase.plan.passage_end or phrase.plan.end)
      valid, error = generator.solo.harmonic.anchors_valid(phrase.anchors, phrase.state.notes, phrase.plan.end, complete=complete, closure=phrase.plan.closure)
      if not valid:
        raise ValueError(error)