"""Read optional reference MusicXML literally, without repairing it or composing."""
from __future__ import annotations
import argparse, collections, html, os
import hashlib
import json
import pathlib
import fractions
import xml.etree.ElementTree

import generator.training_io
import generator.fretboard
import generator.solo.placement
import generator.solo.types

def events(review):
  offset = fractions.Fraction(review['alignment']['encoded_offset'] or 0)
  return [dict(start=str(fractions.Fraction(n['start']) - offset), end=str(fractions.Fraction(n['end']) - offset),
    pitch=n['sounding_midi'], spelling=n['pitch'], measure=n['measure']) for n in review['sounding_events']]

def rests(review):
  offset = fractions.Fraction(review['alignment']['encoded_offset'] or 0)
  cursor = -offset
  result = []
  for n in sorted(events(review), key=lambda n: fractions.Fraction(n['start'])):
    start, end = fractions.Fraction(n['start']), fractions.Fraction(n['end'])
    if start > cursor:
      result.append((str(cursor), str(start)))
    cursor = max(cursor, end)
  finish = max((fractions.Fraction(m['start']) + fractions.Fraction(m['duration']) for m in review['measures']), default=fractions.Fraction()) - offset
  if cursor < finish:
    result.append((str(cursor), str(finish)))
  return result

def playability(review, tempo=fractions.Fraction(120)):
  notes = []
  conflicts = []
  for i, n in enumerate(review['sounding_events']):
    p = n['pitch']
    shift = fractions.Fraction(n['sounding_midi']) - fractions.Fraction(n['encoded_midi'])
    if shift % 12:
      raise ValueError('Guitar comparison expects octave-only transposition')
    pitch = generator.solo.types.Pitch(p['step'].lower(), fractions.Fraction(p['alter']), p['octave'] + int(shift / 12))
    note = generator.solo.types.Note(fractions.Fraction(n['start']), fractions.Fraction(n['end']), pitch, event_id=str(i))
    notes.append(note)
    if not generator.fretboard.single_positions_for_pitch(fractions.Fraction(pitch.number())):
      conflicts.append({'measure': n['measure'], 'written': p, 'sounding_midi': pitch.number(),
        'constraint': 'No standard-tuning position on strings 1–6, frets 1–12, without open strings.'})
  overlaps = [{'first': a.event_id, 'second': b.event_id} for a, b in zip(notes, notes[1:]) if a.end > b.start]
  placed = generator.solo.placement.place(tuple(notes), tempo) if notes and not conflicts and not overlaps else None
  return {'conflicts': conflicts, 'overlaps': overlaps, 'complete_path_found': placed is not None,
    'range': [min(n.pitch.number() for n in notes), max(n.pitch.number() for n in notes)] if notes else [],
    'positions': [list(n.position[:2]) for n in placed[1]] if placed else [],
    'limitation': 'Legal placement only. Shift cost is a preference, not a verified human speed limit.'}

def musical_summary(review):
  ns = events(review)
  durations = collections.Counter(str(fractions.Fraction(n['end']) - fractions.Fraction(n['start'])) for n in ns)
  intervals = collections.Counter(str(fractions.Fraction(b['pitch']) - fractions.Fraction(a['pitch'])) for a, b in zip(ns, ns[1:]))
  holds = []
  for event, n in zip(review['sounding_events'], ns):
    if fractions.Fraction(n['end']) - fractions.Fraction(n['start']) < 2:
      continue
    spans = [s for i in event['notation_events'] for s in review['notes'][i]['harmonic_spans']]
    holds.append({**n, 'harmonic_spans': spans})
  return {'attacks': len(ns), 'duration_counts': dict(durations), 'interval_counts': dict(intervals),
    'rests': rests(review), 'sustained_events': holds,
    'annotations': review['annotations'], 'interpretation': 'Post-generation diagnostics; does not train or reselect music.'}

def compare(generated, reference):
  a, b = events(generated), events(reference)

  def grouped(ns):
    result = collections.defaultdict(list)
    for n in ns:
      result[n['start']].append(n)
    return result

  left, right = grouped(a), grouped(b)
  differences = []
  spelling = []
  equal = 0
  for start in sorted(left.keys() | right.keys(), key=fractions.Fraction):
    x, y = left[start], right[start]

    def sounded(ns):
      return sorted((n['pitch'], n['end']) for n in ns)

    if sounded(x) != sounded(y):
      differences.append({'time': start, 'generated': x, 'reference': y,
        'context': 'reference lead-in before form' if fractions.Fraction(start) < 0 else 'aligned form'})
    else:
      equal += len(x)
      if [n['spelling'] for n in x] != [n['spelling'] for n in y]:
        spelling.append({'time': start, 'generated': x, 'reference': y})
  rest_difference = rests(generated) != rests(reference)
  return {'exact_musical_match': not differences and not rest_difference, 'exact_matching_attacks': equal,
    'relationships': relationships(a, b),
    'differences': differences, 'spelling_differences': spelling, 'rest_difference': rest_difference,
    'generated': musical_summary(generated), 'reference': musical_summary(reference),
    'reference_playability': playability(reference), 'reference_original': reference}

