"""Fit shared ranking from completed ordinary generation, never prefix roll-ins.

This development command reads explicit targets and saved complete alternatives.
It writes a *proposal*, not a released model. A staged data copy may be passed to
the ordinary generator for a separate target-free cold reproduction test.
"""
import argparse
import copy
from fractions import Fraction as F
import hashlib
import json
from pathlib import Path
import shutil
import sys

from audition.fit_selection import build,derive
from audition.reference_review import read
from audition.structured_learning import vector,constraint,solve,sounded,CoverageTrace
from generator import production,score,vocabulary
from generator.solo import construction,selection,policy


def evaluation_vector(value):
  if value['version']!=selection.VERSION:raise ValueError('Completed output uses another evaluator version')
  span=F(value['denominators']['span'])
  features={k:F(v)*span for k,v in value['raw'].items() if not k.startswith('prior/') and F(v)}
  selection.validate_weights(dict.fromkeys(features,0))
  prior=sum(F(v)*span*selection.DEFAULT[k.removeprefix('prior/')] for k,v in value['raw'].items() if k.startswith('prior/'))
  return features,prior


def training_context(data):
  """Changing learned weights is allowed between rounds; musical rules are not."""
  musical=copy.deepcopy(data)
  musical['roles'].pop('provenance',None)
  for role in musical['roles']['roles'].values():role.pop('selection_weights',None)
  return {'evaluator':selection.VERSION,'algorithm':policy.DEFAULT.composition_version,
    'fixed_priors':selection.DEFAULT,'musical_data':vocabulary.fingerprint(musical)}


def previous_constraints(evidence,cases,context):
  if evidence.get('schema')!='complete-candidate-training-1' or evidence.get('context')!=context:
    raise ValueError('Previous round has incompatible evaluation or musical data')
  earlier=evidence['training_dependency']
  if earlier.keys()!=cases.keys():raise ValueError('Previous round has different training cases')
  for name in cases:
    for key in ('source_sha256','target_sha256','role'):
      if cases[name][key]!=earlier[name][key]:raise ValueError('Previous round changed '+name+'/'+key)
  rows=[]
  for row in evidence['constraints']:
    if row['case'] not in cases or row['role']!=cases[row['case']]['role']:
      raise ValueError('Previous constraint has a mismatched case/role')
    selection.validate_weights(dict.fromkeys(row['delta'],0))
    rows.append({**row,'delta':{k:F(v) for k,v in row['delta'].items()},'prior':F(row['prior'])})
  return rows


