"""Export source-only complete candidates; references are not accepted here."""
import argparse
import copy
import dataclasses
from fractions import Fraction as F
import json
from pathlib import Path
import time
import cProfile

from audition.fit_selection import build
from generator import model, musicxml_writer, score, vocabulary
from generator.solo import notation, placement


def export(builder,document,output):
  output.mkdir(parents=True,exist_ok=True)
  source=score.source_staff(document)['measures']
  start_index=score.body_start_measure_index(source)
  if builder.start<0:start_index-=1
  length=F();selected=[]
  for m in source[start_index:]:
    if length>=builder.end-builder.start:break
    selected.append(m);length+=model.fraction(m['duration'])
  if length!=builder.end-builder.start:raise ValueError('Diagnostic export requires complete source measures')
  records=[]
  for index,state in enumerate(builder.completed,1):
    played=placement.place(state.notes,builder.analysis.tempo,profile=builder.harmony.profile)
    if played is None:raise ValueError('Completed candidate is unplayable')
    shifted=tuple(dataclasses.replace(n,start=n.start-builder.start,end=n.end-builder.start) for n in played[1])
    part=copy.deepcopy(document);staff=part['staves'][0]
    part['staves']=[staff];staff['measures']=notation.measures(selected,shifted,'candidate')
    staff['tuplets']=notation.tuplets(staff['measures']);staff['tablature']=True
    for key in ('scale_spans','octave_shifts','slurs','wedges'):staff.pop(key,None)
    path=output/f'candidate-{index}.xml'
    musicxml_writer.write_musicxml(path,part,tablature=True)
    records.append({'file':path.name,'automatic_winner':index==1,'evaluation':dataclasses.asdict(state.evaluation),
      'home':[(str(n.start),str(n.end),n.pitch.number()) for n in state.home],
      'events':[(str(n.start),str(n.end),n.pitch.number()) for n in state.notes],
      'actions':[{'kind':a.kind,'start':str(a.start),'end':str(a.end),'contracts':[c.rule for c in a.contracts]} for a in state.actions]})
  value={'source_only':True,'automatic_winner':'candidate-1.xml','start':str(builder.start),'end':str(builder.end),
    'candidates':records,'frontiers':builder.frontiers,'timings':dict(builder.timings),
    'note':'Diagnostic candidates, not publication or evidence of exact reproduction.'}
  (output/'candidates.json').write_text(json.dumps(value,indent=2,default=str)+'\n')
  return value


def main():
  p=argparse.ArgumentParser(description=__doc__)
  p.add_argument('--source',type=Path,required=True);p.add_argument('--role',choices=('acoustic_solo','electric_solo'),required=True)
  p.add_argument('--data-dir',type=Path,default=Path('data'));p.add_argument('--output',type=Path,required=True)
  p.add_argument('--cache-dir',type=Path);p.add_argument('--bars',type=int,default=4)
  p.add_argument('--profile',action='store_true')
  args=p.parse_args()
  if args.bars<1:p.error('--bars must be positive')
  with vocabulary.using(vocabulary.load(args.data_dir/'composition'),args.role):
    b,document=build(args.source,args.role,args.data_dir,args.cache_dir)
    ends=[x-b.analysis.body_offset for x in b.analysis.boundaries if x-b.analysis.body_offset>0]
    b.end=min(b.end,ends[min(args.bars-1,len(ends)-1)]);b.operators.end=b.end
    started=time.monotonic();seen=set()
    def progress(stage,at,state,proposed,retained):
      if stage=='frontier' and at.denominator==1 and at not in seen:
        seen.add(at);print(f'frontier {at}: {len(proposed)} histories, {len(retained)} retained; {time.monotonic()-started:.1f}s',flush=True)
    b.observe=progress
    args.output.mkdir(parents=True,exist_ok=True)
    profiler=cProfile.Profile() if args.profile else None
    if profiler:profiler.enable()
    try:b.compose()
    finally:
      if profiler:profiler.disable();profiler.dump_stats(args.output/'search.prof')
    export(b,document,args.output)
    print(json.dumps({'completed':len(b.completed),'seconds':b.timings['action_search_seconds'],'winner':str(b.winner_evaluation.total)}))

if __name__=='__main__':main()