"""Interpretable shared action ranking; no song identity or absolute bar features.

Complete actions compete by exact cost per quarter, with bounded backtracking.
The final report exposes the exact same contributions used by search.
Role-specific learned integer coefficients can augment these general priors.
"""
from collections import Counter
from fractions import Fraction as F
import dataclasses
from generator import vocabulary, harmony as chords, model
from generator.solo import expressions, scoring, types

VERSION='action-selection-1'
DEFAULT={'excessive_leap':12,'register_distance':1,'oscillation':6,'repeated_attack':3,
  'unprepared_space':12,'rest_after_rest':24,'source_relation':-2,'theme_relation':-4,
  'sustained_arrival':-1,'directed_motion':-1}


def validate_weights(weights):
  durations={str(F(i,4)).replace('/','_') for i in range(1,17)}|{'1_3','2_3'}
  kinds={'tone','color','color_resolution','anticipation','anticipated_color','approach','enclosure','connection','delayed_approach','space','rest','scalar','arpeggio','return','answer'}
  allowed={f'kind/{k}' for k in kinds}|{f'duration/{d}' for d in durations}
  allowed|={f'harmony/{q}/{p}/{d}' for q in ('major','minor','dominant','half_diminished','other') for p in range(12) for d in durations}
  allowed|={f'motion/{p}/{d}' for p in range(-12,13) for d in durations}
  allowed|={f'metric/{m}/{d}' for m in ('barline','beat','eighth','subdivision') for d in durations}
  allowed|={f'register/{i}' for i in range(-4,5)}
  allowed|={f'obligation/step/{i}' for i in range(-12,13)}
  allowed|={f'obligation/delay/{str(F(i,4)).replace("/","_")}' for i in range(1,21)}
  allowed|={f'obligation/hold/{d}' for d in durations}
  allowed|={f'continuation/{a}/{b}' for a in ('up','down','repeat') for b in ('up','down','repeat')}
  allowed|={'source/rhythm_match','theme/realized_identity',*(f'space/{s}' for s in ('unharmonized','breath','after_arrival','delayed_entrance'))}
  for key,value in weights.items():
    if key not in allowed or not isinstance(value,int) or isinstance(value,bool):
      raise ValueError('Unsupported fitted feature (only shared musical categories): '+key)
    if any(c not in 'abcdefghijklmnopqrstuvwxyz_0123456789/-' for c in key):
      raise ValueError('Noncanonical fitted feature: '+key)
  return dict(weights)


def signature(notes):
  if not notes:return ()
  first=notes[0]
  return tuple((n.start-first.start,n.end-n.start,n.pitch.number()-first.pitch.number()) for n in notes)


