from __future__ import annotations

import dataclasses
import fractions
import functools
import itertools
import logging
import typing

import generator.fretboard
import generator.model
import generator.pitch
import generator.chord_families
import generator.vocabulary

PREFERRED_MAX_FRET = 12
ELECTRIC_STRINGS = ((3, 2, 1), (4, 3, 2, 1))

class VoicingCost(typing.NamedTuple):
  extended_grips: int = 0
  melody_string: int = 0
  covered_melody: int = 0
  covering_tones: int = 0
  string_shortfall: int = 0
  family_penalty: int = 0
  effort: fractions.Fraction = fractions.Fraction()
  missing_tones: int = 0
  def plus(self, other: VoicingCost) -> VoicingCost:
    return VoicingCost(*(a + b for a, b in zip(self, other)))
  @functools.lru_cache(maxsize=16384)
  def moved(self, distance: int) -> VoicingCost:
    return self._replace(effort=self.effort + distance)

def voicing_preference(role, center=None):
  """Immutable role settings for both selection and scored-retrieval caches."""
  settings = generator.vocabulary.role(role)
  policy = settings.get('voicing', {})
  center = settings.get('center', 53 if role.startswith('acoustic') else 60) if center is None else center
  result = tuple(fractions.Fraction(str(v)) for v in (center, policy.get('register_weight', 1), policy.get('high_position_weight', 1), policy.get('comfortable_max_fret', 9)))
  if result[1] < 0 or result[2] < 0 or result[3] < 1:
    raise ValueError('voicing preference requires nonnegative weights and a positive fret threshold')
  return result

@functools.lru_cache(maxsize=8192)
def voicing_preferences(positions, preference):
  if not positions or not preference:
    return fractions.Fraction(), fractions.Fraction()
  center, register_weight, position_weight, threshold = preference
  register = abs(fractions.Fraction(sum(p.midi for p in positions), len(positions)) - center)
  position = max(fractions.Fraction(), fractions.Fraction(sum(p.fret for p in positions), len(positions)) - threshold)
  return register_weight * register, position_weight * position

def voicing_cost(positions: tuple[generator.fretboard.TabPosition, ...], profile: str = '', allowed: tuple[int, ...] = (), melody: tuple[int, ...] = ()) -> VoicingCost:
  families = tuple(generator.vocabulary.accompaniment('electric')['voicing_families']['rank'].items()) if profile == 'electric' else ()
  preference = voicing_preference('electric_melody') if profile == 'electric' and melody else ()
  return _voicing_cost(positions, profile, allowed, melody, families, preference)

@functools.lru_cache(maxsize=8192)
def _voicing_cost(positions, profile, allowed, melody, families, preference):
  frets = {p.fret for p in positions}
  classes = {p.midi % 12 for p in positions}
  above = sum(p.midi > max(melody) for p in positions) if melody else 0
  span = compactness(positions)
  if profile == 'electric':
    effort = 2 * max(0, len(frets) - 1) + 3 * span + sum(voicing_preferences(positions, preference))
  else:
    effort = 2 * max(0, max(frets, default=0) - PREFERRED_MAX_FRET) + span
  family_penalty = dict(families).get(generator.chord_families.electric(positions), 0) if positions and families else 0
  melody_string = min((p.string_number - 1 for p in positions if melody and p.midi == max(melody)), default=0)
  return VoicingCost(extended_grips=generator.fretboard.span_penalty(positions), melody_string=melody_string, string_shortfall=max(0, 4 - len(positions)) if profile == 'electric' else 0, covered_melody=int(bool(above)), covering_tones=above, family_penalty=family_penalty, effort=effort, missing_tones=len(set(allowed) - classes))

def chord_positions(requirements: generator.fretboard.ChordRequirements, prune: bool = True, profile: str = '', melody: tuple[int, ...] = ()) -> tuple[tuple[generator.fretboard.TabPosition, ...], ...]:
  families = tuple(generator.vocabulary.accompaniment('electric')['voicing_families']['rank'].items()) if profile == 'electric' else ()
  preference = voicing_preference('electric_melody') if profile == 'electric' and melody else ()
  return _chord_positions(requirements, prune, profile, melody, families, preference)