def relationships(a, b):
  """Read-only comparisons on the aligned clock; no fitted aggregate loss."""

  def distance(left, right):
    row = list(range(len(right) + 1))
    for i, x in enumerate(left, 1):
      next_row = [i]
      for j, y in enumerate(right, 1):
        next_row.append(min(next_row[-1] + 1, row[j] + 1, row[j - 1] + (x != y)))
      row = next_row
    return row[-1]

  ns = [[dict(start=fractions.Fraction(n['start']), end=fractions.Fraction(n['end']), pitch=fractions.Fraction(n['pitch'])) for n in notes] for notes in (a, b)]
  points = sorted({n[k] for notes in ns for n in notes for k in ('start', 'end')})
  occupied = equal = classes = fractions.Fraction()
  for start, end in zip(points, points[1:]):
    pitches = [{n['pitch'] for n in notes if n['start'] <= start < n['end']} for notes in ns]
    if any(pitches):
      occupied += end - start
      if pitches[0] == pitches[1]:
        equal += end - start
      if {p % 12 for p in pitches[0]} == {p % 12 for p in pitches[1]}:
        classes += end - start
  pitches = [[n['pitch'] for n in notes] for notes in ns]
  rhythms = [[(n['start'], n['end'] - n['start']) for n in notes] for notes in ns]
  return {'pitch_sequence_edit_distance': distance(*pitches),
    'timed_attack_edit_distance': distance(*rhythms),
    'normalization_attacks': max(len(a), len(b)),
    'same_sounding_pitch_beats': str(equal), 'same_pitch_class_beats': str(classes),
    'either_sounding_beats': str(occupied),
    'interpretation': 'Separate descriptive comparisons, not a quality score. Pitch distance ignores timing; timing distance ignores pitch. Held overlap uses the original aligned register and includes lead-ins.'}

