"""Development-only shared integer ranking from COMPLETE ordinary rollouts.

The returned coefficients depend on the training examples. Target event arrays,
source identifiers and diagnostic oracle paths are never production artifacts.
"""
from collections import Counter
from fractions import Fraction as F
import itertools
import math
from generator import vocabulary
from generator.solo import selection


def contextual_candidates(selector,notes):
  """Induce only whitelisted two/three-condition event contexts.

  These counts are training-dependent. No title, time index, song key, event
  array, fingerprint or arbitrary predicate can appear in the rule schema.
  """
  values=Counter()
  for context,ticks in selector.contexts(notes):
    terms=sorted((k,v) for k,v in context.items() if v in selection.CONDITIONS[k])
    for size in (2,3):
      for terms_ in itertools.combinations(terms,size):
        values['joint/'+'&'.join(k+'='+v for k,v in terms_)]+=ticks
  return values


def vector(builder,actions,home=()):
  e=builder.selector.complete(actions,home,start=builder.start,end=builder.end)
  # Equal spans: clear normalization without rounding any musical time.
  span=builder.end-builder.start
  features={k:F(v)*span for k,v in e.raw.items() if not k.startswith('prior/') and v}
  prior=sum(F(v)*span for k,v in e.contributions.items() if k.startswith('prior/'))
  return features,prior


def constraint(role,gold,other,case):
  left,lp=gold;right,rp=other
  delta={k:left.get(k,0)-right.get(k,0) for k in left.keys()|right.keys()}
  delta={k:v for k,v in delta.items() if v}
  selection.validate_weights(dict.fromkeys(delta,0))
  return {'role':role,'delta':delta,'prior':lp-rp,'case':case}


def solve(rows,bound=1024):
  """One shared coefficient table per role; deterministic pinned integer solver."""
  import z3
  if z3.get_version_string()!='4.13.3':raise ValueError('Use pinned z3-solver==4.13.3.0')
  z3.set_param('smt.random_seed',0)
  # Exact zero/one-coefficient cases need no combinatorial cardinality search.
  # Checking every feature's integer interval proves the minimum support and
  # L1 objective when all independent roles have such a solution. This does
  # not claim the resulting sparse model reproduces any complete melody.
  simple={}
  for role in sorted({r['role'] for r in rows}):
    subset=[r for r in rows if r['role']==role]
    if all(r['prior']<=-1 for r in subset):simple[role]={};continue
    choices=[]
    for key in sorted({k for r in subset for k in r['delta']}):
      low,high=-bound,bound
      for r in subset:
        slope=F(r['delta'].get(key,0));limit=-1-F(r['prior'])
        if slope>0:high=min(high,math.floor(limit/slope))
        elif slope<0:low=max(low,math.ceil(limit/slope))
        elif limit<0:low,high=1,0;break
      if low<=high:
        value=low if low>0 else high if high<0 else 0
        choices.append((abs(value),key,value))
    if not choices:break
    _,key,value=min(choices);simple[role]={key:value}
  if set(simple)=={r['role'] for r in rows}:
    for weights in simple.values():selection.validate_weights(weights)
    return {'status':'solved','weights':simple,'bound':bound,
      'nonzero':sum(len(w) for w in simple.values()),'l1':sum(abs(v) for w in simple.values() for v in w.values()),
      'method':'exhaustive zero/one-coefficient integer intervals; minimum support and L1 within that support proven'}
  names=sorted({(r['role'],k) for r in rows for k in r['delta']})
  variables={n:z3.Int('w_'+str(i)) for i,n in enumerate(names)}
  limits=[v>=-bound for v in variables.values()]+[v<=bound for v in variables.values()]
  rational=lambda v:z3.RealVal(str(F(v)))
  constraints=[z3.Sum([variables[(r['role'],k)]*rational(v) for k,v in r['delta'].items()])+rational(r['prior'])<=-1 for r in rows]
  check=z3.Solver();check.set(rlimit=2000000);check.add(*limits)
  for i,c in enumerate(constraints):check.assert_and_track(c,'comparison_'+str(i))
  status=check.check()
  if status==z3.unknown:return {'status':'solver_unknown','stage':'feasibility','reason':check.reason_unknown()}
  if status==z3.unsat:
    return {'status':'incompatible_with_current_features_and_bounds','bound':bound,
      'conflicts':[str(x) for x in check.unsat_core()],
      'identical_feature_pairs':[i for i,r in enumerate(rows) if not r['delta'] and r['prior']>=0]}
  opt=z3.Optimize();opt.add(*limits,*constraints)
  # Deterministic solver work budget. Exhaustion yields no proposed model;
  # elapsed wall time never chooses a musical winner.
  opt.set(rlimit=2000000)
  opt.minimize(z3.Sum([z3.If(v==0,0,1) for v in variables.values()]))
  opt.minimize(z3.Sum([z3.If(v<0,-v,v) for v in variables.values()]))
  if opt.check()!=z3.sat:return {'status':'solver_unknown','reason':opt.reason_unknown()}
  model=opt.model();weights={r:{} for r,_ in names}
  for (role,key),v in variables.items():
    value=model.eval(v,model_completion=True).as_long()
    if value:weights[role][key]=value
  for w in weights.values():selection.validate_weights(w)
  return {'status':'solved','weights':weights,'bound':bound,
    'nonzero':sum(len(w) for w in weights.values()),'l1':sum(abs(v) for w in weights.values() for v in w.values())}


