from __future__ import annotations

import dataclasses
import fractions
import time

import generator.solo.enclosures

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, revise=True, context=None, references=(), stage='elaborated', mode='thematic', enclosures_enabled=True):
  started = time.monotonic()
  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')
  from generator.solo import leads
  context = context if context is not None else generator.solo.realization.SearchContext()
  options = leads.solo_options(analysis, measures, part, end, choices, sequences, rhythms, context, stage, mode)
  result = options[0][1]
  result.diagnostics['strict_composition_seconds'] = time.monotonic()-started
  result = generator.solo.enclosures.refine(result,sequences,enclosures_enabled and stage=='elaborated' and mode=='thematic')
  validate(result)
  return result


def compose_pair(analysis, requests, *, scales, sequences, rhythms, contexts=None, stage='elaborated', mode='thematic', enclosures_enabled=True):
  """Select two recognizable interpretations together, before arranging backing."""
  from generator.solo import leads
  started = time.monotonic()
  pools = {}
  for part, (measures,end,initial) in requests.items():
    timeline = generator.scales.timeline.read({'measures': measures},initial=initial)
    choices = generator.scales.selection.select(timeline,generator.scales.timeline.Solo(0,end,part),scales)
    for c in choices:
      if c.span.chord and c.scale is None:
        raise ValueError(f'{part} solo, source measure {c.span.source_measure}: No exact match within the permitted catalog')
    context = (contexts or {}).get(part) or generator.solo.realization.SearchContext()
    pools[part] = leads.solo_options(analysis,measures,part,end,choices,sequences,rhythms,context,stage,mode)
  def theme_key(solo):
    home=next((p.state.home for p in solo.phrases if p.state and p.state.home),None)
    if not home or not home.notes:return ()
    return tuple((n.start-home.notes[0].start,n.end-n.start,n.pitch.number()-home.notes[0].pitch.number()) for n in home.notes)
  pairs = [(a_cost+e_cost+2*leads.similarity(a,e),a,e) for a_cost,a in pools['acoustic'] for e_cost,e in pools['electric']
    if mode!='thematic' or not theme_key(a) or theme_key(a)!=theme_key(e)]
  if not pairs:
    raise generator.solo.realization.SearchExhausted('Distinct acoustic/electric themes not found within the bounded pair search')
  _,acoustic,electric = min(pairs,key=lambda v:(v[0],generator.solo.policy.stable_rank('solo-pair',(v[1].notes(),v[2].notes()))))
  overlap = leads.similarity(acoustic,electric)
  strict_seconds=time.monotonic()-started
  selected=[]
  for solo in (acoustic,electric):
    solo.diagnostics.update(pair_similarity=overlap,pairs_compared=len(pairs),solo_alternatives=len(pools[solo.part]))
    solo.diagnostics['strict_pair_seconds']=strict_seconds
    solo=generator.solo.enclosures.refine(solo,sequences,enclosures_enabled and stage=='elaborated' and mode=='thematic')
    validate(solo)
    selected.append(solo)
  return tuple(selected)


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, revise_pair=None):
  """Commit only after the next phrase has been evaluated from each ending."""
  result, previous, pending = [], None, None
  for index, plan in enumerate(plans):
    options = realize(plan, previous)
    if pending is not None and pending not in options:
      options = (pending, *options)
    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, None))
        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)
      for future in following:
        left, right = revise_pair(option, future, previous) if revise_pair else (option, future)
        cost = (right.score if shared_passage else (left.score * length + right.score * next_length) / (length + next_length)) + transition_cost(left, right)
        ranked.append((cost + incoming, left, right))
    if not ranked:
      raise failure or generator.solo.realization.SearchExhausted(f'{plan.id}: no feasible phrase transition')
    _, chosen, pending = 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.position not in generator.fretboard.single_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))
  enclosures = tuple(e for p in solo.phrases if p.draft for e in p.draft.enclosures)
  checked = harmony.evaluate(solo.notes(), contracts, until=solo.end, complete=True,enclosures=enclosures)
  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)