from __future__ import annotations

import dataclasses
import pathlib
import types

import generator.solo.policy

import generator.model
import generator.musicxml

@dataclasses.dataclass(frozen=True)
class Sequence:
  id: str
  steps: tuple[int, ...]
  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

@dataclasses.dataclass(frozen=True)
class Path:
  steps: tuple[int, ...]
  sources: tuple[str, ...]
  def offsets(self):
    result = [0]
    for step in self.steps:
      result.append(result[-1] + step)
    return tuple(result)

def family(steps):
  if not steps or all(s == 0 for s in steps):
    return 'repeated'
  turns = sum(a * b < 0 for a, b in zip(steps, steps[1:]))
  if not turns:
    return 'rising' if sum(steps) > 0 else 'falling'
  if max(map(abs, steps)) >= 4:
    return 'leap_return'
  if sum(steps) == 0 and max(map(abs, steps)) == 1:
    return 'neighbor'
  return ('arch' if steps[0] > 0 else 'valley') if turns == 1 else 'wave'

def diverse(items, rank, category, limit):
  """Round-robin eligible musical families before a destructive cutoff."""
  groups = {}
  for item in sorted(items, key=rank):
    groups.setdefault(category(item), []).append(item)
  result, index = [], 0
  while len(result) < limit:
    layer = [values[index] for values in groups.values() if len(values) > index]
    if not layer:
      break
    result.extend(sorted(layer, key=rank)[:limit - len(result)])
    index += 1
  return tuple(result)

@dataclasses.dataclass(frozen=True, init=False)
class Library:
  paths: types.MappingProxyType
  chains: dict
  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()}))
    object.__setattr__(self, 'chains', {})
  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 chained(self, count):
    return self.retrieve(count)
  def retrieve(self, count, displacement=0, direction=0):
    key = (count, displacement, direction)
    if key in self.chains:
      return self.chains[key]
    if not 1 <= count <= 16:
      return ()
    states, remaining = (Path((), ()),), count - 1
    while remaining:
      size = min(3, remaining)
      options = self.connecting(size + 1)
      if not options:
        size = next((n for n in range(size - 1, 0, -1) if self.connecting(n + 1)), 0)
        if not size:
          return ()
        options = self.connecting(size + 1)
      desired = displacement * (count - 1 - remaining + size) / max(1, count - 1)
      def rank(path):
        offsets = path.offsets()
        motion = offsets[-1]
        return abs(motion - desired) + (2 if direction and motion * direction < 0 else 0), max(offsets) - min(offsets), max(map(abs, path.steps), default=0), path.sources
      # Retrieve from the entire authored vocabulary, conditioned on each prefix.
      following = []
      for prior in states:
        if count > 4:
          local = self.retrieve(size + 1, round(desired - prior.offsets()[-1]), direction)
          paths = (Path(prior.steps + s.steps, prior.sources + s.sources) for s in local)
        else:
          paths = (Path(prior.steps + s.steps, prior.sources + (s.id,)) for s in options)
        following.extend(diverse(paths, rank, lambda p: family(p.steps), generator.solo.policy.DEFAULT.retrieval))
      states = diverse(following, rank, lambda p: (family(p.steps), (p.offsets()[-1] > 0) - (p.offsets()[-1] < 0)), generator.solo.policy.DEFAULT.retrieval)
      remaining -= size
    self.chains[key] = states
    return states

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))