"""Explicit, deterministic supervised training; never called by an asset build."""
from __future__ import annotations
import argparse
import dataclasses
import hashlib
import json
import os
from pathlib import Path
import tempfile
import time
from fractions import Fraction as F
from generator import harmony, model, musicxml, parts, training_io, vocabulary
from generator import harmonic_time as timeline
from generator.solo import analysis, learned


@dataclasses.dataclass(frozen=True)
class Span:
  start:F
  end:F
  chord:str


class Context:
  def __init__(self,document):
    self.spans=tuple(Span(s.start,s.end,harmony.symbol(s.chord) if s.chord else 'N.C.')
      for s in timeline.read(document['staves'][0]))
  def at(self,time):return next((s for s in self.spans if s.start<=time<s.end),None)


def load_examples(manifest):
  manifest=Path(manifest);value=json.loads(manifest.read_text())
  if value.get('version')!=1:raise ValueError('Unsupported training manifest')
  examples=[];provenance=[];seen=set();pairs=set()
  for entry in sorted(value['examples'],key=lambda e:e['id']):
    if entry['id'] in seen or entry['slot'] not in (1,2):raise ValueError('Duplicate example or invalid role')
    seen.add(entry['id']);role='acoustic' if entry['slot']==1 else 'electric'
    source=(manifest.parent/entry['source']).resolve();solo=(manifest.parent/entry['solo']).resolve()
    if not solo.is_relative_to(manifest.parent.resolve()):raise ValueError('Training solo must be inside the training directory')
    if not source.is_relative_to((manifest.parent.parent/'songs').resolve()):raise ValueError('Training sources must be inside data/songs')
    document=parts.prepare_document(musicxml.load_musicxml(source));a=analysis.analyze(document)
    from generator.production import SourceAnalysis
    pair=(SourceAnalysis.read(document).musical_key,entry['slot'])
    if pair in pairs:raise ValueError('Duplicate source music and solo slot')
    pairs.add(pair)
    evidence=training_io.read(solo,role,pitch_convention=entry['pitch_convention'])
    spans=[(F(m['start']),F(m['duration'])) for m in evidence['measures']]
    if spans!=[(x,y-x) for x,y in zip(a.boundaries,a.boundaries[1:])]:raise ValueError(entry['id']+': source/solo measure timing differs; no padding or pickup removal allowed')
    if evidence['discrepancies']:raise ValueError(entry['id']+': '+str(evidence['discrepancies']))
    ctx=Context(document)
    training_io.validate_harmony(document,evidence,entry['id'])
    events=[];cursor=F()
    from generator import fretboard
    for n in evidence['sounding_events']:
      start,end=F(n['start']),F(n['end']);number=F(n['sounding_midi'])
      if number.denominator!=1:raise ValueError(entry['id']+': microtonal guitar pitch is unsupported')
      pitch=int(number)
      if not fretboard.single_positions_for_pitch(F(pitch)):raise ValueError(entry['id']+': unplayable sounding pitch '+str(pitch)+' at '+str(start))
      if start<cursor or end<=start:raise ValueError(entry['id']+': non-monophonic/invalid event')
      if start>cursor:events.append((cursor,start,None))
      events.append((start,end,pitch));cursor=end
    if cursor<a.boundaries[-1]:events.append((cursor,a.boundaries[-1],None))
    source_hash=hashlib.sha256(source.read_bytes()).hexdigest()
    examples.append({'id':entry['id'],'song':source_hash,'role':role,'analysis':a,'harmony':ctx,'events':tuple(events)})
    provenance.append({'id':entry['id'],'source_sha256':source_hash,'solo_sha256':evidence['sha256'],
      'pitch_convention':entry['pitch_convention'],'measures':len(spans),'pickup_beats':str(a.body_offset),
      'slot':entry['slot'],'attacks':len(evidence['sounding_events']),'annotations':len(evidence['annotations']),
      'adapter_notes':['Trailing chord-display parentheses interpreted without modifying XML.'] if any(h['display'] and h['display'].endswith(')') and '(' not in h['display'] for h in evidence['harmonies']) else []})
  if not examples:raise ValueError('No training examples')
  return examples,{'description':value['provenance'],'examples':provenance,'training_version':learned.VERSION}


def likelihood(predictor,example):
  history=[];previous=None;cost=0
  for start,end,pitch in example['events']:
    harmonic,root,_=learned.chord_context(example['harmony'],start)
    action=learned.token(end-start,pitch,previous,root)
    relative=start-example['analysis'].body_offset
    c=learned.contexts(harmonic,learned.phase_at(example['analysis'],relative),learned.source_feature(example['analysis'],relative,root),tuple(history[-4:]))
    cost+=predictor.cost(action,c);history.append(action)
    if pitch is not None:previous=pitch
  return cost/len(history)


