from __future__ import annotations

import dataclasses
import pathlib
import types
import functools

import generator.model
import generator.musicxml

@dataclasses.dataclass(frozen=True)
class Sequence:
  id: str
  steps: tuple[int, ...]
  @functools.lru_cache(maxsize=32768)
  def offsets(self) -> tuple[int, ...]:
    result = [0]
    for step in self.steps:
      result.append(result[-1] + step)
    return tuple(result)
  def rank(self) -> tuple:
    offsets = self.offsets()
    return max(map(abs, self.steps)), max(offsets) - min(offsets), sum(a * b < 0 for a, b in zip(self.steps, self.steps[1:])), self.id

def diverse(items, rank, category, limit):
  """Round-robin eligible musical families before a destructive cutoff."""
  groups = {}
  # Rank once: hashing full event provenance and classifying interval paths
  # used to be repeated for every retained layer. Preserve stable tie order.
  decorated = sorted(((rank(item), i, item) for i,item in enumerate(items)),key=lambda v:v[:2])
  for value in decorated:
    groups.setdefault(category(value[2]), []).append(value)
  result, index = [], 0
  while len(result) < limit:
    layer = [values[index] for values in groups.values() if len(values) > index]
    if not layer:
      break
    # Equal ranks retain category insertion order, just as the original stable
    # layer sort did; the input index is only for the initial sort.
    result.extend(v[2] for v in sorted(layer,key=lambda v:v[0])[:limit-len(result)])
    index += 1
  return tuple(result)

@dataclasses.dataclass(frozen=True, init=False)
class Library:
  paths: types.MappingProxyType
  def __init__(self, sequences: tuple[Sequence, ...]):
    paths, identities = {}, set()
    for sequence in sequences:
      if sequence.id in identities or not 1 <= len(sequence.steps) <= 3 or any(not 1 <= abs(step) <= 7 for step in sequence.steps):
        raise ValueError(f'{sequence.id}: invalid or duplicate sequence')
      identities.add(sequence.id)
      paths.setdefault(len(sequence.steps) + 1, []).append(sequence)
    object.__setattr__(self, 'paths', types.MappingProxyType({count: tuple(sorted(items, key=lambda s: s.rank())) for count, items in paths.items()}))
  def connecting(self, count: int, displacement: int | None = None) -> tuple[Sequence, ...]:
    return tuple(s for s in self.paths.get(count, ()) if displacement is None or s.offsets()[-1] == displacement)

def read(path: pathlib.Path, identifier: str, count: int) -> Sequence:
  document = generator.musicxml.load_musicxml(path)
  measures = document['staves'][0]['measures']
  if len(measures) != 1 or len(measures[0]['voices']) != 1 or measures[0]['time_signature'] != f'{count}/4':
    raise ValueError(f'{path}: expected one {count}/4 sequence measure and one voice')
  notes = measures[0]['voices'][0]['notes']
  if len(notes) != count:
    raise ValueError(f'{path}: expected {count} sequence notes')
  degrees = []
  for index, event in enumerate(notes):
    if event.get('grace') or len(event['pitches']) != 1 or generator.model.fraction(event['onset']) != index or generator.model.fraction(event['duration']) != 1:
      raise ValueError(f'{path}: sequence requires consecutive untied quarter notes')
    pitch = event['pitches'][0]
    if generator.model.alteration(pitch['alter']) or 'tie_from' in pitch or 'tie_to' in pitch:
      raise ValueError(f'{path}: sequences must use untied C-major scale degrees')
    degrees.append(7 * pitch['octave'] + generator.model.LETTERS.index(pitch['letter']))
  if degrees[0] != 28:
    raise ValueError(f'{path}: sequence must begin on C4')
  return Sequence(identifier, tuple(b - a for a, b in zip(degrees, degrees[1:])))

def load_library(root: pathlib.Path) -> Library:
  sequences = []
  for path in sorted(root.rglob('*.xml')):
    relative = path.relative_to(root)
    if len(relative.parts) != 2 or relative.parts[0] not in ('2_4', '3_4', '4_4'):
      raise ValueError(f'{path}: sequence sources belong under 2_4, 3_4, or 4_4')
    sequences.append(read(path, relative.as_posix(), int(relative.parts[0][0])))
  if not sequences:
    raise ValueError(f'{root}: sequence library is empty or missing')
  return Library(tuple(sequences))