"""Explicit development-only shared-rule derivation and fitting.

Targets are read only by this command. Production imports neither this module
nor its artifacts. A derivation uses the exact same enumerated Actions as the
normal composer; accepting injected pitches is not a construction proof.
"""
from __future__ import annotations
import argparse
import bisect
import dataclasses
from fractions import Fraction as F
import hashlib
import heapq
import json
from pathlib import Path
import time

from audition.reference_review import read
from audition.solo_compare import events
from generator import accompaniment, musicxml, parts, production, resources, score, vocabulary
from generator.arrange import form, generate
from generator.scales import selection, timeline
from generator.solo import analysis, construction, harmonic, profiles


def build(source_path,role,data_dir,cache_dir):
  """Read original song/form and production data; never reference harmony."""
  document=parts.prepare_document(musicxml.load_musicxml(source_path))
  a=analysis.analyze(document);part=role.split('_')[0]
  source=score.source_staff(document)['measures']
  sections=form.full_arrangement_sections(source,part);si=3 if part=='acoustic' else 4
  section=sections[si]
  effective=generate.section_source(source,section,sections[si+1])
  measures=effective[section.source_measure_index:section.source_measure_index+section.measure_count]
  initial=accompaniment.carried_chords(effective)[section.source_measure_index]
  from generator.solo.arrangement import regions
  first,_,offset,_=regions(document)[si]
  start,end=next((a-offset,b-offset) for a,b,r in production.windows(document,part) if r==role)
  if start<0:
    template=generate._arrange_document(document,part,{},owned={'slash'})
    measures=[template['staves'][0]['measures'][first-1],*measures]
    initial=accompaniment.carried_chords(template['staves'][0]['measures'])[first-1]
  sc,seq,rh=resources.load(data_dir,cache_dir=cache_dir)
  choices=selection.select(timeline.read({'measures':measures},initial=initial),timeline.Solo(F(),end-start,role),sc)
  choices=tuple(dataclasses.replace(c,span=dataclasses.replace(c.span,start=c.span.start+start,end=c.span.end+start)) for c in choices)
  h=harmonic.Harmony(choices,profiles.for_part(part,a))
  return construction.Builder(a,measures,h,seq,rh,end,'development',start=start),document


def canonical(notes):
  return tuple((n.start,n.end,n.pitch.number()) for n in notes)


def derive(builder,target):
  """Stateful target-constrained coverage probe; NEVER production success.

  Separate histories carry actual pending contracts. This bounded diagnostic
  can report exhaustion, not prove a phrase impossible. No notes are injected.
  """
  encoded=tuple((F(n['start']),F(n['end']),int(n['pitch'])) for n in events(target))
  starts=tuple(n[0] for n in encoded);encoded_set=set(encoded)
  queue=[builder.start];chart={builder.start:[construction.State()]};queued={builder.start}
  furthest=builder.start;last=construction.State();cuts=0
  while queue:
    at=heapq.heappop(queue);states=chart.pop(at);furthest=max(furthest,at)
    unique={builder.identity(s):s for s in states};states=list(unique.values())
    cuts+=max(0,len(states)-16);states=states[:16]
    if at==builder.end:
      for state in states:
        checked=builder.harmony.evaluate(state.notes,expressions=tuple(c for a in state.actions for c in a.contracts),complete=True,until=at)
        if checked.valid and canonical(state.notes)==encoded:
          return state.actions,{'complete':True,'covered_until':str(at),'actions':len(state.actions),
            'status':'statefully_constructible','diagnostic_only':True,'pruned_histories':cuts}
    left=bisect.bisect_left(starts,at)
    def matches(action):
      if any((c.target.start,c.target.end,c.target.pitch.number()) not in encoded_set for c in action.contracts):return False
      expected=encoded[left:bisect.bisect_left(starts,action.end)]
      return not (left and encoded[left-1][1]>at) and canonical(action.notes)==expected
    first=encoded[left] if left<len(encoded) else None
    actions=builder.operators.at(at,eligible=matches,first=first)
    for state in states:
      last=state
      for action in actions:
        next_state=builder.advance(state,None,action)
        if next_state is None:continue
        chart.setdefault(action.end,[]).append(next_state)
        if action.end not in queued:heapq.heappush(queue,action.end);queued.add(action.end)
  next_event=next((n for n in encoded if n[0]>=furthest),None)
  return last.actions,{'complete':False,'covered_until':str(furthest),
    'status':'coverage_search_exhausted' if cuts else 'not_proposed_or_rejected_by_rule',
    'next_event':tuple(map(str,next_event)) if next_event else None,
    'pruned_histories':cuts,'diagnostic_only':True,'action_coverage':dict(builder.operators.coverage)}


def fit(examples,passes,on_checkpoint=None,resources=None):
  # Imported only by this explicitly requested development operation.
  from audition.structured_learning import fit_complete
  return fit_complete(examples,passes,on_checkpoint,resources)


def main():
  p=argparse.ArgumentParser(description=__doc__)
  p.add_argument('operation',choices=('explain','fit'))
  p.add_argument('--source',action='append',required=True,metavar='CASE=PATH')
  p.add_argument('--reference',action='append',required=True,metavar='CASE=PATH')
  p.add_argument('--role',action='append',required=True,metavar='CASE=acoustic_solo|electric_solo')
  p.add_argument('--form-start',action='append',default=[],metavar='CASE=MEASURE')
  p.add_argument('--data-dir',type=Path,default=Path('data'))
  p.add_argument('--cache-dir',type=Path)
  p.add_argument('--output',type=Path,required=True)
  p.add_argument('--passes',type=int,default=8)
  p.add_argument('--solver-path',type=Path,help='Development-only pinned solver installation directory')
  args=p.parse_args()
  if args.solver_path:
    import sys
    sys.path.insert(0,str(args.solver_path.resolve()))
  if not 1<=args.passes<=200:p.error('passes must be 1–200')
  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('source/reference/role cases must match')
  args.output.mkdir(parents=True,exist_ok=True);results={};examples=[]
  data=vocabulary.load(args.data_dir/'composition');began=time.monotonic()
  for name in sorted(sources):
    with vocabulary.using(data,roles[name]):
      b,_=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')
      path,result=derive(b,target)
      results[name]={**result,'reference_sha256':target['sha256'],'source_sha256':hashlib.sha256(Path(sources[name]).read_bytes()).hexdigest(),
        'derivation':[{'kind':a.kind,'start':str(a.start),'end':str(a.end),'events':[(str(n.start),str(n.end),n.pitch.number()) for n in a.notes],'rules':[c.rule for c in a.contracts]} for a in path]}
      if args.operation=='fit' and result['complete']:examples.append((name,roles[name],b,path))
      print(name,result['status'],result['covered_until'],flush=True)
  result={'schema':'shared-phrase-learning-2','development_only':True,'coverage_seconds':time.monotonic()-began,'cases':results}
  if args.operation=='fit':
    if not all(r['complete'] for r in results.values()):result['fit_status']='blocked_by_construction_coverage'
    else:
      def checkpoint(value):
        path=args.output/'checkpoint.json';temporary=path.with_suffix('.tmp')
        temporary.write_text(json.dumps(value,indent=2)+'\n');temporary.replace(path)
      fitting_start=time.monotonic()
      result['fit']=fit(examples,args.passes,checkpoint,data)
      result['fit_seconds']=time.monotonic()-fitting_start
  result['total_seconds']=time.monotonic()-began
  (args.output/'coverage.json').write_text(json.dumps(result,indent=2)+'\n')
  return 0 if all(r['complete'] for r in results.values()) and result.get('fit',{'complete':True})['complete'] else 2

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