"""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
import fractions
import pathlib
import generator.model
import generator.solo.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 = generator.solo.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 = generator.solo.expressions.parsed_chord(future.chord) if future else None

  def root(c): return int(generator.model.NATURAL[c['root']['letter']] + generator.model.alteration(c['root']['alter'])) % 12 if c else 0

  return (generator.solo.expressions.quality(current) if current else 'unknown',
    generator.solo.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(fractions.Fraction, 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(pathlib.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=fractions.Fraction), '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())}}