"""Source-phrase variations: edit actual music, then realize complete routes.

The external lead references are audition controls, never production seeds.
All notes use sounding coordinates and all time uses quarter-note Fractions.
"""
from __future__ import annotations

import dataclasses
import fractions

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

F = fractions.Fraction


def signature(notes):
  return tuple((n.start, n.end, n.pitch.number()) for n in notes)


def shift(notes, octave):
  return tuple(dataclasses.replace(n, pitch=dataclasses.replace(n.pitch, octave=n.pitch.octave+octave),
    position=None, role='source', event_id=n.event_id or n.sources[0]) for n in notes)


def available(harmony, start, end, *, members=False):
  span = harmony.at(start)
  allowed = types.permitted(harmony.choices, start, end)
  if members:
    allowed &= harmony.members(start, end)
  if not span or not span.scale:
    return ()
  return tuple(p for i in range(-35, 28) if 41 <= (p := types.scale_pitch(span.scale, i)).number() <= 76 and p.number()%12 in allowed)


def adapt(notes, harmony):
  """Document required harmonic rewrites; never repair a sequence one step at a time."""
  result, edits = [], []
  for index, note in enumerate(notes):
    terminal = index == len(notes)-1
    choices = available(harmony,note.start,note.end,members=terminal or note.end-note.start >= 2)
    pieces = [(note.start,note.end,choices)]
    if not choices:
      pieces = [(a,b,available(harmony,a,b,members=terminal or b-a >= 2)) for _,a,b in harmony.segments(note)]
    for j,(start,end,domain) in enumerate(pieces):
      if not domain:
        if harmony.at(start) and harmony.at(start).scale:
          return (), ()
        edits.append(types.DraftEdit('missing_harmony',start,end,(note.event_id,),(), 'explicit no-chord space'))
        continue
      pitch = min(domain,key=lambda p:(abs(p.number()-note.pitch.number()),p.number()))
      identifier = note.event_id if j == 0 else note.event_id+f':harmonic:{j}'
      actual = dataclasses.replace(note,start=start,end=end,pitch=pitch,event_id=identifier)
      result.append(actual)
      if (start,end,pitch) != (note.start,note.end,note.pitch):
        edits.append(types.DraftEdit('harmonic_adaptation',start,end,(note.event_id,),(identifier,),
          f'{note.pitch.label()} to {pitch.label()}: sounding duration / {harmony.at(start).chord}'))
  if result:
    last=result[-1]
    domain=available(harmony,last.start,last.end,members=True)
    if not domain:
      return (),()
    if last.pitch.number()%12 not in harmony.members(last.start,last.end):
      pitch=min(domain,key=lambda p:(abs(p.number()-last.pitch.number()),p.number()))
      result[-1]=dataclasses.replace(last,pitch=pitch)
      edits.append(types.DraftEdit('harmonic_destination',last.start,last.end,(last.event_id,),(last.event_id,),
        f'{last.pitch.label()} to {pitch.label()}: final sounding arrival before release'))
  return tuple(result),tuple(edits)