@functools.lru_cache(maxsize=256)
def _chord_positions(requirements, prune, profile, melody, families, preference):
  result = []
  options = [tuple(generator.fretboard.TabPosition(string, fret, open_midi + fret) for fret in range(generator.fretboard.MIN_FRET, generator.fretboard.MAX_FRET + 1) if (open_midi + fret) % 12 in requirements.allowed or open_midi + fret in melody) for string, open_midi in generator.fretboard.STANDARD_TUNING]

  def visit(index, selected, covered):
    if len(set(requirements.required) - covered) > len(options) - index:
      return
    if index == len(options):
      if selected and generator.fretboard.has_fret_anchor(selected) and requirements.covers(covered) and (not requirements.basses or min(selected, key=lambda p: p.midi).midi % 12 in requirements.basses):
        result.append(tuple(selected))
      return
    visit(index + 1, selected, covered)
    for position in options[index]:
      if position.midi % 12 not in covered and generator.fretboard.within_fret_span((*selected, position)):
        visit(index + 1, (*selected, position), covered | {position.midi % 12})

  if profile:
    by_string = {string: options[i] for i, (string, _) in enumerate(generator.fretboard.STANDARD_TUNING)}
    # A lower source melody can be the top of a three-note adaptation on
    # strings 2-4. It need not be buried beneath an added first-string tone.
    for strings in (*ELECTRIC_STRINGS, (4, 3, 2)) if melody else ELECTRIC_STRINGS:
      for selected in itertools.product(*(by_string[s] for s in strings)):
        if not generator.fretboard.within_fret_span(selected) or not generator.fretboard.has_fret_anchor(selected) or not requirements.covers({p.midi % 12 for p in selected}) or not set(melody) <= {p.midi for p in selected}:
          continue
        if requirements.basses and (selected[0].midi % 12 not in requirements.basses or selected[0].midi != min(p.midi for p in selected)):
          continue
        result.append(selected)
  else:
    visit(0, (), set())
  if not prune:
    return tuple(sorted(result))
  # Independent attacks with equal centers have identical movement costs.
  # Held strings require the full candidate set instead of this reduction.
  best = {}
  for positions in result:
    center = position_center(positions)
    rank = (_voicing_cost(positions, profile, requirements.allowed, melody, families, preference), positions)
    if center not in best or rank < best[center][0]:
      best[center] = (rank, positions)
  return tuple(sorted(item[1] for item in best.values()))

@functools.lru_cache(maxsize=16384)
def movement(previous: tuple[generator.fretboard.TabPosition, ...], current: tuple[generator.fretboard.TabPosition, ...]) -> int:
  if not previous or not current:
    return 0
  return sum(abs(a - b) for a, b in zip(position_center(previous), position_center(current)))

@functools.lru_cache(maxsize=4096)
def position_center(positions: tuple[generator.fretboard.TabPosition, ...]) -> tuple[int, int]:
  frets = [p.fret for p in positions] or [0]
  strings = [p.string_number for p in positions]
  return min(frets) + max(frets), min(strings) + max(strings) - 2

def compactness(positions) -> int:
  frets = [p.fret for p in positions]
  return max(frets) - min(frets) if frets else 0

def extend_positions(held, pending):
  if not pending:
    yield tuple(sorted(held.items()))
    return
  identifier, options = pending[0]
  occupied = {p.string_number for p in held.values()}
  for position in options:
    if position.string_number not in occupied and generator.fretboard.within_fret_span((*held.values(), position)):
      yield from extend_positions({**held, identifier: position}, pending[1:])

@dataclasses.dataclass
class Placement:
  cost: VoicingCost
  active: tuple
  previous: Placement | None
  choices: tuple

