from __future__ import annotations

import dataclasses
import functools
import itertools
import typing

import generator.acoustic
import generator.fretboard
import generator.harmony

MAX_ANCHOR_FRET = 12

class Voicing(typing.NamedTuple):
  bass: generator.fretboard.TabPosition
  string4: generator.fretboard.TabPosition
  string3: generator.fretboard.TabPosition
  string2: generator.fretboard.TabPosition

def playable(voicing: Voicing) -> bool:
  return voicing.bass.string_number in (5, 6) and tuple(p.string_number for p in voicing[1:]) == (4, 3, 2) and all(generator.fretboard.MIN_FRET <= p.fret <= generator.fretboard.MAX_FRET for p in voicing) and min(p.fret for p in voicing) <= MAX_ANCHOR_FRET and generator.fretboard.within_fret_span(voicing) and voicing.bass.midi == min(p.midi for p in voicing)

@functools.lru_cache(maxsize=256)
def candidates(requirements: generator.fretboard.ChordRequirements, complete_triad: bool = False) -> tuple[Voicing, ...]:
  options = {s: tuple(generator.fretboard.TabPosition(s, f, midi + f) for f in range(generator.fretboard.MIN_FRET, generator.fretboard.MAX_FRET + 1) if (midi + f) % 12 in requirements.allowed) for s, midi in generator.fretboard.STANDARD_TUNING if s != 1}
  result = []
  for bass in (*options[5], *options[6]):
    if requirements.basses and bass.midi % 12 not in requirements.basses:
      continue
    for upper in itertools.product(*(options[s] for s in (4, 3, 2))):
      voicing = Voicing(bass, *upper)
      if playable(voicing) and requirements.covers({p.midi % 12 for p in voicing}):
        result.append(voicing)
  if complete_triad:
    complete = [v for v in result if set(requirements.allowed) <= {p.midi % 12 for p in v}]
    result = complete or result
  return tuple(sorted(result))

@functools.lru_cache(maxsize=16384)
def movement(previous: Voicing, current: Voicing) -> int:
  return sum(abs(a.fret - b.fret) for a, b in zip(previous, current)) + abs(previous.bass.string_number - current.bass.string_number)

class Cost(typing.NamedTuple):
  travel: int = 0
  missing_tones: int = 0

@dataclasses.dataclass
class Placement:
  cost: Cost
  voicing: Voicing | None = None
  previous: Placement | None = None

def advance(options: tuple[Voicing, ...], requirements: generator.fretboard.ChordRequirements, states: list[Placement]) -> list[Placement]:
  result = []
  for voicing in options:
    missing = len(set(requirements.allowed) - {p.midi % 12 for p in voicing})
    cost, index = min((Cost(p.cost.travel + (movement(p.voicing, voicing) if p.voicing else 0), p.cost.missing_tones + missing), i) for i, p in enumerate(states))
    result.append(Placement(cost, voicing, states[index]))
  return result

def select(spans: list[generator.acoustic.HarmonySpan]) -> list[Voicing]:
  states = [Placement(Cost())]
  for span in spans:
    try:
      requirements = generator.fretboard.chord_requirements(span.harmony)
      options = candidates(requirements, generator.harmony.is_triad(span.harmony))
      if not options:
        raise ValueError(f'no acoustic grip on strings 2-4 plus bass 5 or 6 within frets {generator.fretboard.MIN_FRET}-{generator.fretboard.MAX_FRET}, three-fret span, and at least one fret <= {MAX_ANCHOR_FRET}')
    except ValueError as error:
      raise ValueError(f'{span.location}, {generator.harmony.symbol(span.harmony)}: {error}') from error
    states = advance(options, requirements, states)
  selected = min(states, key=lambda p: (p.cost, p.voicing))
  result = []
  while selected.previous is not None:
    result.append(selected.voicing)
    selected = selected.previous
  return list(reversed(result))

def place(document: dict) -> None:
  """Plan every iteration before materializing its four string-voice journeys."""
  for measure in document['staves'][0]['measures']:
    for state in (measure, *measure.get('attributes', [])):
      if 'staff_details' in state:
        generator.fretboard.validate_tuning(state, measure['id'])
  choices = []
  for _, group in itertools.groupby(generator.acoustic.harmony_spans(document), key=lambda span: span.iteration):
    spans = list(group)
    choices.extend(zip(spans, select(spans)))
  for span, voicing in choices:
    for note in span.notes:
      positions = (voicing.bass,) if note['_guitar_lane'] == 'bass' else voicing[1:]
      for pitch in generator.fretboard.chord_pitches(span.harmony, positions):
        note['pitches'].append({**pitch, 'id': f'{note["id"]}_s{pitch["tablature"]["string"]}', 'duration': note['duration'], 'written': note['written']})
      for field in ('_harmony', '_guitar_profile', '_guitar_lane'):
        note.pop(field, None)
  document['staves'][0]['tablature'] = True