"""Bounded target-driven ornaments after strict thematic pair selection.

The main tune and its home identity remain strictly mapped. Every exception is
an atomic upper/lower/target gesture, never a replacement policy for single notes.
"""
from __future__ import annotations

import dataclasses
import math
import time
from fractions import Fraction as F

import generator.model
from generator.solo import types, harmonic, placement, policy, drafts, routes, phrasing


def mapped_pitch(number, start, end, harmony):
  if number%12 not in types.permitted(harmony.choices,start,end):return None
  span=harmony.at(start)
  return types.scale_pitch(span.scale,types.scale_index(span.scale,number))


def approach(target, harmony):
  """Spelling follows the sounding collection, or explicitly spells resolution."""
  t=target.start; number=target.pitch.number()
  upper=next((p for n in range(number+1,number+4) if (p:=mapped_pitch(n,t-1,t-F(1,2),harmony))),None)
  if upper is None:return None
  lower=mapped_pitch(number-1,t-F(1,2),t,harmony)
  if lower is None:
    if target.pitch.alter>0:
      lower=dataclasses.replace(target.pitch,alter=target.pitch.alter-1)
    else:
      letters='cdefgab'; index=letters.index(target.pitch.letter)
      letter=letters[(index-1)%7]; octave=target.pitch.octave-(index==0)
      alter=number-1-generator.model.NATURAL[letter]-12*(octave+1)
      lower=types.Pitch(letter,F(alter),octave)
  if not all(placement.positions(p.number()) for p in (upper,lower)):return None
  identity=target.event_id+':enclosure'
  return (types.Note(t-1,t-F(1,2),upper,target.sources,'enclosure_upper',event_id=identity+':upper'),
    types.Note(t-F(1,2),t,lower,target.sources,'enclosure_lower',event_id=identity+':lower'))


def declared_targets(phrase,harmony):
  """Declare eligible existing stops before proposing their ornaments."""
  if not phrase.draft:return ()
  terminal={a.pitch.number():a for a in phrase.anchors if a.role=='terminal'}
  values=[]
  held=phrase.draft.notes[0].event_id if phrase.draft.notes and phrase.draft.notes[0].continuation else None
  for note in drafts.local_sounding(phrase.draft.notes):
    if note.event_id==held:continue
    if note.continuation or note.role=='quote' or note.end-note.start<1:continue
    if note.pitch.number()%12 not in harmony.members(note.start,note.end):continue
    anchor=terminal.get(note.pitch.number())
    role='terminal' if anchor and anchor.earliest<=note.start<=anchor.latest else 'intermediate'
    rank=(role!='terminal',phrasing.destination_cost(note.pitch,note.start,note.end,harmony,terminal=role=='terminal'),-(note.end-note.start),note.start)
    values.append((rank,note,role))
  return tuple((n,role) for _,n,role in sorted(values)[:policy.DEFAULT.enclosure_targets])


def replace_window(notes,start,end,replacement):
  """Only rewrite continuously sounding material; preserve meaningful space."""
  cursor=start;prefix=[];suffix=[]
  for note in notes:
    if note.end<=start:prefix.append(note);continue
    if note.start>=end:suffix.append(note);continue
    if note.start>cursor:return None
    cursor=max(cursor,note.end)
    if note.start<start:prefix.append(dataclasses.replace(note,end=start))
    if note.end>end:return None  # never clip a pre-existing target sustain
  if cursor<end:return None
  return (*prefix,*replacement,*suffix)


def editable(draft,start,end):
  ids=set((*draft.structural,*draft.protected,*(e for f in draft.features for e in f.events)))
  return not any(n.start<end and n.end>start and (n.event_id in ids or n.role=='quote') for n in draft.notes) and not any(t.start<end and t.end>start for t in draft.sustained_tails)


def movement(notes):
  return sum(max(0,abs(b.pitch.number()-a.pitch.number())-3)/12 for a,b in zip(notes,notes[1:]))/max(1,len(notes)-1)


