"""Canonical sounded-event evaluation, shared by inference and training.

Constructor names are not musical evidence. Complete equal-span candidates
are compared using this record; local retrieval is explicitly only a heuristic.
"""
from collections import Counter
from fractions import Fraction as F
import dataclasses
import functools
from types import MappingProxyType
from generator import vocabulary, harmony as chords, model
from generator.solo import expressions, scoring, types

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

CONDITIONS={
  'quality':frozenset(('major','minor','dominant','half_diminished','other')),
  'degree':frozenset(map(str,range(12))),
  'motion':frozenset(map(str,range(-12,13))),
  'metric':frozenset(('barline','beat','eighth','subdivision')),
  'duration':frozenset(str(F(i,4)).replace('/','_') for i in range(1,17))|{'1_3','2_3'},
  'source_pair':frozenset(('yes','no')),
}

def rule_conditions(key):
  if not key.startswith('joint/'):return None
  pairs=tuple(p.split('=',1) for p in key[6:].split('&'))
  if not 2<=len(pairs)<=3 or any(len(p)!=2 for p in pairs):raise ValueError('Expected two or three contextual conditions')
  if len({p[0] for p in pairs})!=len(pairs) or list(pairs)!=sorted(pairs):raise ValueError('Noncanonical contextual rule')
  if any(k not in CONDITIONS or v not in CONDITIONS[k] for k,v in pairs):raise ValueError('Forbidden contextual rule condition')
  return dict(pairs)


def validate_weights(weights):
  durations={str(F(i,4)).replace('/','_') for i in range(1,17)}|{'1_3','2_3'}
  allowed={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():
    joint=rule_conditions(key)
    if (key not in allowed and joint is None) or not isinstance(value,int) or isinstance(value,bool) or abs(value)>1024:
      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)


def matches_signature(notes,expected):
  """Exact signature equality, stopping at the first differing feature."""
  if len(notes)!=len(expected):return False
  if not notes:return True
  first=notes[0];pitch=first.pitch.number()
  return all(n.pitch.number()-pitch==relative and n.end-n.start==duration and n.start-first.start==onset
    for n,(onset,duration,relative) in zip(notes,expected))