class Selector:
  def __init__(self,harmony,analysis,center,weights=None):
    self.harmony=harmony;self.analysis=analysis;self.center=center
    self.weights=validate_weights(weights if weights is not None else vocabulary.role().get('selection_weights',{}))
    self.source_intervals=Counter(b.pitch.number()-a.pitch.number() for a,b in zip(analysis.melody,analysis.melody[1:]))
    self.source_pairs=Counter((b.pitch.number()-a.pitch.number(),c.pitch.number()-b.pitch.number())
      for a,b,c in zip(analysis.melody,analysis.melody[1:],analysis.melody[2:]))
    self.source_rhythms=Counter(n.end-n.start for n in analysis.melody)
    self.static={}

  def base(self,action):
    if action.key in self.static:return self.static[action.key]
    values=Counter();prior=Counter();ticks=int((action.end-action.start)*12)
    values['kind/'+action.kind]=ticks
    for contract in action.contracts:
      if contract.target.start>=action.end and contract.events:
        # Pending destinations must be visible to selection. Identical heard
        # attacks with different future obligations are different actions.
        last=contract.events[-1];target=contract.target
        values[f'obligation/step/{max(-12,min(12,target.pitch.number()-last.pitch.number()))}']+=ticks
        values['obligation/delay/'+str(target.start-last.start).replace('/','_')]+=ticks
        values['obligation/hold/'+str(target.end-target.start).replace('/','_')]+=ticks
    if not action.notes:
      values['space/'+('unharmonized' if action.kind=='space' else 'breath')]=ticks
    for n in action.notes:
      duration=n.end-n.start;length=int(duration*12)
      label=str(duration).replace('/','_')
      values['duration/'+label]+=length
      beat=(n.start%1)
      metric='barline' if n.start+self.analysis.body_offset in self.analysis.boundaries else 'beat' if beat==0 else 'eighth' if beat==F(1,2) else 'subdivision'
      values['metric/'+metric+'/'+label]+=length
      band=max(-4,min(4,(n.pitch.number()-self.center+6)//12))
      values[f'register/{band}']+=length
      region=self.harmony.at(n.start)
      chord=expressions.parsed_chord(region.chord)
      root=int(model.NATURAL[chord['root']['letter']]+model.alteration(chord['root']['alter']))
      degree=(n.pitch.number()-root)%12
      values[f'harmony/{expressions.quality(chord)}/{degree}/{label}']+=length
      prior['register_distance']+=length*max(0,abs(n.pitch.number()-self.center)-7)
      if duration>=2 and n.pitch.number()%12 in self.harmony.members(n.start,n.end):
        prior['sustained_arrival']+=length
      values['source/rhythm_match']+=length if duration in self.source_rhythms else 0
    self.static[action.key]=(values,prior)
    if len(self.static)>8192:self.static.pop(next(iter(self.static)))
    return values,prior

  def features(self,action,recent=(),home=(),previous_kind=''):
    if any(n.continuation for n in action.notes):action=dataclasses.replace(action,notes=types.sounding_notes(action.notes))
    static,base=self.base(action);values=Counter(static);prior=Counter(base)
    if not action.notes:
      if action.kind=='space':return values,prior
      ticks=int((action.end-action.start)*12)
      if previous_kind in ('rest','space'):prior['rest_after_rest']+=ticks
      elif recent and recent[-1].end-recent[-1].start>=F(3,2):values['space/after_arrival']+=ticks
      elif recent:prior['unprepared_space']+=ticks
      else:values['space/delayed_entrance']+=ticks
      return values,prior
    previous=recent[-1] if recent else None
    intervals=[b.pitch.number()-a.pitch.number() for a,b in zip(recent,recent[1:])]
    for n in action.notes:
      length=int((n.end-n.start)*12)
      if previous:
        delta=n.pitch.number()-previous.pitch.number()
        token=max(-12,min(12,delta));label=str(n.end-n.start).replace('/','_')
        values[f'motion/{token}/{label}']+=length
        if intervals:
          # Direction changes are contextual, not a blanket leap prohibition.
          sign=lambda x: 'up' if x>0 else 'down' if x<0 else 'repeat'
          values[f'continuation/{sign(intervals[-1])}/{sign(delta)}']+=length
          if len(intervals)>=2 and intervals[-2]==-intervals[-1]==delta and delta:prior['oscillation']+=length
        prior['excessive_leap']+=length*max(0,abs(delta)-7)
        prior['repeated_attack']+=length if delta==0 and n.start==previous.end else 0
        # Ordered relationships, not the set of available intervals. Weight
        # by sounding duration so subdivisions cannot multiply a reward.
        if intervals and (intervals[-1],delta) in self.source_pairs:prior['source_relation']+=length
        if 0<abs(delta)<=2 and intervals and delta*intervals[-1]>0:prior['directed_motion']+=length
        intervals.append(delta)
      previous=n
    if action.kind=='return' and len(home)>=3:
      if signature(action.notes)==signature(home):
        recently_heard=any(signature(recent[i:i+len(home)])==signature(home)
          for i in range(max(0,len(recent)-len(home)+1)))
        if not recently_heard:prior['theme_relation']+=int((action.end-action.start)*12)
    if recent and home:
      joined=(*recent,*action.notes)
      count=min(len(home),len(joined))
      if count>=3 and signature(joined[-count:])==signature(home[:count]):
        values['theme/realized_identity']+=12
    return values,prior

  def evaluate(self,action,recent=(),home=(),previous_kind=''):
    features,prior=self.features(action,recent,home,previous_kind)
    contributions={'prior/'+k:v*DEFAULT[k] for k,v in prior.items() if v}
    contributions.update({k:v*self.weights.get(k,0) for k,v in features.items() if v})
    weights={'prior/'+k:v for k,v in DEFAULT.items()}|self.weights
    raw={'prior/'+k:v for k,v in prior.items()}|dict(features)
    span=action.end-action.start
    contributions={k:F(v)/span for k,v in contributions.items()}
    raw={k:F(v)/span for k,v in raw.items()}
    total=sum(contributions.values(),F())
    return scoring.Evaluation(VERSION,True,raw,weights,contributions,total,
      {'span':str(action.end-action.start),'ticks_per_quarter':12,'attacks':len(types.sounding_notes(action.notes))})

  def cost(self,action,recent=(),home=(),previous_kind=''):
    """Identical selection total without serializing every rejected breakdown."""
    features,prior=self.features(action,recent,home,previous_kind)
    numerator=sum(v*DEFAULT[k] for k,v in prior.items())
    numerator+=sum(v*self.weights.get(k,0) for k,v in features.items())
    return F(numerator)/(action.end-action.start)