def atomic_json(path,value):
  path=Path(path);path.parent.mkdir(parents=True,exist_ok=True)
  with tempfile.NamedTemporaryFile('w',dir=path.parent,delete=False) as stream:
    json.dump(value,stream,sort_keys=True,separators=(',',':'));stream.write('\n');name=stream.name
  os.chmod(name,0o644)
  os.replace(name,path)


def register_example(manifest,source,solo,slot,pitch_convention):
  """Copy and validate a new pair before atomically publishing its manifest."""
  import shutil
  manifest=Path(manifest).resolve();root=manifest.parent
  sources=root.parent/'songs';root.mkdir(parents=True,exist_ok=True);sources.mkdir(parents=True,exist_ok=True)
  value=json.loads(manifest.read_text()) if manifest.exists() else {'version':1,'provenance':'User-authored paired song/solo training corpus.','examples':[]}
  source,solo=Path(source).resolve(),Path(solo).resolve()
  destinations=((source,sources/source.name),(solo,root/solo.name));created=[];temporary=None
  try:
    for original,destination in destinations:
      if destination.exists():
        if original.read_bytes()!=destination.read_bytes():raise ValueError(str(destination)+': already exists with different contents; no overwrite performed')
      else:shutil.copyfile(original,destination);created.append(destination)
    entry={'id':source.stem+'/solo'+str(slot),'source':'../songs/'+source.name,'solo':solo.name,
      'slot':slot,'pitch_convention':pitch_convention}
    if any(e['id']==entry['id'] for e in value['examples']):raise ValueError('Example already registered: '+entry['id'])
    value['examples'].append(entry)
    with tempfile.NamedTemporaryFile('w',dir=root,suffix='.json',delete=False) as stream:
      json.dump(value,stream);temporary=Path(stream.name)
    load_examples(temporary)
    atomic_json(manifest,value)
  except BaseException:
    for p in created:p.unlink()
    raise
  finally:
    if temporary:temporary.unlink(missing_ok=True)


def main():
  parser=argparse.ArgumentParser(description=__doc__)
  parser.add_argument('--manifest',type=Path,default=Path('data/training/manifest.json'))
  parser.add_argument('--output',type=Path,default=Path('data/models/solo.json'))
  parser.add_argument('--report',type=Path,default=Path('build/training/report.json'))
  parser.add_argument('--leave-one-song-out',action='store_true')
  parser.add_argument('--add-source',type=Path,help='Register a source song (copied to data/songs if needed).')
  parser.add_argument('--add-solo',type=Path,help='Register its authored solo (copied unchanged to data/training).')
  parser.add_argument('--slot',type=int,choices=(1,2),help='Arrangement mapping only: 1 acoustic, 2 electric; learning is pooled.')
  parser.add_argument('--pitch-convention',choices=('guitar-written','encoded'),default='guitar-written')
  args=parser.parse_args();began=time.monotonic()
  if any((args.add_source,args.add_solo,args.slot)):
    if not all((args.add_source,args.add_solo,args.slot)):parser.error('--add-source, --add-solo and --slot must be supplied together')
    register_example(args.manifest,args.add_source,args.add_solo,args.slot,args.pitch_convention)
  examples,provenance=load_examples(args.manifest)
  value=learned.fit(examples,provenance);predictor=learned.Model(value)
  report={'version':learned.VERSION,'training':provenance,'model_sha256':vocabulary.fingerprint(value),
    'contexts':len(value['contexts']),'in_sample_event_nll':{e['id']:likelihood(predictor,e) for e in examples},
    'limitations':'A small corpus does not establish generalization. Likelihood is teacher-forced, not free generation or musical preference.'}
  if args.leave_one_song_out:
    report['held_out_event_nll']={}
    for song in sorted({e['song'] for e in examples}):
      remaining=[e for e in examples if e['song']!=song]
      if not remaining:continue
      fitted=learned.Model(learned.fit(remaining,{'evaluation_only':True}))
      report['held_out_event_nll'].update({e['id']:likelihood(fitted,e) for e in examples if e['song']==song})
  atomic_json(args.output,value);report['training_seconds']=time.monotonic()-began;atomic_json(args.report,report)
  print(f'Trained {len(examples)} solos: {len(value["contexts"])} contexts in {report["training_seconds"]:.3f}s; {args.output}')


if __name__=='__main__':main()