"""Bounded search of joint musical actions; one normal production path.

Selection compares complete actions by cost per quarter. Actions are indivisible;
cross-bar paths and resolution contracts never restart at search boundaries.
"""
from __future__ import annotations
from collections import Counter
import dataclasses
from fractions import Fraction as F
import time

from generator import vocabulary
from generator.solo import types, groups, operators, selection, sequences, scoring

class SearchExhausted(ValueError):pass

@dataclasses.dataclass(frozen=True)
class State:
  total: F = F()
  recent: tuple = ()
  home: tuple = ()
  last_kind: str = ''
  parent: object = None
  action: object = None
  evaluation: object = None
  tie: tuple = ()
  pending: tuple = ()
  alternatives: tuple = ()

class Builder:
  def __init__(self,analysis,measures,harmony,sequences,rhythms,end,seed,enabled=True,*,start=F()):
    self.analysis=analysis;self.harmony=harmony;self.start=start;self.end=end
    self.operators=operators.Operators(harmony,analysis,end,start,enabled,sequences,rhythms)
    self.selector=selection.Selector(harmony,analysis,harmony.profile.reference or 60)
    self.rules=vocabulary.data('solo')['operators'];self.decisions=[];self.failures=Counter();self.coverage={};self.timings=Counter()

  def candidates(self,at,state):
    candidates=self.operators.at(at)
    if state.home and state.last_kind!='return':
      candidates+=self.operators.related(at,state.home,'return')+self.operators.related(at,state.home,'answer')
    evaluated=[]
    for a in candidates:
      if state.pending and self.advance(state,None,a) is None:continue
      cost=self.selector.cost(a,state.recent,state.home,state.last_kind)
      evaluated.append((cost,a))
    # Rank complete actions by cost per quarter before bounded repair.
    # Accumulated prefix totals are diagnostic, not a global search objective.
    retained=sequences.diverse(evaluated,
      lambda v:(v[0],v[1].key),
      lambda v:(v[1].kind,v[1].end-v[1].start),self.rules['per_state'])
    for _,a in retained:self.operators.coverage['retained:'+a.kind]+=1
    return tuple((self.selector.evaluate(a,state.recent,state.home,state.last_kind),a) for _,a in retained)

  def advance(self,state,evaluation,action):
    for contract in state.pending:
      target=contract.target
      if action.start<=target.start<action.end:
        # Reject a missed committed attack/hold before revalidating all the
        # harmonic relationships. Partial target evidence remains possible.
        if not any(n.event_id==target.event_id and n.start==target.start
          and n.end==min(action.end,target.end) and n.pitch==target.pitch for n in action.notes):return None
    obligations=(*state.pending,*action.contracts);pending=()
    if not state.pending:
      pending=tuple(c for c in action.contracts if c.target.end>action.end)
    elif obligations:
      heard={n.event_id:n for c in state.pending for n in c.events if n.start<action.start}
      heard.update({n.event_id:n for n in action.notes})
      check=self.harmony.evaluate(tuple(sorted(heard.values(),key=lambda n:n.start)),
        expressions=obligations,until=action.end,complete=action.end==self.end)
      if not check.valid:return None
      pending=check.pending
    recent=(*state.recent,*action.notes)[-8:];home=state.home
    if not home and len(recent)>=4:
      proposed=recent[:4]
      if proposed[-1].end-proposed[0].start<=8:home=proposed
    return State(state.total+(evaluation.total if evaluation else 0),recent,home,action.kind,state,action,evaluation,(),pending)

  def compose(self):
    self.decisions=[];began=time.monotonic();backtracks=0
    # Compare full actions by cost per exact quarter, with a bounded repair
    # frontier. This is a local policy, not a claim of globally optimal melody.
    stack=[(self.start,State(),None)];rankings={}
    while stack:
      at,state,iterator=stack[-1]
      if at==self.end:final=state;break
      if iterator is None:
        options=sorted(self.candidates(at,state),key=lambda v:(v[0].total,v[1].key))
        rankings[id(state)]=tuple({'kind':a.kind,'start':str(a.start),'end':str(a.end),
          'notes':[(str(n.start),str(n.end),n.pitch.number()) for n in a.notes],
          'evaluation':dataclasses.asdict(e)} for e,a in options[:4])
        iterator=iter(options);stack[-1]=(at,state,iterator)
      for evaluation,action in iterator:
        next_state=self.advance(state,evaluation,action)
        if next_state is not None:
          next_state=dataclasses.replace(next_state,alternatives=rankings[id(state)])
          stack.append((action.end,next_state,None));break
      else:
        stack.pop();backtracks+=1
        if backtracks>self.rules['backtracks']:raise SearchExhausted(f'at {at}: action repair budget exhausted')
    else:raise SearchExhausted('no complete action path')
    self.timings['backtracks']=backtracks
    self.timings['action_search_seconds']=time.monotonic()-began
    path=[];state=final
    while state.action:
      path.append(state);state=state.parent
    path.reverse();phrases=[];pending=[]
    self.selected_actions=tuple(s.action for s in path)
    for state in path:
      action=state.action;pending.append(state)
      self.operators.coverage['selected:'+action.kind]+=1
      self.decisions.append({'start':str(action.start),'end':str(action.end),'selected':action.kind,
        'evaluation':dataclasses.asdict(state.evaluation),'prefix_total':state.total,'alternatives':state.alternatives})
      # Reporting groups the selected thought; it does not insert events.
      held=bool(action.notes and action.notes[-1].end-action.notes[-1].start>=2)
      breath=not action.notes and action.end-action.start>=1
      if held or breath or action.end==self.end:
        phrases.append(self.phrase(pending,len(phrases)));pending=[]
    for name,count in self.operators.coverage.items():
      stage,kind=name.split(':',1);self.coverage.setdefault(kind,{})[stage]=count
    self.timings['selected_cost']=final.total
    return tuple(phrases)

  def phrase(self,states,index):
    actions=tuple(s.action for s in states);ns=tuple(n for a in actions for n in a.notes)
    start,end=actions[0].start,actions[-1].end
    last=ns[-1] if ns else None
    held=last and last.end-last.start>=2
    ending='handoff' if end==self.end else 'settle' if held else 'open'
    anchor=None
    if held:
      region=self.harmony.at(last.start)
      anchor=types.Anchor(f'gesture_{index}:arrival',last.start,last.start,last.pitch,
        frozenset({last.pitch.number()%12}),self.harmony.profile.register,last.end-last.start,
        'terminal' if ending=='handoff' else 'intermediate',region.chord,False,last.end,ending)
    raw=Counter();contributions=Counter();weights={}
    for s in states:
      raw.update(s.evaluation.raw);contributions.update(s.evaluation.contributions);weights.update(s.evaluation.weights)
    evaluation=scoring.Evaluation(selection.VERSION,True,dict(raw),weights,dict(contributions),sum(s.evaluation.total for s in states),
      {'span':str(end-start),'attacks':len(ns),'ticks_per_quarter':12})
    direction='up' if len(ns)>1 and ns[-1].pitch.number()>ns[0].pitch.number() else 'down'
    rhythm=tuple(g for a in actions for g in groups.describe(tuple((n.start-a.start,n.end-a.start) for n in a.notes),origin=a.start,source=a.kind))
    return types.Phrase(f'gesture_{index}',start,end,ns,actions[0].kind,direction,ending,anchor,
      tuple(c for a in actions for c in a.contracts),rhythm,'home',evaluation.total,evaluation,tuple(a.source or a.kind for a in actions))