def nearest_positions(states: dict) -> list[list[tuple | None]]:
  """Exact Manhattan distance transform on the small fret/string position grid."""
  initial = min(((p.cost, i, p) for i, p in enumerate(states.values()) if not p.active), key=lambda item: item[:2], default=None)
  grid = [[initial for _ in range(11)] for _ in range(2 * generator.fretboard.MAX_FRET + 1)]
  for index, previous in enumerate(states.values()):
    if not previous.active:
      continue
    x, y = position_center(tuple(p for _, p in previous.active))
    candidate = (previous.cost, index, previous)
    if grid[x][y] is None or candidate[:2] < grid[x][y][:2]:
      grid[x][y] = candidate

  def relax(x, y, a, b):
    source = grid[a][b]
    if source is not None:
      cost, index, previous = source
      candidate = (cost.moved(1), index, previous)
      if grid[x][y] is None or candidate[:2] < grid[x][y][:2]:
        grid[x][y] = candidate

  for y in range(11):
    for x in range(1, len(grid)):
      relax(x, y, x - 1, y)
    for x in range(len(grid) - 2, -1, -1):
      relax(x, y, x + 1, y)
  for x in range(len(grid)):
    for y in range(1, 11):
      relax(x, y, x, y - 1)
    for y in range(9, -1, -1):
      relax(x, y, x, y + 1)
  return grid

@dataclasses.dataclass(frozen=True)
class GuitarTone:
  pitch: dict
  end: fractions.Fraction
  positions: tuple[generator.fretboard.TabPosition, ...]

@dataclasses.dataclass(frozen=True)
class AttackCandidate:
  active: tuple
  choices: tuple
  shape: tuple[generator.fretboard.TabPosition, ...]
  cost: VoicingCost

@dataclasses.dataclass(frozen=True)
class EventPlacement:
  note: dict
  exact: tuple
  added: tuple[generator.fretboard.TabPosition, ...] = ()

def playable_choices(choices) -> bool:
  return all(generator.fretboard.playable_attack((*((position, generator.model.fraction(tone['duration'])) for tone, position in choice.exact), *((position, generator.model.fraction(choice.note['duration'])) for position in choice.added))) for choice in choices)

def continued_sources(notes):
  result = {p['tie_from'] for n in notes for p in n['pitches'] if 'tie_from' in p}
  result.update(f'{n["_continue_from"]}_s{s}' for n in notes if '_continue_from' in n for s, _ in generator.fretboard.STANDARD_TUNING)
  return result

def has_attack(notes: list[dict]) -> bool:
  return any('_harmony' in n and '_continue_from' not in n or any('tie_from' not in p for p in n['pitches']) for n in notes)

def attack_candidates(notes: list[dict], tones: dict[str, GuitarTone], active: dict, held: dict, preserve_strings: bool, fallback: bool = False):
  pending = []
  chords = [n for n in notes if '_harmony' in n]
  chord = chords[0] if chords else None
  profile = chord.get('_guitar_profile', '') if chord else ''
  if len(chords) > 1:
    raise ValueError('simultaneous generated chords require acoustic harmony-span placement')
  for note in notes:
    for tone in note['pitches']:
      options = tones[tone['id']].positions
      if source := tone.get('tie_from'):
        options = tuple(p for p in options if p == active.get(source))
      pending.append((tone['id'], options))
  if fallback and chord and chord.get('_harmonize'):
    chords, chord, profile = [], None, ''
  requirements = generator.fretboard.chord_requirements(chord['_harmony'], respect_bass=profile != 'electric') if chord else generator.fretboard.ChordRequirements()
  melody = tuple(sorted({p.midi for _, options in pending for p in options})) if profile == 'electric' else ()
  fresh = has_attack(notes)
  candidates = chord_positions(requirements, prune=not (preserve_strings or held or pending or profile), profile=profile, melody=melody) if chord else [()]
  for exact in extend_positions(held, sorted(pending, key=lambda p: len(p[1]))):
    exact_positions = dict(exact)
    for positions in candidates:
      if profile:
        if any(p not in positions for p in exact_positions.values()):
          continue
        assignment = dict(exact)
        choices = []
        compatible = True
        for note in notes:
          added = tuple(p for p in positions if p not in exact_positions.values() and note in chords)
          if any(p.midi % 12 not in requirements.allowed for p in added):
            compatible = False
            break
          if '_continue_from' in note:
            previous = {p for i, p in active.items() if i.startswith(note['_continue_from'] + '_s')}
            if previous != set(added):
              compatible = False
              break
          assignment.update((f'{note["id"]}_s{p.string_number}', p) for p in added)
          choices.append(EventPlacement(note, tuple((p, exact_positions[p['id']]) for p in note['pitches']), added))
        if not compatible or not playable_choices(choices):
          continue
        cost = voicing_cost(positions, profile, requirements.allowed, melody) if fresh else VoicingCost()
        yield AttackCandidate(tuple(sorted(assignment.items())), tuple(choices), positions, cost)
      else:
        shape = tuple(p for _, p in exact) + positions
        if len({p.string_number for p in shape}) != len(shape) or not generator.fretboard.within_fret_span(shape):
          continue
        assignment = tuple(sorted((*exact, *((f'{chord["id"]}_s{p.string_number}', p) for p in positions)))) if chord else exact
        choices = tuple(EventPlacement(n, tuple((p, exact_positions[p['id']]) for p in n['pitches']), positions if n is chord else ()) for n in notes)
        if not playable_choices(choices):
          continue
        cost = voicing_cost(shape, profile, requirements.allowed, melody) if fresh else VoicingCost()
        yield AttackCandidate(assignment, choices, shape, cost)

