"""Explicit, deterministic supervised training; never called by an asset build."""
from __future__ import annotations
import argparse
import collections
import dataclasses
import hashlib
import json
import os
import pathlib
import tempfile
import time
import fractions
import generator.harmony
import generator.model
import generator.musicxml
import generator.parts
import generator.training_io
import generator.vocabulary
import generator.harmonic_time
import generator.solo.analysis
import generator.solo.learned
import shutil
import generator.production
import generator.solo.authored
import generator.fretboard

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

class Context:
  def __init__(self, document):
    self.spans = tuple(Span(s.start, s.end, generator.harmony.symbol(s.chord) if s.chord else 'N.C.')
      for s in generator.harmonic_time.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 = pathlib.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 = generator.parts.prepare_document(generator.musicxml.load_musicxml(source))
    a = generator.solo.analysis.analyze(document)
    pair = (generator.production.SourceAnalysis.read(document).musical_key, entry['slot'])
    if pair in pairs:
      raise ValueError('Duplicate source music and solo slot')
    pairs.add(pair)
    evidence = generator.training_io.read(solo, role, pitch_convention=entry['pitch_convention'])
    spans = [(fractions.Fraction(m['start']), fractions.Fraction(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)
    generator.training_io.validate_harmony(document, evidence, entry['id'])
    checked = generator.solo.authored.validate_evidence(document, evidence)
    events = []
    cursor = fractions.Fraction()
    for n in evidence['sounding_events']:
      start, end = fractions.Fraction(n['start']), fractions.Fraction(n['end'])
      number = fractions.Fraction(n['sounding_midi'])
      if number.denominator != 1:
        raise ValueError(entry['id'] + ': microtonal guitar pitch is unsupported')
      pitch = int(number)
      if not generator.fretboard.single_positions_for_pitch(fractions.Fraction(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),
      'harmonic_validation': {'valid': checked.valid, 'treatments': dict(sorted(collections.Counter(c.rule for c in checked.contracts).items()))}})
    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': generator.solo.learned.VERSION}

def likelihood(predictor, example):
  history = []
  previous = None
  cost = 0
  for start, end, pitch in example['events']:
    harmonic, root, _ = generator.solo.learned.chord_context(example['harmony'], start)
    action = generator.solo.learned.token(end - start, pitch, previous, root)
    relative = start - example['analysis'].body_offset
    c = generator.solo.learned.contexts(harmonic, generator.solo.learned.phase_at(example['analysis'], relative), generator.solo.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 = pathlib.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."""
  manifest = pathlib.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 = pathlib.Path(source).resolve(), pathlib.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 = pathlib.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=pathlib.Path, default=pathlib.Path('data/training/manifest.json'))
  parser.add_argument('--output', type=pathlib.Path, default=pathlib.Path('data/models/solo.json'))
  parser.add_argument('--report', type=pathlib.Path, default=pathlib.Path('build/training/report.json'))
  parser.add_argument('--leave-one-song-out', action='store_true')
  parser.add_argument('--add-source', type=pathlib.Path, help='Register a source song (copied to data/songs if needed).')
  parser.add_argument('--add-solo', type=pathlib.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 = generator.solo.learned.fit(examples, provenance)
  predictor = generator.solo.learned.Model(value)
  report = {'version': generator.solo.learned.VERSION, 'training': provenance, 'model_sha256': generator.vocabulary.fingerprint(value),
    'contexts': len(value['contexts']), 'in_sample_event_nll': {e['id']: likelihood(predictor, e) for e in examples},
    'harmonic_validation': {e['id']: e['harmonic_validation'] 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 = generator.solo.learned.Model(generator.solo.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()