def identifying(notes, source):
  """A complete heard figure, chosen by its musical content rather than its index."""
  if not notes:
    return ()
  count = min(4,max(1,len(notes)//2))
  windows = [notes[i:i+count] for i in range(len(notes)-count+1)]
  def rank(group):
    intervals = tuple(b.pitch.number()-a.pitch.number() for a,b in zip(group,group[1:]))
    rhythm = {(n.end-n.start) for n in group}
    # Favor a rhythm/turn with information; leave long holds available for development.
    return (-len(rhythm),-len(set(intervals)),sum(max(F(),n.end-n.start-1) for n in group),group[0].start)
  group = min(windows,key=rank)
  return (types.IdentifyingFeature(source+':figure',tuple(n.event_id for n in group),
    tuple((n.start-group[0].start,n.end-n.start) for n in group),
    tuple(b.pitch.number()-a.pitch.number() for a,b in zip(group,group[1:])),source),)


def variations(notes, source, harmony):
  """One purposeful timing or pitch edit; no generic replacement openings."""
  result=[]
  if source.pickup and notes and source.start==0:
    length=-source.pickup[0].start
    if notes[0].end-length>=F(1,2):
      pickup=tuple(dataclasses.replace(n,start=n.start+length,end=n.end+length) for n in source.pickup)
      changed=(*pickup,dataclasses.replace(notes[0],start=length),*notes[1:])
      if harmony.evaluate(changed,complete=True).valid:
        result.append((changed,(types.DraftEdit('pickup_entrance',F(),notes[0].end,(notes[0].event_id,),
          tuple(n.event_id for n in pickup)+(notes[0].event_id,),source.id+': bring source pickup inside the solo and shorten its first hold'),)))
  for i,(before,note) in enumerate(zip(notes,notes[1:]),1):
    if before.end == note.start and before.end-before.start >= 1 and note.end-note.start >= F(1,2):
      for displacement in (F(-1,2),F(1,2)):
        time = note.start+displacement
        if time-before.start < F(1,2) or note.end-time < F(1,2):
          continue
        changed=(*notes[:i-1],dataclasses.replace(before,end=time),dataclasses.replace(note,start=time),*notes[i+1:])
        if harmony.evaluate(changed,complete=True).valid:
          edit=types.DraftEdit('rhythmic_paraphrase',before.start,note.end,(before.event_id,note.event_id),(before.event_id,note.event_id),
            f'{source.id}: retain both pitches and their order; move shared boundary {displacement} beat')
          result.append((changed,(edit,)))
    # Structural color is chosen in relation to the original turn, never a new random contour.
    if 0 < i < len(notes)-1:
      after=notes[i+1]
      for p in available(harmony,note.start,note.end,members=note.end-note.start>=1):
        if p == note.pitch or abs(p.number()-note.pitch.number())>2:
          continue
        direction=lambda a,b:(b>a)-(b<a)
        if (direction(before.pitch.number(),p.number()),direction(p.number(),after.pitch.number())) != (direction(before.pitch.number(),note.pitch.number()),direction(note.pitch.number(),after.pitch.number())):
          continue
        changed=(*notes[:i],dataclasses.replace(note,pitch=p),*notes[i+1:])
        if harmony.evaluate(changed,complete=True).valid:
          edit=types.DraftEdit('melodic_paraphrase',note.start,note.end,(note.event_id,),(note.event_id,),f'{source.id}: {note.pitch.label()} to {p.label()}, retain rhythm and turn')
          result.append((changed,(edit,)))
  # Broaden or release an ending as an explicit whole-phrase alternative.
  if notes and notes[-1].end-notes[-1].start >= F(3,2):
    tail=dataclasses.replace(notes[-1],end=notes[-1].end-F(1,2))
    result.append(((*notes[:-1],tail),(types.DraftEdit('ending_release',tail.end,notes[-1].end,(tail.event_id,),(tail.event_id,),source.id+': half-beat breath after the arrival'),)))
  return tuple(result)


def rhythm_catalog(library, meter):
  # Keep complete snippet objects. Slices below retain real rests and durations.
  return drafts.catalog(library,meter)


def connection_patterns(start,end,entry,destination,rhythms,meter):
  length=end-start
  spans={( (F(),length), ):'held'}
  for identifier,pattern,total in rhythm_catalog(rhythms,meter):
    if total < length:
      continue
    # Retrieve every exact-duration window, including offbeat and cross-bar windows.
    for offset in sorted({F(),*(a for a,_ in pattern)}):
      if offset+length>total:
        continue
      if any(a<offset<b or a<offset+length<b for a,b in pattern):
        continue
      selected=tuple((a-offset,b-offset) for a,b in pattern if offset<=a and b<=offset+length)
      if 2<=len(selected)<=7 and selected[0][0]==0:
        spans.setdefault(selected,identifier)
  ranked=sequences.diverse(tuple(spans.items()),
    lambda item:(abs(len(item[0])-float(length)*1.5),item[1]),
    lambda item:(len(item[0]),any(a[1]!=b[0] for a,b in zip(item[0],item[0][1:]))),8)
  result=[]
  for pattern,identifier in ranked:
    events=tuple(dataclasses.replace(entry,start=start+a,end=start+b,role='connection',
      event_id=f'{entry.event_id}:route:{a}',position=None) for a,b in pattern)
    if len(events)>=2:
      result.append(((*events,destination),identifier))
  return result


def elaborate(notes, features, harmony, library, rhythms, meter, source, *, sustain_entry=False):
  """Edit an exposed held tail or a complete interior span, with both endpoints known."""
  fixed={e for f in features for e in f.events}
  windows=[]
  for i,(note,target) in enumerate(zip(notes,notes[1:])):
    if note.event_id in fixed or note.end != target.start:
      continue
    if note.end-note.start>=F(3,2):
      start=max(note.start+F(1,2),target.start-3)
      windows.append((i,start,target,False))
  # A source-held ending can arrive later via a related in-scale turn.
  if notes and notes[-1].event_id not in fixed and notes[-1].end-notes[-1].start>=F(5,2):
    tail=notes[-1]
    for duration in (F(1,2),F(1),F(3,2),F(2)):
      target=dataclasses.replace(tail,start=tail.end-duration,event_id=tail.event_id+':arrival')
      start=max(tail.start+F(1,2),target.start-3)
      if target.start-start>=1:
        windows.append((len(notes)-1,start,target,True))
  result=[]
  for i,start,target,terminal in windows:
    entry=notes[i];end=target.start
    prefix=dataclasses.replace(entry,end=start)
    seed=dataclasses.replace(entry,start=start,event_id=entry.event_id+':tail',role='connection')
    for pattern,rhythm_id in connection_patterns(start,end,seed,target,rhythms,meter):
      for route in routes.candidates(pattern,{target.event_id},harmony,library,(41,76),'develop '+source.id):
        if sustain_entry:
          # The route owns its complete sounding entry, including the old hold.
          joined=dataclasses.replace(route.notes[0],start=entry.start,event_id=entry.event_id)
          route=dataclasses.replace(route,notes=(joined,*route.notes[1:]))
          actual=(*notes[:i],*route.notes,*notes[i+1 if terminal else i+2:])
        else:
          actual=(*notes[:i],prefix,*route.notes,*notes[i+1 if terminal else i+2:])
        if not harmony.evaluate(actual,complete=True).valid:
          continue
        if any(abs(b.pitch.number()-a.pitch.number())>7 for a,b in zip(route.notes,route.notes[1:])):
          continue
        route=dataclasses.replace(route,sources=(*route.sources,'rhythm:'+rhythm_id))
        edit=types.DraftEdit('complete_route',start,target.start,(entry.event_id,),tuple(n.event_id for n in route.notes),
          source.id+'; '+rhythm_id+'; '+('later terminal arrival on ' if terminal else 'lead into ')+target.pitch.label())
        result.append((actual,edit,route))
  # Whole routes compete; bounded diversity retains different locations and rhythms.
  return sequences.diverse(result,lambda v:(abs(len(v[0])*4/float(source.end-source.start)-5),
    sum(abs(b.pitch.number()-a.pitch.number()) for a,b in zip(v[2].notes,v[2].notes[1:])),v[2].id),
    lambda v:(v[1].start,v[2].destination.end-v[2].destination.start,len(v[2].notes)),8)


def make_phrase(source, original, refined, notes, edits, route_list, octave, part, harmony, tempo, seed):
  from generator.solo import realization
  features=identifying(refined,source.id)
  # Only fully retained source turns become committed anchor events.
  anchors=[]
  for i,n in enumerate(notes):
    if n.pitch.number()%12 not in harmony.members(n.start,n.end):
      continue
    role='terminal' if i==len(notes)-1 else 'opening' if i==0 else 'intermediate'
    if role=='intermediate' and n.end-n.start<2:
      continue
    anchors.append(types.Anchor(n.event_id+':anchor',n.start,n.start,n.pitch,harmony.members(n.start,n.end),
      (41,76),n.end-n.start,role,harmony.at(n.start).chord,hold_until=n.end))
  if notes and (not anchors or anchors[-1].role!='terminal'):
    return None
  identifier=part+'_'+source.id+'_'+policy.stable_rank(seed,signature(notes))[:12]
  motif=types.Motif(source.id,tuple(dataclasses.replace(n,start=n.start-source.start,end=n.end-source.start) for n in refined),
    source.end-source.start,tuple(s.chord for s in harmony.spans if s.start<source.end and s.end>source.start),'song melody '+source.family)
  home=realization.make_identity(identifier,source.id,refined,source.start,source.end,'source_variation',harmony)
  transform=types.Transformation(identifier,source.id,source.family,'source_variation',source.start,source.end,
    tuple(n.event_id for n in notes),tuple((n.start-source.start,n.end-source.start) for n in notes),preserves=(),source=source.id)
  draft=types.MelodicDraft(identifier,source.start,source.end,notes,(transform,),
    tuple(a.id.removesuffix(':anchor') for a in anchors),(),tuple(anchors),home,'source_variation',features=features,
    reference=source.id,reference_motif=motif,source_phrase=source,octave_transfer=octave,
    source_links=tuple((n.event_id,n.sources) for n in notes))
  spaces=[];cursor=source.start
  for n in notes:
    if n.start>cursor:spaces.append(types.ActivitySpan(cursor,n.start,'space',F(),F()))
    spaces.append(types.ActivitySpan(n.start,n.end,'active',F(5,4),F()))
    cursor=n.end
  if cursor<source.end:spaces.append(types.ActivitySpan(cursor,source.end,'space',F(),F()))
  if not notes:spaces=[types.ActivitySpan(source.start,source.end,'space',F(),F())]
  plan=types.PhrasePlan(identifier,source.start,source.end,motif,'source_variation','source_variation',None,(41,76),F(5,4),
    source.end-cursor,tuple(spaces),(), 'breath' if cursor<source.end else 'sustain','closed',seed,source.id,source.start,source.end, (tuple(anchors),))
  phrase=realization.checked_phrase(plan,draft,notes,(),types.ContinuationState(''),harmony,tempo,edits)
  if phrase is None:
    return None
  # Source correspondence is evaluated once. Do not score the imported tune's
  # characteristic leaps as generic contour defects or reward filling its rests.
  source_by_id={n.event_id:n for n in original}
  retained={n.event_id:n for n in notes if n.event_id in source_by_id}
  deviation=sum(min(1,abs(n.pitch.number()-retained[k].pitch.number())/5)+
    min(1,float(abs(n.start-retained[k].start))/2)+min(1,float(abs((n.end-n.start)-(retained[k].end-retained[k].start)))/2) for k,n in source_by_id.items() if k in retained)/max(1,len(source_by_id))
  exposure=harmony.evaluate(notes,complete=True).exposure
  density=len(notes)*4/max(.001,float(source.end-source.start))
  score=3*deviation+2*exposure+.35*max(0,4-density,density-6)
  source_draft=dataclasses.replace(draft,id=identifier+':source',notes=original,features=(),anchors=(),structural=())
  refined_draft=dataclasses.replace(draft,id=identifier+':variation',notes=refined,features=features,anchors=(),structural=())
  return dataclasses.replace(phrase,score=score,initial=source_draft,refined=refined_draft,routes=route_list)


def candidates(source, harmony, sequences_library, rhythms, meter, part, tempo, seed, *, stage='elaborated'):
  result=[]
  for octave in range(-3,4):
    original=shift(source.notes,octave)
    if any(not generator.fretboard.single_positions_for_pitch(F(n.pitch.number())) for n in original):
      continue
    adapted,required=adapt(original,harmony)
    if original and not adapted and any(harmony.at(n.start) and harmony.at(n.start).scale for n in original):
      continue
    variants=list(variations(adapted,dataclasses.replace(source,pickup=shift(source.pickup,octave)),harmony))
    if required or not adapted:
      variants.append((adapted,()))
    # A whole preserved phrase is the baseline, not an accidental production fallback.
    if not variants:
      continue
    variants=sequences.diverse(variants,lambda v:(len(v[1]),policy.stable_rank(seed,signature(v[0]))),
      lambda v:(v[1][0].kind if v[1] else 'harmonic_adaptation',v[1][0].start if v[1] else source.start),8)
    for refined,changes in variants:
      options=[(refined,(),())]
      if stage=='elaborated':
        connections=elaborate(refined,identifying(refined,source.id),harmony,sequences_library,rhythms,meter,source)
        options.extend((notes,(edit,),(route,)) for notes,edit,route in connections)
        # A second, distinct source-held span may develop the same phrase.
        # Reuse complete routes; never elaborate an already elaborated event.
        for first_notes,first_edit,first_route in connections[:2]:
          fixed=(*identifying(refined,source.id),types.IdentifyingFeature('route',
            tuple(n.event_id for n in first_route.notes)+first_edit.removed,(),(),source.id))
          second=elaborate(first_notes,fixed,harmony,sequences_library,rhythms,meter,source)
          options.extend((actual,(first_edit,edit),(first_route,route)) for actual,edit,route in second[:2])
      for notes,added,route_list in options:
        phrase=make_phrase(source,original,refined,notes,(*required,*changes,*added),route_list,octave,part,harmony,tempo,seed)
        if phrase:
          result.append(phrase)
  # Preserve register AND edit-location diversity, rather than eight near-copies.
  unique={signature(p.draft.notes):p for p in result}
  return sequences.diverse(tuple(unique.values()),lambda p:(p.score,policy.stable_rank(seed,p.draft.id)),
    lambda p:p.draft.octave_transfer,policy.DEFAULT.source_variations)


def clipped_sources(analysis,end):
  result=[]
  for source in analysis.source_phrases:
    if source.start>=end:break
    finish=min(source.end,end)
    notes=tuple(dataclasses.replace(n,start=max(F(),n.start),end=min(n.end,finish)) for n in source.notes if n.start<finish and n.end>0)
    result.append(dataclasses.replace(source,end=finish,notes=notes,rests=tuple((a,min(b,finish)) for a,b in source.rests if a<finish)))
  return tuple(result)


def transition(first,second):
  if not first or not first.decisions or not second.decisions:return 0.
  a=first.state.last;b=second.decisions[0].notes[0]
  leap=max(0,abs(a.pitch.number()-b.pitch.number())-7)/12
  same_family=first.draft.source_phrase.family==second.draft.source_phrase.family
  repeated_end=not same_family and first.anchors and second.anchors and first.anchors[-1].pitch==second.anchors[-1].pitch and first.anchors[-1].duration==second.anchors[-1].duration
  return .2*leap+.15*bool(repeated_end)


def solo_options(analysis,measures,part,end,choices,library,rhythms,context,stage,mode='thematic'):
  from generator.solo import realization, themes
  if mode not in ('thematic','paraphrase'):raise ValueError('Unknown solo treatment: '+mode)
  context.resources[id(library)]=library
  context.resources[id(rhythms)]=rhythms
  harmony=harmonic.Harmony(choices)
  sources=themes.windows(analysis,measures,end,harmony) if mode=='thematic' else clipped_sources(analysis,end)
  seed=policy.stable_rank('source_variations', (analysis.melody,choices,policy.SEED_NAMESPACE))
  states=[(0.,())]
  for source in sources:
    meter=measures[min(len(measures)-1,int(source.start//generator.model.fraction(measures[0]['duration'])))]['time_signature']
    following=[]
    for cost,history in states:
      incoming=history[-1].state if history else types.ContinuationState('')
      key=(mode,source,incoming if mode=='thematic' else None,policy.stable_rank("harmony",choices),part,analysis.tempo,seed,stage,id(library),id(rhythms),policy.DEFAULT)
      if key in context.results:
        pool=context.results[key];context.diagnostics['cache_hits']+=1
      else:
        if mode=='thematic' and any(c.scale and c.span.start<source.end and c.span.end>source.start for c in choices):
          pool=()
          for batch in range(2):
            pool=themes.phrase_candidates(source,analysis,incoming,part,harmony,library,rhythms,meter,analysis.tempo,context,seed,stage=stage,batch=batch,final=source.end==end)
            if pool:break
            context.diagnostics['theme_replans']=context.diagnostics.get('theme_replans',0)+1
        elif mode=='thematic':
          pool=themes.silence(source,incoming,part,harmony,analysis.tempo,seed)
        else:
          pool=candidates(source,harmony,library,rhythms,meter,part,analysis.tempo,seed,stage=stage)
        context.results[key]=pool;context.diagnostics['generated_candidates']+=len(pool)
      context.diagnostics.setdefault('source_candidate_pools', {}).setdefault(source.id,[]).extend(
        {'id':p.draft.id,'score':p.score,'octave':p.draft.octave_transfer,'attacks':len(p.draft.notes),'edits':[e.kind for e in p.edits],'source_family':source.family,'theme':p.state.home.id if p.state.home else None} for p in pool)
      for candidate in pool:
        # The thematic constructor has already placed all decisions from the
        # actual incoming state; do not collapse its motif groups into one bar.
        if mode=='thematic':
          total=cost+candidate.score*float(source.end-source.start)/float(end)+transition(history[-1] if history else None,candidate)
          following.append((total,(*history,candidate)))
          continue
        last=incoming.last
        placements=placement.alternatives(tuple(n for d in candidate.decisions for n in d.notes),analysis.tempo,last,limit=2)
        for travel,notes in placements:
          decision=dataclasses.replace(candidate.decisions[0],notes=notes) if candidate.decisions else None
          state=dataclasses.replace(candidate.state,notes=notes,last=notes[-1] if notes else last,travel=travel)
          placed=dataclasses.replace(candidate,decisions=(decision,) if decision else (),state=state)
          total=cost+candidate.score*float(source.end-source.start)/float(end)+.05*travel/max(1,len(notes))/2000+transition(history[-1] if history else None,placed)
          following.append((total,(*history,placed)))
    if not following:
      raise realization.SearchExhausted(f'{part} {source.id}: {mode} feature, harmonic destination or continuation failed within the bounded search; not proof of incompatibility')
    states=sequences.diverse(following,lambda v:(v[0],policy.stable_rank(seed,tuple((p.draft.id,p.state.last.position if p.state.last else None) for p in v[1]))),
      lambda v:(v[1][-1].state.home.id if mode=='thematic' and v[1][-1].state.home else tuple(p.draft.octave_transfer for p in v[1])),policy.DEFAULT.chorus_frontier)
  return tuple((cost,types.Solo(part,choices,phrases,analysis.tempo,end,{**context.diagnostics,'construction':'thematic_recomposition' if mode=='thematic' else 'source_phrase_variations','policy':policy.describe()['fingerprint']},
    tuple(types.MelodicPassagePlan(p.plan.passage_id,p.plan.start,p.plan.end,(p.plan.id,),seed) for p in phrases))) for cost,phrases in states)


def similarity(first,second):
  """Compare composition independently of octave and fingering; callbacks are explicit."""
  def events(solo):
    callbacks={e for p in solo.phrases if p.draft for f in p.draft.features if f.id.endswith(':callback') for e in f.events}
    return tuple(n for p in solo.phrases if p.plan.relationship!='callback' for d in p.decisions for n in d.notes if n.event_id not in callbacks)
  a,b=events(first),events(second)
  def tokens(notes):
    return {(n.start,n.end,n.pitch.number()%12) for n in notes}
  left,right=tokens(a),tokens(b)
  literal=len(left&right)/max(1,min(len(left),len(right)))
  def runs(notes):
    return {tuple((n.start-notes[i].start,n.end-n.start,n.pitch.number()-notes[i].pitch.number()) for n in notes[i:i+4]) for i in range(len(notes)-3)}
  left,right=runs(a),runs(b)
  interval=len(left&right)/max(1,min(len(left),len(right)))
  return .65*literal+.35*interval