def guitar_events(document: dict) -> tuple[list, dict[str, GuitarTone]]:
  events, tones = [], {}
  absolute = fractions.Fraction()
  previous = {}
  for measure in document['staves'][0]['measures']:
    if measure.get('section_label'):
      previous.clear()
    changes = sorted(measure.get('attributes', []), key=lambda a: generator.model.fraction(a['onset']))
    state = {k: measure[k] for k in ('transpose', 'staff_details')}
    for change in [state, *changes]:
      state.update({k: change[k] for k in state if k in change})
      generator.fretboard.validate_tuning(state, measure['id'])
    for voice in measure['voices']:
      for order, note in enumerate(voice['notes']):
        if note.get('_rhythm_continuation'):
          source = previous.get(voice['id'])
          if source and '_harmony' in source:
            note['_continue_from'] = source['id']
            for string, _ in generator.fretboard.STANDARD_TUNING:
              tones[f'{source["id"]}_s{string}'].pitch['tie_to'] = f'{note["id"]}_s{string}'
          elif not note['pitches']:
            note.pop('_harmony', None)
            note['kind'] = 'rest'
        previous[voice['id']] = note
        onset = generator.model.fraction(note['onset'])
        state = {k: measure[k] for k in ('transpose', 'staff_details')}
        for change in changes:
          if generator.model.fraction(change['onset']) <= onset:
            state.update({k: change[k] for k in state if k in change})
        for tone in note['pitches']:
          sounding = generator.pitch.sounding_pitch(tone, state)
          options = generator.fretboard.positions_for_pitch(sounding)
          if tab := tone.get('tablature'):
            options = tuple(p for p in options if (p.string_number, p.fret) == (tab['string'], tab['fret']))
          if not options:
            raise ValueError(f'{measure["id"]}, event {note["id"]}: sounding pitch {sounding} has no string/fret in {generator.fretboard.MIN_FRET}-{generator.fretboard.MAX_FRET}')
          tones[tone['id']] = GuitarTone(tone, absolute + onset + generator.model.fraction(tone['duration']), options)
        if '_harmony' in note:
          for string, _ in generator.fretboard.STANDARD_TUNING:
            tones[f'{note["id"]}_s{string}'] = GuitarTone({}, absolute + onset + generator.model.fraction(note['duration']), ())
        if note['pitches'] or '_harmony' in note:
          events.append((absolute + onset, not bool(note.get('grace')), order if note.get('grace') else -1, note))
    absolute += generator.model.fraction(measure['duration'])
  return sorted(events, key=lambda e: e[:3]), tones

def retain_tone(tone: GuitarTone, onset: fractions.Fraction) -> bool:
  return tone.end > onset or tone.end == onset and bool(tone.pitch.get('tie_to'))

