"""Bounded, source-only measure rhythms; acoustic layers are selected jointly."""
import dataclasses
import fractions
import itertools

import generator.accompaniment
import generator.model
import generator.solo.analysis
import generator.solo.policy
import generator.vocabulary

@dataclasses.dataclass(frozen=True, order=True)
class Attack:
  lane: str
  start: fractions.Fraction
  end: fractions.Fraction

@dataclasses.dataclass(frozen=True)
class Rhythm:
  events: tuple
  families: tuple
  contributions: tuple
  exceptions: tuple = ()
  @property
  def cost(self):
    return sum(value for _, value in self.contributions)
  @property
  def signature(self):
    return tuple((e.lane, str(e.start), str(e.end - e.start)) for e in self.events)

def melody_contexts(document):
  melody = generator.solo.analysis.attacks(document)
  result, cursor = [], fractions.Fraction()
  for m in document['staves'][0]['measures']:
    length = generator.model.fraction(m['duration'])
    result.append(tuple((n.start - cursor, n.end - cursor) for n in melody if n.start < cursor + length and n.end > cursor))
    cursor += length
  return tuple(result)

def spans(measure, prior):
  length = generator.model.fraction(measure['duration'])
  boundaries = sorted({fractions.Fraction(), length, *(generator.model.fraction(h['onset']) for h in measure['harmonies'])})
  result = []
  for a, b in zip(boundaries, boundaries[1:]):
    chord = generator.accompaniment.active_chord(measure, a, prior)
    if result and generator.accompaniment.same_chord(result[-1][2], chord):
      result[-1] = (result[-1][0], b, chord)
    else:
      result.append((a, b, chord))
  return tuple(result)

def lane_score(events, lane, melody, harmony, length):
  """Named source-context preferences, independent of the deterministic seed."""
  attacks = [a for a, b in melody if a >= 0]
  occupied = sum((max(0, min(length, b) - max(0, a)) for a, b in melody), fractions.Fraction())
  space = max(0, length - occupied)
  busy = len(attacks) >= length
  if lane == 'bass':
    target = min(2, sum(c is not None and b - a >= 1 for a, b, c in harmony))
  elif lane == 'upper':
    target = 2.5 if space >= length / 2 else 1.5 if busy else 2
  else:
    target = .5 if busy else 1.5 if space >= length / 2 else 1
  # A half-attack tolerance makes neighboring densities genuine alternatives.
  # The objective does not always fill up to a layer's permitted maximum.
  activity = .35 * max(0, abs(len(events) - target) - .5)
  support, articulation = 0., 0.
  for e in events:
    # Supporting an entrance is different from doubling every moving note.
    entrance = any(a == e.start and (i == 0 and a >= 0 or i > 0 and a - melody[i - 1][1] >= fractions.Fraction(1, 2)) for i, (a, _) in enumerate(melody))
    moving = any(a <= e.start < b and b - a <= 1 for a, b in melody)
    gap = not any(a <= e.start < b for a, b in melody)
    if lane != 'bass':
      support += -.2 if entrance else .25 if moving and e.start in attacks else -.15 if gap else 0
      articulation += .12 if busy and e.end - e.start > 1 else .08 if not busy and e.end - e.start < 1 else 0
    elif e.start and e.start.denominator != 1:
      support += .15
  divisor = max(1, len(events))
  return (('activity', activity), ('melody_support', support / divisor), ('duration_contrast', articulation / divisor))

def representatives(candidates, limit, seed):
  """Keep counts, placements and duration families before the final cutoff."""
  unique = {}
  for candidate in candidates:
    key = candidate.signature
    rank = (candidate.cost, generator.solo.policy.stable_rank(seed, key))
    if key not in unique or rank < unique[key][0]:
      unique[key] = rank, candidate
  ordered = [v for _, v in sorted(unique.values(), key=lambda item: item[0])]
  groups = {}
  shapes = set()
  for c in ordered:
    counts = tuple((lane, sum(e.lane == lane for e in c.events)) for lane, _ in c.families)
    shape = (counts, c.families, any(e.start.denominator != 1 for e in c.events))
    # Round-robin attack counts first. Otherwise many cosmetic variants of
    # the cheapest density can exclude every one- or three-attack idea.
    if shape not in shapes:
      groups.setdefault(counts, []).append(c)
      shapes.add(shape)
  result = []
  while groups and len(result) < limit:
    for key in list(groups):
      result.append(groups[key].pop(0))
      if not groups[key]:
        del groups[key]
      if len(result) == limit:
        break
  return tuple(result)