class Selector:
  @property
  def weights(self):return self._weights
  @weights.setter
  def weights(self,value):self._weights=MappingProxyType(validate_weights(value))

  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={}
    self.home_signatures={}
    self.cost_static={};self.cost_weights=None

  def contexts(self,notes,recent=()):
    history=list(recent[-2:])
    for n in notes:
      region=self.harmony.at(n.start)
      if region and region.chord:
        chord=expressions.parsed_chord(region.chord)
        root=int(model.NATURAL[chord['root']['letter']]+model.alteration(chord['root']['alter']))
        delta=n.pitch.number()-history[-1].pitch.number() if history else None
        previous=history[-1].pitch.number()-history[-2].pitch.number() if len(history)>1 else None
        metric='barline' if n.start+self.analysis.body_offset in self.analysis.boundaries else 'beat' if n.start%1==0 else 'eighth' if n.start%1==F(1,2) else 'subdivision'
        context={'quality':expressions.quality(chord),'degree':str((n.pitch.number()-root)%12),
          'metric':metric,'duration':str(n.end-n.start).replace('/','_'),
          'source_pair':'yes' if (previous,delta) in self.source_pairs else 'no'}
        if delta is not None:context['motion']=str(max(-12,min(12,delta)))
        yield context,int((n.end-n.start)*12)
      history.append(n)

  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)
    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
    for n in action.notes:
      f,p=self.note_base(n.start,n.end,n.pitch.number())
      values.update(f);prior.update(p)
    self.static[action.key]=(values,prior)
    if len(self.static)>8192:self.static.pop(next(iter(self.static)))
    return values,prior

  @functools.lru_cache(maxsize=32768)
  def note_base(self,start,end,number):
    values=Counter();prior=Counter()
    duration=end-start;length=int(duration*12)
    label=str(duration).replace('/','_')
    values['duration/'+label]+=length
    beat=(start%1)
    metric='barline' if 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,(number-self.center+6)//12))
    values[f'register/{band}']+=length
    region=self.harmony.at(start)
    if not region or not region.chord:return values,prior
    chord=expressions.parsed_chord(region.chord)
    root=int(model.NATURAL[chord['root']['letter']]+model.alteration(chord['root']['alter']))
    degree=(number-root)%12
    values[f'harmony/{expressions.quality(chord)}/{degree}/{label}']+=length
    prior['register_distance']+=length*max(0,abs(number-self.center)-7)
    if duration>=2 and number%12 in self.harmony.members(start,end):
      prior['sustained_arrival']+=length
    values['source/rhythm_match']+=length if duration in self.source_rhythms else 0
    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)
    # Silence is measured as intervals in the sounded stream, independent of
    # the number of rest actions used to encode it. Missing harmony is distinct.
    cursor=action.start;before=recent[-1] if recent else None
    for n in (*action.notes,None):
      until=n.start if n else action.end
      if until>cursor:
        f,p=self.space_features(cursor,until,None if before is None else before.end-before.start>=F(3,2))
        values.update(f);prior.update(p)
      if n:cursor=n.end;before=n
    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 len(home)>=3:
      key=id(home)
      if key not in self.home_signatures:
        self.home_signatures[key]=(home,signature(home))
        if len(self.home_signatures)>32:self.home_signatures.pop(next(iter(self.home_signatures)))
      expected=self.home_signatures[key][1]
      joined=(*recent,*action.notes)
      for i in range(max(0,len(recent)-len(home)+1),len(joined)-len(home)+1):
        part=joined[i:i+len(home)]
        if part[0].start>home[-1].end and matches_signature(part,expected):
          values['theme/realized_identity']+=12
          prior['theme_relation']+=12
    joint=[(key,rule_conditions(key)) for key in self.weights if key.startswith('joint/')]
    if joint:
      for context,ticks in self.contexts(action.notes,recent):
        for key,conditions in joint:
          if all(context.get(k)==v for k,v in conditions.items()):values[key]+=ticks
    return values,prior

  def complete(self,actions,home=(),*,start=None,end=None):
    """The sole authoritative score. Ties/rests/partitions are canonicalized."""
    from generator.solo.operators import Action
    actions=tuple(actions)
    first=start if start is not None else actions[0].start
    last=end if end is not None else actions[-1].end
    notes=types.sounding_notes(tuple(n for a in actions for n in a.notes))
    contracts=tuple(dict.fromkeys(c for a in actions for c in a.contracts))
    return self.evaluate(Action(first,last,notes,'sounded_phrase',contracts),home=home)

  def extend(self,previous,old_notes,action,home=(),*,old_home=(),start=F(),contracts=()):
    """Exact sufficient-statistic update, checked against complete() in tests.

    Reevaluate the seam and the tail containing any possible new home match.
    Previously completed events are not rescored on every search expansion.
    """
    from generator.solo.operators import Action
    if previous is None or home!=old_home or any(n.continuation for n in action.notes):
      notes=types.sounding_notes((*old_notes,*action.notes))
      return self.evaluate(Action(start,action.end,notes,'sounded_phrase',contracts),home=home)
    span=F(previous.denominators['span'])
    raw=Counter({k:v*span for k,v in previous.raw.items() if not k.startswith('obligation/')})
    index=max(0,len(old_notes)-max(3,len(home)+2))
    tail=old_notes[index:];context=old_notes[max(0,index-3):index]
    begin=tail[0].start if index else start
    before=Action(begin,action.start,tail,'sounded_phrase')
    after=Action(begin,action.end,(*tail,*action.notes),'sounded_phrase')
    for sign,part in ((-1,before),(1,after)):
      f,p=self.features(part,context,home)
      for k,v in f.items():
        if not k.startswith('obligation/'):raw[k]+=sign*v
      for k,v in p.items():raw['prior/'+k]+=sign*v
    obligations,_=self.base(Action(start,action.end,(),'sounded_phrase',contracts))
    raw.update({k:v for k,v in obligations.items() if k.startswith('obligation/')})
    weights={'prior/'+k:v for k,v in DEFAULT.items()}|self.weights
    duration=action.end-start
    raw={k:F(v)/duration for k,v in raw.items() if v}
    contributions={k:v*weights.get(k,0) for k,v in raw.items()}
    return scoring.Evaluation(VERSION,True,raw,weights,contributions,sum(contributions.values(),F()),
      {'span':str(duration),'ticks_per_quarter':12,'attacks':len(old_notes)+len(action.notes)})

  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() if v}|{k:v for k,v in features.items() if v}
    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)

  @functools.lru_cache(maxsize=32768)
  def transition_cost(self,delta,preceding,earlier,length,label,contiguous):
    token=max(-12,min(12,delta))
    sign=lambda x:'up' if x>0 else 'down' if x<0 else 'repeat'
    cost=length*self.weights.get(f'motion/{token}/{label}',0)
    cost+=length*DEFAULT['excessive_leap']*max(0,abs(delta)-7)
    if delta==0 and contiguous:cost+=length*DEFAULT['repeated_attack']
    if preceding is not None:
      cost+=length*self.weights.get(f'continuation/{sign(preceding)}/{sign(delta)}',0)
      if (preceding,delta) in self.source_pairs:cost+=length*DEFAULT['source_relation']
      if 0<abs(delta)<=2 and delta*preceding>0:cost+=length*DEFAULT['directed_motion']
      if earlier is not None and earlier==-preceding==delta and delta:cost+=length*DEFAULT['oscillation']
    return cost

  @functools.lru_cache(maxsize=32768)
  def seam_cost(self,prefix,context,contiguous):
    def total(initial):
      pitches=list(initial);value=0
      for index,(number,length,label,touch) in enumerate(prefix):
        if pitches:
          value+=self.transition_cost(number-pitches[-1],pitches[-1]-pitches[-2] if len(pitches)>1 else None,
            pitches[-2]-pitches[-3] if len(pitches)>2 else None,length,label,contiguous if index==0 and initial else touch)
        pitches.append(number)
      return value
    return total(context)-total(())

  @functools.lru_cache(maxsize=32768)
  def unary_cost(self,start,end,number):
    features,prior=self.note_base(start,end,number)
    return sum(v*DEFAULT[k] for k,v in prior.items())+sum(v*self.weights.get(k,0) for k,v in features.items())

  def intrinsic_cost(self,action):
    """Exact ordinary retrieval cost without constructing discarded reports.

    Unary, transition, silence and pending-target contributions are identical
    to features(action). Contextual learned conjunctions use the full path.
    """
    notes=types.sounding_notes(action.notes) if any(n.continuation for n in action.notes) else action.notes
    cost=0;history=[];cursor=action.start;settled=None;previous=None
    for n in notes:
      if n.start>cursor:cost+=self.gap_cost(cursor,n.start,settled)
      number=n.pitch.number();length=int((n.end-n.start)*12);label=str(n.end-n.start).replace('/','_')
      cost+=self.unary_cost(n.start,n.end,number)
      if history:
        cost+=self.transition_cost(number-history[-1],history[-1]-history[-2] if len(history)>1 else None,
          history[-2]-history[-3] if len(history)>2 else None,length,label,n.start==previous.end)
      history.append(number);previous=n;cursor=n.end;settled=n.end-n.start>=F(3,2)
    if action.end>cursor:cost+=self.gap_cost(cursor,action.end,settled)
    ticks=int((action.end-action.start)*12)
    for contract in action.contracts:
      if contract.target.start>=action.end and contract.events:
        last=contract.events[-1];target=contract.target
        cost+=ticks*self.weights.get(f'obligation/step/{max(-12,min(12,target.pitch.number()-last.pitch.number()))}',0)
        cost+=ticks*self.weights.get('obligation/delay/'+str(target.start-last.start).replace('/','_'),0)
        cost+=ticks*self.weights.get('obligation/hold/'+str(target.end-target.start).replace('/','_'),0)
    return cost

  def fast_cost(self,action,recent=(),context=None):
    """Unnormalized endpoint-bucket retrieval, with cached internal routes.

    Only the first three transitions depend on incoming history. Final ranking
    always uses complete(); this heuristic does not compare unequal spans.
    """
    # Learned conjunctions include incoming context; do not use a shortcut
    # that forgets it. This small fallback is explicit and profileable.
    if any(n.continuation for n in action.notes) or any(k.startswith('joint/') for k in self.weights):
      f,p=self.features(action,recent)
      return sum(v*DEFAULT[k] for k,v in p.items())+sum(v*self.weights.get(k,0) for k,v in f.items())
    if self.cost_weights is not self.weights:
      self.cost_static.clear();self.transition_cost.cache_clear();self.seam_cost.cache_clear();self.gap_cost.cache_clear();self.unary_cost.cache_clear();self.cost_weights=self.weights
    key=id(action)
    if key not in self.cost_static:
      value=self.intrinsic_cost(action)
      prefix=tuple((n.pitch.number(),int((n.end-n.start)*12),str(n.end-n.start).replace('/','_'),
        i>0 and action.notes[i-1].end==n.start) for i,n in enumerate(action.notes[:3]))
      self.cost_static[key]=(action,value,prefix)
      if len(self.cost_static)>32768:self.cost_static.pop(next(iter(self.cost_static)))
    _,value,prefix=self.cost_static[key]
    if context is None:context=(tuple(n.pitch.number() for n in recent[-3:]),recent[-1].end if recent else None)
    if context[0] and prefix:value+=self.seam_cost(prefix,context[0],context[1]==action.notes[0].start)
    gap_end=action.notes[0].start if action.notes else action.end
    if recent and gap_end>action.start:
      # Incoming sustained arrivals change the role of a leading rest, just
      # as incoming pitches change the first intervals. Do not cache it as a
      # fresh silent entrance in every action.
      value+=self.gap_cost(action.start,gap_end,recent[-1].end-recent[-1].start>=F(3,2))-self.gap_cost(action.start,gap_end,None)
    return value

  @functools.lru_cache(maxsize=8192)
  def space_features(self,start,end,settled):
    values=Counter();prior=Counter();delayed=0
    boundaries=sorted({start,end,*[s.start for s in self.harmony.spans if start<s.start<end],
      *[s.end for s in self.harmony.spans if start<s.end<end]})
    for left,right in zip(boundaries,boundaries[1:]):
      ticks=int((right-left)*12);region=self.harmony.at(left)
      label='unharmonized' if not region or not region.scale else 'delayed_entrance' if settled is None else 'after_arrival' if settled else 'breath'
      values['space/'+label]+=ticks
      if label=='breath':prior['unprepared_space']+=ticks
      if label=='delayed_entrance':delayed+=ticks
    prior['unprepared_space']+=max(0,delayed-12)
    return values,prior

  @functools.lru_cache(maxsize=8192)
  def gap_cost(self,start,end,settled):
    f,p=self.space_features(start,end,settled)
    return sum(v*DEFAULT[k] for k,v in p.items())+sum(v*self.weights.get(k,0) for k,v in f.items())