def advance_attacks(candidates, states: dict) -> dict:
  nearest = nearest_positions(states)
  result = {}
  for candidate in candidates:
    x, y = position_center(candidate.shape)
    cost, _, previous = nearest[x][y]
    cost = cost.plus(candidate.cost)
    if candidate.active not in result or cost < result[candidate.active].cost:
      result[candidate.active] = Placement(cost, candidate.active, previous, candidate.choices)
  return result

def advance_positions(notes: list[dict], onset: fractions.Fraction, tones: dict[str, GuitarTone], states: dict, preserve_strings: bool, fallback: bool = False) -> dict:
  continued = continued_sources(notes)
  fresh = has_attack(notes)
  held_ids = {i for previous in states.values() for i, _ in previous.active if retain_tone(tones[i], onset)}
  result = {}

  def keep(candidate, previous, cost):
    cost = cost.plus(candidate.cost)
    if candidate.active not in result or cost < result[candidate.active].cost:
      result[candidate.active] = Placement(cost, candidate.active, previous, candidate.choices)

  if not held_ids and not continued:
    result = advance_attacks(attack_candidates(notes, tones, {}, {}, preserve_strings, fallback), states)
  else:
    for previous in states.values():
      active = dict(previous.active)
      held = {i: p for i, p in previous.active if i not in continued and i in held_ids}
      for candidate in attack_candidates(notes, tones, active, held, preserve_strings, fallback):
        distance = movement(tuple(p for _, p in previous.active), candidate.shape) if fresh else 0
        keep(candidate, previous, previous.cost.moved(distance))
  if not result:
    if not fallback and any(n.get('_harmonize') for n in notes):
      return advance_positions(notes, onset, tones, states, preserve_strings, True)
    raise ValueError(f'event {notes[0]["id"]}, onset {onset}: no guitar placement; single notes need frets {generator.fretboard.MIN_FRET}-{generator.fretboard.MAX_SINGLE_FRET}; chords need frets {generator.fretboard.MIN_FRET}-{generator.fretboard.MAX_FRET}, an anchor <= {generator.fretboard.MAX_ANCHOR_FRET}, and at most {generator.fretboard.MAX_FRET_SPAN} consecutive fret positions; held/tied strings must remain fixed')
  return result

def place_guitar(document: dict) -> None:
  """Choose electric voicings or exact/held melody placements."""
  events, tones = guitar_events(document)
  states = {(): Placement(VoicingCost(), (), None, ())}
  groups = [(key[0], [item[3] for item in values]) for key, values in itertools.groupby(events, key=lambda e: e[:3])]
  for index, (onset, notes) in enumerate(groups):
    next_onset = groups[index + 1][0] if index + 1 < len(groups) else None
    preserve_strings = next_onset is not None and any('_harmony' in n and onset + generator.model.fraction(n['duration']) > next_onset for n in notes)
    states = advance_positions(notes, onset, tones, states, preserve_strings)
  selected = min(states.values(), key=lambda s: (s.cost, s.active))
  skipped = []
  while selected.previous is not None:
    for choice in selected.choices:
      note = choice.note
      for tone, position in choice.exact:
        tone['tablature'] = {'string': position.string_number, 'fret': position.fret}
      if definition := note.pop('_harmony', None):
        for pitch in generator.fretboard.chord_pitches(definition, choice.added):
          note['pitches'].append({**pitch, 'id': f'{note["id"]}_s{pitch["tablature"]["string"]}', 'duration': note['duration'], 'written': note['written']})
        if note.get('_harmonize') and not choice.added and not note.get('_rhythm_continuation'):
          skipped.append(note['id'])
    selected = selected.previous
  for measure in document['staves'][0]['measures']:
    for voice in measure['voices']:
      for note in voice['notes']:
        for field in ('_guitar_profile', '_guitar_lane', '_harmonize', '_continue_from'):
          note.pop(field, None)
  if skipped:
    logging.getLogger(__name__).warning('%s: %s melody harmonizations kept as single notes; first event %s', document['metadata']['title'], len(skipped), skipped[-1])
  document['staves'][0]['tablature'] = True
  generator.fretboard.validate_document(document)