"""Explicit, deterministic supervised training; never called by an asset build."""
from __future__ import annotations
import argparse
import collections
import dataclasses
import fractions
import hashlib
import json
import os
import pathlib
import shutil
import tempfile
import time

import generator.fretboard
import generator.harmonic_time
import generator.harmony
import generator.musicxml
import generator.parts
import generator.solo.authored
import generator.solo.learned
import generator.source
import generator.training_io
import generator.vocabulary

@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.source.analyze(document)
    pair = (generator.source.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()