def run(generated_dir, reference_paths, starts, pitch_convention, output, baseline_dir=None):
  generated_dir, output = pathlib.Path(generated_dir), pathlib.Path(output)
  manifest = json.loads((generated_dir / 'manifest.solo').read_text())
  output.mkdir(parents=True, exist_ok=True)
  baseline_manifest = json.loads((pathlib.Path(baseline_dir) / 'manifest.solo').read_text()) if baseline_dir else None
  result = {'comparison_only': True, 'feeds_generation': False, 'solos': {}}
  for role, path in sorted(reference_paths.items()):
    item = manifest['solos'][role]
    actual = generator.training_io.read(generated_dir / item['file'], role, item['form_start_measure'])
    expected = generator.training_io.read(path, role, starts.get(role), pitch_convention)
    review = compare(actual, expected)
    baseline = None
    if baseline_manifest:
      old = baseline_manifest['solos'][role]
      baseline = generator.training_io.read(pathlib.Path(baseline_dir) / old['file'], role, old['form_start_measure'])
      review['baseline'] = musical_summary(baseline)
      review['baseline_relationships'] = relationships(events(baseline), events(expected))
    review['production_trace'] = item['trace']
    review['production_phrases'] = item.get('phrase_report', {})
    candidates = review['production_phrases'].get('selection_trace', {}).get('opening_candidates', [])
    # Preserve the real role cutoff; reference material beyond it is not trimmed.
    phrases = review['production_phrases'].get('phrases', [])
    if phrases:
      cutoff = max(fractions.Fraction(p['end']) for p in phrases)
      review['production_end'] = str(cutoff)
      review['reference_events_after_production_end'] = [n for n in events(expected) if fractions.Fraction(n['end']) > cutoff]
    result['solos'][role] = review
  output.mkdir(parents=True, exist_ok=True)
  (output / 'comparison.json').write_text(json.dumps(result, indent=2) + '\n')
  rows = ['<!doctype html><meta charset="utf-8"><title>Solo comparison</title>',
    '<style>body{font:16px system-ui;max-width:1100px;margin:2rem auto}td,th{padding:.5rem;text-align:left}pre{white-space:pre-wrap}</style>',
    '<h1>Completed solos and reference scores</h1><p>Differences are diagnostic. This comparison is read-only; registered authored scores are explicit repertoire inputs, and pooled examples supply training evidence. No listening verdict is claimed.</p>',
    '<p><a href="comparison.json">Exact timing, notes, harmony, annotations, placement, and phrase reports</a></p>']
  for role, r in result['solos'].items():
    actual_path = generated_dir / manifest['solos'][role]['file']
    reference_path = pathlib.Path(reference_paths[role])
    rows += ['<h2>' + html.escape(role) + '</h2>', f'<p>Exact event match: <strong>{"yes" if r["exact_musical_match"] else "no"}</strong>. {r["exact_matching_attacks"]} exactly matching attacks; {len(r["differences"])} differing onset positions. This is not a musical-quality score.</p>',
      '<p><a href="' + html.escape(os.path.relpath(actual_path, output)) + '">Generated MusicXML</a> · <a href="' + html.escape(os.path.relpath(reference_path, output)) + '">Untouched reference MusicXML</a></p>']
    if baseline_manifest:
      baseline_path = pathlib.Path(baseline_dir) / baseline_manifest['solos'][role]['file']
      rows.append('<p><a href="' + html.escape(os.path.relpath(baseline_path, output)) + '">Previous generated MusicXML</a></p>')
    rows.append('<table><tr><th>Diagnostic against reference</th><th>Current</th><th>Baseline</th></tr>')
    for key in ('pitch_sequence_edit_distance', 'timed_attack_edit_distance', 'same_sounding_pitch_beats', 'same_pitch_class_beats', 'either_sounding_beats'):
      rows.append('<tr><td>' + key.replace('_', ' ') + '</td><td>' + str(r['relationships'][key]) + '</td><td>' + str(r.get('baseline_relationships', {}).get(key, '—')) + '</td></tr>')
    rows.append('</table><p>Lower edit distance means fewer insertions, deletions or substitutions on that axis. Sounding agreement is measured in quarter-note beats, without octave shifting or fitted alignment. These do not measure musical quality.</p>')
    selected = collections.Counter(c['kind'] for c in r['production_phrases'].get('expressions', []))
    rows.append('<p>Validated production treatments: ' + html.escape(str(dict(selected))) + '. Registered authored music is preserved; harmonic relationships are recognized from notes and chords without requiring explanatory annotations.</p>')
    if r.get('reference_events_after_production_end'):
      rows.append('<p>Reference material extends beyond this exported production allocation: ' + str(len(r['reference_events_after_production_end'])) + ' events. The reference and source timing remain unchanged.</p>')
    funnel = r['production_phrases'].get('selection_trace', {}).get('operation_coverage', {})
    if funnel:
      stages = tuple(dict.fromkeys(k for counts in funnel.values() for k in counts))
      rows.append('<details><summary>Construction coverage (not similarity scoring)</summary><table><tr><th>Purpose</th>' + ''.join('<th>' + k.replace('_', ' ') + '</th>' for k in stages) + '</tr>')
      for purpose, counts in sorted(funnel.items()):
        rows.append('<tr><td>' + html.escape(purpose) + '</td>' + ''.join('<td>' + str(counts.get(k, 0)) + '</td>' for k in stages) + '</tr>')
      rows.append('</table></details>')
    candidates = r['production_phrases'].get('selection_trace', {}).get('opening_candidates', [])
    if candidates:
      rows.append('<details open><summary>Completed source-only opening candidates</summary><p>All candidates were composed without the references. Selection accounts for later phrases; the final choice need not have the cheapest opening alone.</p><table><tr><th>Candidate</th><th>Final choice</th><th>Purpose</th><th>Score contributions</th></tr>')
      for index, c in enumerate(candidates, 1):
        rows.append('<tr>' + ''.join('<td>' + html.escape(str(v)) + '</td>' for v in (index, c['selected_in_final_solo'], c['purpose'], c['evaluation']['contributions'])) + '</tr>')
      rows.append('</table></details>')
    rows.append('<details><summary>Exact note differences</summary><table><tr><th>Beat relative to form</th><th>Generated</th><th>Reference</th></tr>')
    for d in r['differences']:

      def describe(ns): return '; '.join(f'MIDI {n["pitch"]}, until {n["end"]}' for n in ns) or 'No attack'

      rows.append('<tr>' + ''.join('<td>' + html.escape(v) + '</td>' for v in (d['time'], describe(d['generated']), describe(d['reference']))) + '</tr>')
    rows.append('</table></details>')
  (output / 'index.html').write_text('\n'.join(rows))
  return result

def main():
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument('--generated-dir', type=pathlib.Path, required=True)
  parser.add_argument('--reference', action='append', default=[], metavar='SONG/ROLE=PATH')
  parser.add_argument('--form-start', action='append', default=[], metavar='SONG/ROLE=MEASURE')
  parser.add_argument('--reference-pitch', choices=('encoded', 'guitar-written'), default='encoded')
  parser.add_argument('--output', type=pathlib.Path, required=True)
  parser.add_argument('--baseline-dir', type=pathlib.Path, help='Optional previous completed solo exports; never generation inputs.')
  parser.add_argument('--require-exact', action='store_true', help='Fail unless every requested reference has identical complete sounding events and rests.')
  args = parser.parse_args()
  if args.require_exact and not args.reference:
    parser.error('--require-exact requires at least one explicit --reference')
  result = run(args.generated_dir, dict(x.split('=', 1) for x in args.reference),
    {k: int(v) for k, v in (x.split('=', 1) for x in args.form_start)}, args.reference_pitch, args.output, args.baseline_dir)
  for role, review in result['solos'].items():
    print(f'{role}: exact={review["exact_musical_match"]}; matching attacks={review["exact_matching_attacks"]}; differing onsets={len(review["differences"])}; rests differ={review["rest_difference"]}')
  return 1 if args.require_exact and any(not r['exact_musical_match'] for r in result['solos'].values()) else 0

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