"""Constrained, deterministically seeded autoregressive musical sampling.

The pooled model proposes joint note/rest, duration and pitch events. Complete
expression actions include their real resolution. Guitar and chord/function
validity precede sampling. No target access, hand-authored melody scorer or
legacy fallback. A bounded restart handles an exhausted/entirely silent sample.
"""
from __future__ import annotations
import collections
import dataclasses
import hashlib
import random
import math
from fractions import Fraction as F
from generator import model as music, vocabulary
from generator.solo import expressions, groups, learned, types


class SearchExhausted(ValueError):pass


@dataclasses.dataclass(frozen=True)
class State:
  events:tuple=()
  notes:tuple=()
  contracts:tuple=()
  history:tuple=()
  previous:int|None=None
  cost:float=0.


def event_key(events):return tuple((a,b,-1 if p is None else p) for a,b,p in events)


class Decoder:
  def __init__(self,analysis,harmony,model,role,start,end,*,expressions_enabled=True,attempts=4,seed=''):
    if type(attempts) is not int or not 1<=attempts<=16:raise ValueError('attempts must be between 1 and 16')
    self.analysis,self.harmony,self.model,self.role=analysis,harmony,model,role
    self.start,self.end=start,end
    self.attempts=attempts
    source=tuple((str(n.start),str(n.end),n.pitch.number()) for n in analysis.melody)
    self.seed=seed or hashlib.sha256(repr((source,role,learned.VERSION)).encode()).hexdigest()
    self.expressions_enabled=expressions_enabled
    self.low,self.high=harmony.profile.register
    limits=vocabulary.data('roles')['difficulty']
    self.durations=tuple(d for d in model.durations if 0<d<=4 and d.denominator in (1,2,4) and
      (d.denominator!=4 or analysis.tempo<=limits['sixteenth_max_bpm']))
    self.allow_triplets=analysis.tempo<=limits['triplet_max_bpm']
    self._actions={};self._context={};self._pitches={};self.trace=[];self.counts=collections.Counter()

  def context(self,time):
    if time not in self._context:
      harmonic,root,future=learned.chord_context(self.harmony,time)
      self._context[time]=(harmonic,root,learned.phase_at(self.analysis,time),learned.source_feature(self.analysis,time,root),future)
    return self._context[time]

  def pitch(self,number,time):
    span=self.harmony.at(time)
    identity=(number,span.start if span else None)
    if identity not in self._pitches:self._pitches[identity]=self._spell(number,time)
    return self._pitches[identity]

  def _spell(self,number,time):
    span=self.harmony.at(time)
    if span and span.chord!='N.C.':
      chord=expressions.parsed_chord(span.chord)
      from generator import harmony as chords
      for degree,alter in chords.degrees(chord).items():
        p=expressions.degree_pitch(chord,degree,alter,number)
        if p.number()==number:return p
      for letter,alter in zip(music.LETTERS,span.key['alterations']):
        accidental=music.fraction(alter)
        if (music.NATURAL[letter]+accidental)%12==number%12:
          return types.Pitch(letter,accidental,int((number-music.NATURAL[letter]-accidental)//12)-1)
    names=(('C',0),('D',-1),('D',0),('E',-1),('E',0),('F',0),('F',1),('G',0),('A',-1),('A',0),('B',-1),('B',0))
    letter,alter=names[number%12]
    return types.Pitch(letter.lower(),F(alter),number//12-1)

  def note(self,a,b,p,role='learned'):
    pitch=p if isinstance(p,types.Pitch) else self.pitch(p,a)
    return types.Note(a,b,pitch,role=role,event_id=f'{a}:{b}:{pitch.number()}')

  def permission(self,n):
    """Find justified color/anticipation support, never authorize by label."""
    # Source spans are contiguous. These are precisely the ordinary-event
    # branches of Harmony.evaluate, without rebuilding an empty contract map.
    if self.harmony.at(n.start) and n.end<=self.harmony.spans[-1].end:
      if n.pitch.number()%12 in self.harmony.members(n.start,n.end):return ()
      if n.end-n.start<2 and n.pitch.number()%12 in self.harmony.permitted(n.start,n.end):return ()
    if not self.expressions_enabled:return None
    span=self.harmony.at(n.start)
    if not span or span.chord=='N.C.':return None
    exposed=tuple(region for region,_,_ in self.harmony.segments(n)
      if n.pitch.number()%12 not in region.roles.basic|region.roles.specified)
    for rule in vocabulary.data('harmony')['colors']:
      if rule.get('resolve') or n.end-n.start>F(rule['max_beats']):continue
      # A held chord member can become a color only after a chord change.
      # Test the regions needing permission, not just the attack's harmony.
      if any(region.chord=='N.C.' or not expressions.eligible_color(rule,region,self.harmony) or
        expressions.degree_pitch(expressions.parsed_chord(region.chord),rule['degree'],rule['alter'],n.pitch.number()).number()!=n.pitch.number()
        for region in exposed):continue
      c=types.ExpressionContract(n.event_id+':color','color',rule['id'],(n,),n,span.chord)
      if self.harmony.evaluate((n,),complete=True,expressions=(c,)).valid:return (c,)
    for arriving in self.harmony.spans:
      if n.start<arriving.start<n.end and arriving.start-n.start<=F(vocabulary.data('harmony')['anticipation_max']):
        c=types.ExpressionContract(n.event_id+':anticipation','anticipation','early_chord_tone',(n,),n,arriving.chord,arriving.start)
        if self.harmony.evaluate((n,),complete=True,expressions=(c,)).valid:return (c,)
    return None

  def actions(self,time):
    if time in self._actions:return self._actions[time]
    result=[]
    def add(events,notes,contracts=()):result.append((events,notes,contracts))
    for duration in self.durations:
      end=time+duration
      if end>self.end:continue
      add(((time,end,None),),())
      for number in range(self.low,self.high+1):
        n=self.note(time,end,number);support=self.permission(n)
        if support is not None:add(((time,end,number),),(n,),support)
    if self.expressions_enabled:
      limits=vocabulary.data('roles')['difficulty']
      for rule in vocabulary.data('lead_ins')['patterns']:
        difficulty=rule.get('difficulty')
        if difficulty and self.analysis.tempo>limits[difficulty+'_max_bpm']:continue
        # Connections require a concrete preceding note, checked at expansion.
        durations=tuple(map(F,rule['durations']));gaps=tuple(map(F,rule.get('internal_gaps',['0']*(len(durations)-1))))
        target_time=time+sum(durations,F())+sum(gaps,F())+F(rule['gap'])
        span=self.harmony.at(target_time)
        if not span or span.chord=='N.C.':continue
        holds=(F(1,3),) if difficulty=='triplet' else (F(1,2),F(1),F(2))
        for duration in holds:
          end=target_time+duration
          if end>self.end:continue
          for number in range(self.low,self.high+1):
            if number%12 not in self.harmony.members(target_time,end):continue
            target=self.note(target_time,end,number,'destination');cursor=time;notes=[];events=[]
            for i,(delta,d) in enumerate(zip(rule['semitones'],durations)):
              n=self.note(cursor,cursor+d,expressions.chromatic(target.pitch,delta),rule['kind'])
              notes.append(n);events.append((n.start,n.end,n.pitch.number()));cursor=n.end
              gap=gaps[i] if i<len(gaps) else F(rule['gap'])
              if gap:events.append((cursor,cursor+gap,None));cursor+=gap
            if any(not self.low<=n.pitch.number()<=self.high for n in notes):continue
            if not groups.valid(tuple((a,b) for a,b,p in (*events,(target_time,end,number)))):continue
            contract=types.ExpressionContract(f'{time}:{number}:{rule["id"]}',rule['kind'],rule['id'],tuple(notes),target,span.chord)
            ns=(*notes,target)
            if rule['kind']=='connection':
              # Validation waits for the actual incoming event.
              add((*events,(target_time,end,number)),ns,(contract,))
            elif self.harmony.evaluate(ns,complete=True,expressions=(contract,)).valid:
              add((*events,(target_time,end,number)),ns,(contract,))
    # Complete ordinary triplet groups; do not leave a third-beat residual.
    if self.allow_triplets and time+1<=self.end:
      for number in range(self.low,self.high+1):
        for direction in (-1,0,1):
          numbers=[number];cursor=time
          for _ in range(2):
            nxt=next((p for p in range(numbers[-1]+direction,numbers[-1]+direction*4,direction or 1)
              if self.low<=p<=self.high and p%12 in self.harmony.permitted(cursor+F(1,3),cursor+F(2,3))),None) if direction else number
            if nxt is None:break
            numbers.append(nxt);cursor+=F(1,3)
          if len(numbers)!=3:continue
          ns=tuple(self.note(time+i*F(1,3),time+(i+1)*F(1,3),p) for i,p in enumerate(numbers))
          if self.harmony.evaluate(ns,complete=True).valid:add(tuple((n.start,n.end,n.pitch.number()) for n in ns),ns)
    self._actions[time]=tuple(result);self.counts['valid_structures']+=len(result)
    return self._actions[time]

  def extend(self,state,action):
    events,notes,contracts=action
    # Training tokenization merges contiguous silence. Never manufacture a
    # rest-to-rest transition absent from that event representation.
    if state.events and state.events[-1][2] is None and events[0][2] is None:return None
    if notes and state.previous is not None and abs(notes[0].pitch.number()-state.previous)>12:return None
    if contracts and contracts[0].kind=='connection':
      if not state.notes:return None
      c=dataclasses.replace(contracts[0],preparation=state.notes[-1]);contracts=(c,)
      _,error=expressions.validate((state.notes[-1],*notes),contracts,self.harmony)
      if error:return None
    history=state.history;previous=state.previous;cost=state.cost
    for a,b,p in events:
      harmonic,root,phase,source,_=self.context(a)
      action=learned.token(b-a,p,previous,root)
      context=learned.contexts(harmonic,phase,source,history)
      cost+=self.model.cost(action,context)
      if p is not None:
        if previous is None:cost+=self.model.register_cost(p)
        previous=p
      history=(*history,action)[-4:]
    return State((*state.events,*events),(*state.notes,*notes),(*state.contracts,*contracts),history,previous,cost)

  def compose(self):
    for attempt in range(self.attempts):
      rng=random.Random(self.seed+':'+str(attempt));state=State();time=self.start
      trace=[]
      while time<self.end:
        actions=self.actions(time)
        options=[s for a in actions if (s:=self.extend(state,a)) is not None]
        # Different operator names cannot multiply the probability of identical
        # sounded music. Prefer plain permission when no exception is needed.
        unique={}
        for s in options:
          signature=event_key(s.events[len(state.events):])
          if signature not in unique or len(s.contracts)<len(unique[signature].contracts):unique[signature]=s
        options=sorted(unique.values(),key=lambda s:event_key(s.events[len(state.events):]))
        self.counts['scored_continuations']+=len(options)
        if not options:break
        low=min(s.cost-state.cost for s in options)
        weights=[math.exp(-(s.cost-state.cost-low)) for s in options]
        draw=rng.random()*sum(weights);selected=options[-1]
        for s,w in zip(options,weights):
          draw-=w
          if draw<=0:selected=s;break
        trace.append({'time':str(time),'valid_joint_actions':len(options),
          'selected_events':[(str(a),str(b),p) for a,b,p in selected.events[len(state.events):]],
          'negative_log_likelihood':selected.cost-state.cost})
        state=selected;time=state.events[-1][1]
      if time==self.end and state.notes:
        self.trace=trace;self.completed=(state,);self.counts['attempts']=attempt+1
        return state
    raise SearchExhausted(f'{self.role}: no complete non-silent sample within {self.attempts} deterministic attempts; end={self.end}')