def sounded(actions):
  from generator.solo.types import sounding_notes
  return tuple((n.start,n.end,n.pitch.number()) for n in sounding_notes(tuple(n for a in actions for n in a.notes)))


class CoverageTrace:
  """Read-only observer. Targets diagnose cuts, never supply search decisions."""
  def __init__(self,gold):
    self.gold=sounded(gold);self.records=[]
  def matches(self,notes,end):
    # These are committed whole actions, not a truncated playback preview.
    # A quarter-note attack cannot reproduce a half-note target merely because
    # its first quarter sounds the same. Future actions do not reopen it.
    from generator.solo.types import sounding_notes
    actual=tuple((n.start,n.end,n.pitch.number()) for n in sounding_notes(notes))
    expected=tuple((a,b,p) for a,b,p in self.gold if a<end)
    return actual==expected
  def __call__(self,stage,at,state,proposed,retained):
    if stage=='successors':
      if not self.matches(state.notes,at):return
      eligible=[a for a in proposed if self.matches((*state.notes,*a.notes),a.end)]
      survivors=[a for a in retained if self.matches((*state.notes,*a.notes),a.end)]
    else:
      eligible=[s for s in proposed if self.matches(s.notes,s.action.end if s.action else at)]
      survivors=[s for s in retained if self.matches(s.notes,s.action.end if s.action else at)]
    if eligible and not survivors:
      self.records.append({'time':str(at),'status':'pruned_at_'+stage,'matching_proposals':len(eligible)})
    elif stage=='successors' and not eligible:
      self.records.append({'time':str(at),'status':'not_proposed_in_actual_incoming_state'})


def fit_complete(examples,passes,on_checkpoint=None,resources=None):
  weights={role:dict(b.selector.weights) for _,role,b,_ in examples}
  rows=[];history=[];last_solver={};exact=False;comparisons=[];induced=[]
  for iteration in range(passes):
    cases=[];errors=0
    for name,role,b,gold in examples:
      b.selector.weights=weights[role]
      trace=CoverageTrace(gold);b.observe=trace
      # Honor the explicitly selected development data directory during lazy
      # operator expansion too, not just when constructing the builder.
      with vocabulary.using(resources or vocabulary.current()['data'],role):
        b.compose()  # No target argument, prefix, candidate injection or oracle.
      b.observe=None
      state=__import__('generator.solo.construction',fromlist=['State']).State()
      for a in gold:state=b.advance(state,None,a)
      positive=vector(b,gold,state.home)
      correct=sounded(b.selected_actions)==sounded(gold)
      cases.append({'case':name,'exact_complete_output':correct,
        'completed_candidates':len(b.completed),'winner':str(b.winner_evaluation.total),
        'earliest_observed_failures':trace.records[:8],
        'target_in_completed_pool':any(sounded(s.actions)==sounded(gold) for s in b.completed)})
      if not correct:errors+=1
      for s in b.completed:
        if sounded(s.actions)!=sounded(gold):
          comparisons.append((name,role,b,gold,state.home,s.actions,s.home))
          rows.append(constraint(role,positive,vector(b,s.actions,s.home),name))
      print(name,'exact' if correct else 'complete-output mismatch',flush=True)
    history.append({'iteration':iteration+1,'incorrect_complete_outputs':errors,'cases':cases})
    if not errors:exact=True
    result={'weights':weights,'history':history,'complete':exact,'solver':last_solver,
      'induced_rules':induced,
      'training_dependency':'Coefficients learned from supplied authored examples; runtime target-free is not training independence.',
      'release_gate':'Ordinary cold full-song outputs must independently pass exact comparison; derivation is not success.'}
    if on_checkpoint:on_checkpoint(result)
    if exact:break
    last_solver=solve(rows)
    if last_solver['status']!='solved':
      # Add a context only for a demonstrated feature collision. General
      # ranking failure is not permission to expand vocabulary indiscriminately.
      collisions=last_solver.get('identical_feature_pairs',[])
      if collisions and len(induced)<8 and iteration+1<passes:
        i=collisions[0];name,role,b,gold,gh,other,oh=comparisons[i]
        g=contextual_candidates(b.selector,tuple(n for a in gold for n in a.notes))
        o=contextual_candidates(b.selector,tuple(n for a in other for n in a.notes))
        keys=sorted((k for k in g.keys()|o.keys() if g.get(k,0)!=o.get(k,0) and k not in weights[role]),key=lambda k:(k.count('&'),k))
        if keys:
          key=keys[0]
          for r in weights:weights[r][key]=0
          induced.append({'rule':key,'reason':'complete-candidate feature collision','training_case':name,
            'schema':'two or three whitelisted chord/relative-motion/rhythm/source conditions; no song identity'})
          rows=[]
          for name,r,b,gold,gh,other,oh in comparisons:
            b.selector.weights=weights[r]
            rows.append(constraint(r,vector(b,gold,gh),vector(b,other,oh),name))
          continue
      result['solver']=last_solver;break
    if iteration+1<passes:weights={r:last_solver['weights'].get(r,{}) for r in weights}
  # Do not call an untested final solver update a successful checkpoint.
  result['solver']=last_solver
  result['constraints']=[{**r,'delta':{k:str(v) for k,v in r['delta'].items()},'prior':str(r['prior'])} for r in rows]
  return result