def candidates(phrase,harmony,library,tempo,rejections):
  draft=phrase.draft
  if not draft:return ()
  if draft.notes and draft.notes[0].continuation:
    # This local budget does not reopen a committed preceding phrase. Never
    # invent a new attack by detaching its already-sounding entrance.
    rejections.append({'target':draft.notes[0].event_id,'reason':'cross-phrase sustain retained; no local rewrite of its context'})
    return ()
  notes=types.sounding_notes(draft.notes);result=[]
  for target,role in declared_targets(phrase,harmony):
    start=target.start-1
    if start<phrase.plan.start or not editable(draft,start,target.start):
      rejections.append({'target':target.event_id,'reason':'protected material or no editable approach window'});continue
    pair=approach(target,harmony)
    if not pair:
      rejections.append({'target':target.event_id,'reason':'no mapped upper neighbor or playable approach'});continue
    upper,lower=pair
    contract=types.EnclosureContract(target.event_id+':enclosure',upper,lower,target,harmony.at(target.start).chord,role,
      types.SustainedTail(target.start,target.end,target.pitch.number()))
    proposals=[(start,pair,None)]
    # Reopen up to two beats before the approach as a whole in-map route. Keep
    # the incoming attack and fixed upper endpoint, not independently chosen steps.
    for first in reversed(notes):
      if not start-2<=first.start<start or first.end>start or not editable(draft,first.end,target.start):continue
      templates=tuple(n for n in notes if first.start<=n.start<start)
      if not templates or any(a.end!=b.start for a,b in zip(templates,templates[1:])) or templates[-1].end!=start:continue
      for route in routes.candidates((*templates,upper),{upper.event_id},harmony,library,(41,76),'mapped approach into enclosure')[:policy.DEFAULT.enclosure_routes]:
        proposals.append((first.start,(*route.notes,lower),route))
      break
    accepted=0
    for begin,replacement,route in proposals[:1+policy.DEFAULT.enclosure_routes]:
      actual=replace_window(notes,begin,target.start,replacement)
      if actual is None:continue
      valid,error=drafts.contract_valid(draft,actual)
      if not valid:continue
      checked=harmony.evaluate(actual,phrase.state.treatments if phrase.state else (),until=phrase.plan.end,complete=True,enclosures=(contract,))
      if not checked.valid:continue
      valid,error=harmonic.anchors_valid(phrase.anchors,actual,phrase.plan.end,complete=True,closure=phrase.plan.closure)
      if not valid:continue
      placed=placement.place(actual,tempo)
      if placed is None:continue
      old_travel=placement.place(notes,tempo)[0]
      delta=.6*(movement(actual)-movement(notes))+.025*(placed[0]-old_travel)/max(1,len(notes))/2000
      density=4*len(actual)/float(phrase.plan.end-phrase.plan.start)
      before_density=4*len(notes)/float(phrase.plan.end-phrase.plan.start)
      delta+=.1*(max(0,density-6)-max(0,before_density-6))
      transforms=tuple(dataclasses.replace(t,events=tuple(n.event_id for n in actual if t.start<=n.start<t.end),
        rhythm=tuple((n.start-t.start,n.end-t.start) for n in actual if t.start<=n.start<t.end)) for t in draft.transformations)
      changed=dataclasses.replace(draft,notes=actual,transformations=transforms,enclosures=(contract,))
      decisions=tuple(dataclasses.replace(d,notes=tuple(n for n in actual if d.start<=n.start<d.end)) for d in phrase.decisions)
      if tuple(n for d in decisions for n in d.notes)!=actual:continue
      retained=tuple(r for r in phrase.routes if all(any((n.event_id,n.start,n.end,n.pitch)==(a.event_id,a.start,a.end,a.pitch) for a in actual) for n in r.notes))
      result.append((delta,dataclasses.replace(phrase,draft=changed,decisions=decisions,before_enclosure=draft,
        routes=(*retained,*((route,) if route else ())))))
      accepted+=1
    if not accepted:rejections.append({'target':target.event_id,'reason':'no valid complete enclosure within local route budget'})
  return tuple(result)