def lane_candidates(measure, prior, lane, melody, policy, seed):
  length = generator.model.fraction(measure['duration'])
  harmony = spans(measure, prior)
  active = [(a, b, c) for a, b, c in harmony if c]
  if not active:
    return Rhythm((), ((lane, 'space'),), (('activity', 0.),)),
  low, high = policy['layers'][lane]['attacks']
  # Genuinely different harmony gets a gesture, not its own density budget.
  required = [(a, b) for a, b, c in active if b - a >= (2 if lane == 'bass' else 1)]
  if lane == 'bass' and active[0][0] == 0 and not any(a == 0 for a, _ in required):
    required.insert(0, active[0][:2])
  if lane == 'electric' and len(active) == 1 and generator.accompaniment.same_chord(active[0][2], prior):
    required = []
  # With very fast harmonic rhythm, establish the edges and longest spans.
  # Short intervening chords may be silent; no incompatible sustain is kept.
  if len(required) > high + 1:
    edges = [required[0], required[-1]]
    required = sorted(edges + sorted(required[1:-1], key=lambda ab: (-(ab[1] - ab[0]), ab[0]))[:high - 1])
  extra = int(len(required) > high)
  grid = fractions.Fraction(policy['grid'])
  positions = {i * grid for i in range(int(length / grid) + 1) if i * grid < length}
  positions.update(a for a, _, _ in active)
  positions = sorted(t for t in positions if any(a <= t < b for a, b, _ in active))
  result = []
  for count in range(low, high + extra + 1):
    for onsets in itertools.combinations(positions, count):
      if lane == 'bass' and active[0][0] == 0 and (not onsets or onsets[0] != 0):
        continue
      if any(not any(a <= t < b for t in onsets) for a, b in required):
        continue
      if any(b - a < fractions.Fraction(policy['attack_spacing']) for a, b in zip(onsets, onsets[1:])):
        continue
      for family in policy['layers'][lane]['families']:
        durations = policy['families'][family]
        events = []
        for i, a in enumerate(onsets):
          end = next(b for start, b, c in active if start <= a < b)
          following = onsets[i + 1] if i + 1 < len(onsets) else length
          duration = fractions.Fraction(durations[i % len(durations)])
          events.append(Attack(lane, a, min(a + duration, following, end)))
        if policy.get('allowed_durations') and any(e.end - e.start not in tuple(map(fractions.Fraction, policy['allowed_durations'])) for e in events):
          continue
        exceptions = ((lane, 'extra attack for harmonic changes'),) if count > high else ()
        contributions = (*lane_score(events, lane, melody, harmony, length), ('extra_attacks', .5 * bool(exceptions)))
        result.append(Rhythm(tuple(events), ((lane, family if count else 'space'),), contributions, exceptions))
  if not result:
    # A short partial bar can have fewer gestures, not invented subdivisions.
    if length < 1 and not policy.get('allowed_durations'):
      a, b, _ = active[0]
      return Rhythm((Attack(lane, a, b),), ((lane, 'partial'),), (('activity', 0.),)),
    raise ValueError(f'measure {measure["number"]}: no {lane} rhythm within the bounded attack/spacing policy')
  return representatives(result, policy['lane_candidates'], seed + '/' + lane)

def crowding(events):
  times = sorted({e.start for e in events})
  cost = .22 * sum(y - x < 1 for x, y in zip(times, times[1:]))
  # Three successive half-beat changes are a busy figure even though each
  # layer is sparse. Coincident attacks count as one hand action.
  return cost + .5 * sum(z - x <= 1 for x, y, z in zip(times, times[1:], times[2:]))

def candidates(measure, prior, part, melody, seed):
  policy = generator.vocabulary.accompaniment(part)
  if part == 'electric':
    return lane_candidates(measure, prior, 'electric', melody, policy, seed)
  bass = lane_candidates(measure, prior, 'bass', melody, policy, seed)
  upper = lane_candidates(measure, prior, 'upper', melody, policy, seed)
  result = []
  for a, b in itertools.product(bass, upper):
    costs = tuple((lane + '_' + k, v) for lane, c in (('bass', a), ('upper', b)) for k, v in c.contributions)
    result.append(Rhythm(tuple(sorted((*a.events, *b.events))), (*a.families, *b.families), (*costs, ('combined_crowding', crowding((*a.events, *b.events)))), (*a.exceptions, *b.exceptions)))
  return representatives(result, policy['measure_candidates'], seed)

def transition(previous, current):
  if previous is None:
    return 0.
  # Recurrence is allowed; penalize repeating an entire bar more than sharing
  # a duration family. This stays separate from the source-context objective.
  same_counts = all(sum(e.lane == lane for e in previous.events) == sum(e.lane == lane for e in current.events) for lane, _ in current.families)
  return .45 * (previous.signature == current.signature) + .12 * (previous.families == current.families) + .08 * same_counts

def select(layers, seed, left=None, right=None):
  """One state per concrete bar rhythm; compare complete equal-span paths."""
  states = [(0., (), (), left)]
  for i, layer in enumerate(layers):
    following = []
    for c in layer:
      options = [(cost + c.cost + transition(previous, c), (*ties, generator.solo.policy.stable_rank(seed, (i, c.signature))), (*path, c), c) for cost, ties, path, previous in states]
      following.append(min(options, key=lambda item: item[:2]))
    states = following
  return min(states, key=lambda item: (item[0] + (transition(item[3], right) if right else 0), item[1]))[2]