def main():
  p=argparse.ArgumentParser(description=__doc__)
  for name in ('source','reference','role'):p.add_argument('--'+name,action='append',required=True,metavar='CASE=VALUE')
  p.add_argument('--form-start',action='append',default=[])
  p.add_argument('--data-dir',type=Path,default=Path('data'))
  p.add_argument('--completed-dir',type=Path,required=True,help='Completed ordinary assets directory containing analysis/*.roles')
  p.add_argument('--output',type=Path,required=True);p.add_argument('--cache-dir',type=Path)
  p.add_argument('--solver-path',type=Path,required=True)
  p.add_argument('--stage-data',action='store_true',help='Write an isolated data copy for ordinary inference; does not publish it')
  p.add_argument('--previous-round',type=Path,action='append',default=[],help='Accumulate complete-output constraints from a compatible training.json; never a production input')
  args=p.parse_args();sys.path.insert(0,str(args.solver_path.resolve()))
  pairs=lambda values:dict(v.split('=',1) for v in values)
  sources,refs,roles=map(pairs,(args.source,args.reference,args.role));starts=pairs(args.form_start)
  if sources.keys()!=refs.keys() or sources.keys()!=roles.keys():p.error('Cases must agree')
  args.output.mkdir(parents=True,exist_ok=True);rows=[];cases={};contexts=[];data=vocabulary.load(args.data_dir/'composition')
  for name in sorted(sources):
    with vocabulary.using(data,roles[name]):
      b,document=build(sources[name],roles[name],args.data_dir,args.cache_dir)
      target=read(refs[name],name,int(starts[name]) if name in starts else None,'guitar-written')
      gold,coverage=derive(b,target)
      if not coverage['complete']:raise ValueError(name+': '+coverage['status'])
      s=construction.State()
      for a in gold:s=b.advance(s,None,a)
      positive=vector(b,gold,s.home)
      contexts.append((name,roles[name],b,gold))
      path=args.completed_dir/'analysis'/(score.slugify(Path(sources[name]).stem)+'.roles')
      generated=json.loads(path.read_text())[roles[name]]
      if generated['trace']['source']!=production.SourceAnalysis.read(document).musical_key:
        raise ValueError(name+': completed candidates belong to different canonical source music')
      if generated['trace']['role']!=roles[name]:raise ValueError(name+': wrong completed role')
      if generated['trace']['vocabulary']!=vocabulary.fingerprint(production.dependencies(roles[name],data)):
        raise ValueError(name+': completed candidates use different production data/settings')
      candidates=generated['details']['selection_trace']['completed_candidates']
      negatives=0;exact=False
      for c in candidates:
        stream=tuple((F(a),F(z),int(n)) for a,z,n in c['events'])
        same=stream==sounded(gold)
        if c['rank']==1:exact=same
        if not same:rows.append(constraint(roles[name],positive,evaluation_vector(c['evaluation']),name));negatives+=1
      cases[name]={'source':sources[name],'reference':refs[name],'role':roles[name],
        'source_sha256':hashlib.sha256(Path(sources[name]).read_bytes()).hexdigest(),
        'target_sha256':target['sha256'],'completed_artifact_sha256':hashlib.sha256(path.read_bytes()).hexdigest(),
        'coverage':coverage,'ordinary_winner_exact':exact,'complete_negative_candidates':negatives}
      print(name,exact,negatives,flush=True)
  context=training_context(data);parents=[]
  for path in args.previous_round:
    raw=path.read_bytes();rows.extend(previous_constraints(json.loads(raw),cases,context))
    parents.append({'path':str(path),'sha256':hashlib.sha256(raw).hexdigest()})
  unique={}
  for row in rows:
    key=(row['role'],row['case'],tuple(sorted(row['delta'].items())),row['prior'])
    unique[key]=row
  rows=[unique[key] for key in sorted(unique)]
  result=solve(rows)
  evidence={'schema':'complete-candidate-training-1','training_dependency':cases,'solver':result,
    'context':context,'previous_rounds':parents,
    'constraints':[{**r,'delta':{k:str(v) for k,v in r['delta'].items()},'prior':str(r['prior'])} for r in rows],
    'released':False,'acceptance':'Exact four full outputs from ordinary cold target-free generation, plus transfer; not satisfied by this training solve.'}
  if result['status']=='solved':
    # A zero here proves this proposed model cannot reproduce the full output.
    # A nonzero result proves neither full retention nor exact generation.
    # This is read-only diagnosis, not prefix fitting or candidate injection.
    cuts={}
    for name,role,b,gold in contexts:
      with vocabulary.using(data,role):
        b.selector.weights=result['weights'].get(role,{})
        trace=CoverageTrace(gold);b.observe=trace
        retained=b.candidates(b.start,construction.State());b.observe=None
        count=sum(trace.matches(a.notes,a.end) for _,a in retained)
        cuts[name]={'matching_first_successors':count,'observed_failures':trace.records,
          'full_reproduction_already_excluded':count==0}
    evidence['post_fit_cutoff_diagnosis']=cuts
  (args.output/'training.json').write_text(json.dumps(evidence,indent=2)+'\n')
  if result['status']=='solved':
    (args.output/'proposed-weights.json').write_text(json.dumps(result['weights'],indent=2)+'\n')
    if args.stage_data:
      staged=args.output/'data'
      if staged.exists():raise ValueError('Refusing to overwrite a previous staged data proposal')
      shutil.copytree(args.data_dir,staged)
      path=staged/'composition'/'roles.json';value=json.loads(path.read_text())
      for role,weights in result['weights'].items():value['roles'][role]['selection_weights']=weights
      value['provenance']='Shared bounded coefficients learned from explicit authored training examples; see adjacent development training.json. No runtime target reads or stored phrases.'
      path.write_text(json.dumps(value,indent=2)+'\n')
  print(json.dumps(result,indent=2))
  return 0 if result['status']=='solved' else 2

if __name__=='__main__':raise SystemExit(main())