def strict(solo):
  """Recover the exact selected pre-enclosure notes, including their placement."""
  phrases=[]
  for phrase in solo.phrases:
    draft=phrase.before_enclosure
    if not draft:phrases.append(phrase);continue
    decisions=tuple(dataclasses.replace(d,notes=tuple(n for n in draft.notes if d.start<=n.start<d.end)) for d in phrase.decisions)
    state=dataclasses.replace(phrase.state,notes=draft.notes,last=draft.notes[-1] if draft.notes else phrase.state.last) if phrase.state else None
    phrases.append(dataclasses.replace(phrase,draft=draft,decisions=decisions,state=state))
  return dataclasses.replace(solo,phrases=tuple(phrases),diagnostics={**solo.diagnostics,'enclosures_enabled':False})


def refine(solo,library,enabled=True):
  started=time.monotonic();harmony=harmonic.Harmony(solo.choices);rejections=[]
  baseline=[]
  for phrase in solo.phrases:
    draft=dataclasses.replace(phrase.draft,notes=tuple(n for d in phrase.decisions for n in d.notes)) if phrase.draft else None
    baseline.append(dataclasses.replace(phrase,draft=draft,before_enclosure=draft))
  baseline=tuple(baseline)
  if not enabled:
    return dataclasses.replace(solo,phrases=baseline,diagnostics={**solo.diagnostics,'enclosures_enabled':False,'enclosure_seconds':time.monotonic()-started})
  maximum=math.ceil(policy.DEFAULT.enclosure_max_per_32_bars*float(solo.end)/128)
  low,high=2*float(solo.end)/128,3*float(solo.end)/128
  states=[(0.,(),())];options_count=0
  def sparse(count):return .10*max(0,low-count)+.12*max(0,count-high)
  for phrase in baseline:
    options=((0.,phrase),*candidates(phrase,harmony,library,solo.tempo,rejections));options_count+=len(options)-1
    following=[]
    for cost,chosen,arrivals in states:
      for delta,option in options:
        enclosure=option.draft.enclosures[0] if option.draft and option.draft.enclosures else None
        if enclosure and len(arrivals)>=maximum:continue
        spacing=.08*max(0,1-float(enclosure.target.start-arrivals[-1][0])/16) if enclosure and arrivals else 0.
        repeated=.06*sum(pc==enclosure.target.pitch.number()%12 for _,pc in arrivals) if enclosure else 0.
        history=(*arrivals,(enclosure.target.start,enclosure.target.pitch.number()%12)) if enclosure else arrivals
        following.append((cost+delta+spacing+repeated,(*chosen,option),history))
    states=sorted(following,key=lambda x:(x[0]+sparse(len(x[2])),policy.stable_rank('enclosure-selection',x[2])))[:policy.DEFAULT.enclosure_frontier]
  _,phrases,arrivals=min(states,key=lambda x:(x[0]+sparse(len(x[2])),policy.stable_rank('enclosure-selection',x[2])))
  # Resolve the complete guitar path once, including both sides of phrase joins.
  raw=tuple(n for p in phrases for d in p.decisions for n in d.notes)
  placed=placement.place(raw,solo.tempo)
  if placed is None:raise ValueError('selected enclosure solo has no complete guitar path')
  by_id={n.event_id:n for n in placed[1]};final=[]
  for phrase in phrases:
    decisions=tuple(dataclasses.replace(d,notes=tuple(by_id[n.event_id] for n in d.notes)) for d in phrase.decisions)
    notes=tuple(n for d in decisions for n in d.notes)
    state=dataclasses.replace(phrase.state,notes=notes,last=notes[-1] if notes else phrase.state.last) if phrase.state else None
    final.append(dataclasses.replace(phrase,decisions=decisions,draft=dataclasses.replace(phrase.draft,notes=notes) if phrase.draft else None,state=state))
  diagnostics={**solo.diagnostics,'enclosures_enabled':True,'enclosure_seconds':time.monotonic()-started,
    'enclosure_candidates':options_count,'enclosure_rejections':rejections,'enclosures_selected':len(arrivals)}
  return dataclasses.replace(solo,phrases=tuple(final),diagnostics=diagnostics)