"""Smoothed joint musical-event model. No examples or song lookup at inference.

Counts are sufficient statistics of short, relative contexts, not stored scores.
The hierarchical Dirichlet estimator uses eight parent pseudo-observations.
All paths through the decoder score the same sounding event stream identically.
"""
from __future__ import annotations
import collections
import json
import math
import functools
from fractions import Fraction as F
from pathlib import Path
from generator import model
from generator.solo import expressions

VERSION = 'pooled-events-1'
PRIOR = 8
DURATIONS = ('1/4','1/3','1/2','2/3','3/4','1','3/2','2','5/2','3','4')


def key(value):
  return json.dumps(value,separators=(',',':'))


def chord_context(harmony,time):
  span=harmony.at(time)
  current=expressions.parsed_chord(span.chord) if span and span.chord!='N.C.' else None
  future=next((s for s in harmony.spans if s.start>time and s.chord!='N.C.' and (not span or s.chord!=span.chord)),None)
  following=expressions.parsed_chord(future.chord) if future else None
  def root(c):return int(model.NATURAL[c['root']['letter']]+model.alteration(c['root']['alter']))%12 if c else 0
  return (expressions.quality(current) if current else 'unknown',
    expressions.quality(following) if following else 'end',
    (root(following)-root(current))%12 if current and following else -1),root(current),future


def token(duration,pitch,previous,root):
  return ('rest',str(duration),0,0) if pitch is None else ('note',str(duration),
    pitch-previous if previous is not None else 'start',(pitch-root)%12)


def source_feature(analysis,time,root=0):
  time+=analysis.body_offset
  before=[n for n in analysis.melody if n.start<=time][-2:]
  following=next((n for n in analysis.melody if n.start>time),None)
  direction=0 if len(before)<2 else (before[-1].pitch.number()>before[-2].pitch.number())-(before[-1].pitch.number()<before[-2].pitch.number())
  held=bool(before and before[-1].end>time)
  upcoming=0 if not before or not following else (following.pitch.number()>before[-1].pitch.number())-(following.pitch.number()<before[-1].pitch.number())
  return (direction,held,upcoming,str(before[-1].end-before[-1].start) if before else 'none',
    (before[-1].pitch.number()-root)%12 if before else -1,
    before[-1].pitch.number()-before[-2].pitch.number() if len(before)>1 else 0)


def phase_at(analysis,time):
  absolute=time+analysis.body_offset
  start=next((a for a,b in zip(analysis.boundaries,analysis.boundaries[1:]) if a<=absolute<b),analysis.boundaries[-2])
  return absolute-start


@functools.lru_cache(maxsize=65536)
def contexts(harmonic,phase,source,history):
  # No title, path, absolute bar number, progression ID, or target-derived ID.
  base=(*harmonic,str(phase))
  result=[(),base,(*base,source)]
  # A complete four-event rhythmic/melodic relationship can influence the next
  # event, while unseen histories smoothly return to broader musical evidence.
  for count in range(1,min(4,len(history))+1):
    result.append((*base,source,tuple(history[-count:])))
  return tuple(key(c) for c in result)


class Model:
  def __init__(self,value):
    if value.get('version')!=VERSION:raise ValueError('Missing/incompatible solo model; run python -m generator.train')
    self.value=value
    self.tables={k:(v['total'],v['counts']) for k,v in value['contexts'].items()}
    self.durations=tuple(sorted(map(F,value['durations'])))
    self.rhythms=value['rhythms'];self.steps=value['steps'];self.degrees=value['degrees']
    self._prob={}
    self._totals={id(t):sum(t.values()) for t in (self.rhythms,self.steps,self.degrees)}

  @classmethod
  def load(cls,path):
    try:return cls(json.loads(Path(path).read_text()))
    except FileNotFoundError as e:raise ValueError('Solo model not found; run python -m generator.train') from e

  def probability(self,action,context):
    cache=(action,context)
    if cache in self._prob:return self._prob[cache]
    kind,duration,step,degree=action
    def p(table,item,size):return (table.get(str(item),0)+1)/(self._totals[id(table)]+size)
    probability=p(self.rhythms,kind+':'+duration,2*len(self.durations))
    if kind=='note':probability*=p(self.steps,step,50)*p(self.degrees,degree,12)
    encoded=key(action)
    for c in context:
      if c in self.tables:
        total,counts=self.tables[c]
        probability=(counts.get(encoded,0)+PRIOR*probability)/(total+PRIOR)
    if len(self._prob)>100000:self._prob.clear()
    self._prob[cache]=probability
    return probability

  def cost(self,action,context):return -math.log(self.probability(action,context))

  def register_cost(self,number):
    counts=self.value['registers']
    return -math.log((counts.get(str(number//12),0)+PRIOR)/(sum(counts.values())+PRIOR*4))


def fit(examples,provenance):
  tables=collections.defaultdict(collections.Counter)
  rhythms=collections.Counter();steps=collections.Counter();degrees=collections.Counter()
  registers=collections.Counter();durations=set(DURATIONS)
  for example in examples:
    history=[];previous=None
    for start,end,pitch in example['events']:
      harmonic,root,_=chord_context(example['harmony'],start)
      action=token(end-start,pitch,previous,root)
      context=contexts(harmonic,phase_at(example['analysis'],start-example['analysis'].body_offset),
        source_feature(example['analysis'],start-example['analysis'].body_offset,root),tuple(history[-4:]))
      for c in context:tables[c][key(action)]+=1
      durations.add(str(end-start));rhythms[action[0]+':'+action[1]]+=1
      if pitch is not None:
        steps[str(action[2])]+=1;degrees[str(action[3])]+=1
        registers[str(pitch//12)]+=1;previous=pitch
      history.append(action)
  return {'version':VERSION,'provenance':provenance,'prior':PRIOR,
    'durations':sorted(durations,key=F),'rhythms':dict(rhythms),'steps':dict(steps),'degrees':dict(degrees),
    'registers':dict(registers),
    'contexts':{k:{'total':sum(v.values()),'counts':dict(sorted(v.items()))} for k,v in sorted(tables.items())}}