"""Complete, source-only comping figures. Time is in rational quarter notes."""
import dataclasses
import fractions
import itertools

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

def meter_context(measure):
  """Nominal meter and explicit additive grouping, independent of bar length.

  In particular, a three-quarter-note pickup or a 6/8 bar is not 3/4.
  """
  groups = measure['meter']['groups']
  if not groups or any(g['unit'] != 4 for g in groups):
    raise ValueError(f"source measure {measure['number']}: no comping vocabulary for {measure['time_signature']}")
  beats = tuple(b for g in groups for b in g['beats'])
  total = sum(beats)
  if total not in (3, 4, 5) or (len(beats) > 1 and beats not in ((3, 2), (2, 3))):
    raise ValueError(f"source measure {measure['number']}: unsupported comping grouping {measure['time_signature']}")
  return f'{total}/4', beats if len(beats) > 1 else ()

def grouping_for(run, melody):
  """Choose one source-only quintuple feel, shared by independent roles.

  Written grouping wins. Otherwise compare actual chord changes, entrances
  after an eighth-note breath, then all melodic attacks at the group boundary.
  Counts are evidence, not inferred accents or performance dynamics.
  """
  meter, explicit = meter_context(run[0][0][0])
  if meter != '5/4':
    return (), {'reason': 'simple_meter'}
  if explicit:
    return explicit, {'reason': 'written'}
  evidence = {}
  for boundary in (3, 2):
    changes, entrances, attacks = 0, 0, 0
    for group in run:
      for m, prior, i in group:
        changes += sum(a == boundary and c is not None for a, _, c in spans(m, prior))
        context = melody[i]
        for a, _ in context:
          if a != boundary:
            continue
          attacks += 1
          previous = max((b for x, b in context if x < a), default=fractions.Fraction())
          entrances += a - previous >= fractions.Fraction(1, 2)
    evidence[boundary] = (changes, entrances, attacks)
  boundary = 2 if evidence[2] > evidence[3] else 3
  return (boundary, 5 - boundary), {'reason': 'source' if evidence[2] != evidence[3] else 'default_3+2',
    'boundary_evidence': {str(k): dict(zip(('chord_changes', 'melody_entrances', 'melody_attacks'), v)) for k, v in evidence.items()}}

def patterns_for(policy, measure, lengths, grouping=()):
  meter, explicit = meter_context(measure)
  grouping = explicit or grouping or ((3, 2) if meter == '5/4' else ())
  # Untagged in-memory figures are allowed for independently authored unit
  # tests; the resource loader requires metadata for every production entry.
  return [p for p in policy['patterns'] if tuple(map(fractions.Fraction, p['lengths'])) == tuple(lengths)
    and (p.get('meter') == 'partial' or (p.get('meter', meter) == meter and tuple(p.get('grouping', grouping)) == grouping))]

def relevant_policy(policy, measures):
  """An unused meter's vocabulary must not influence this role's decisions."""
  meters = {meter_context(m)[0] for m in measures}
  lengths = {generator.model.fraction(m['duration']) for m in measures}
  return {**policy, 'patterns': [p for p in policy['patterns'] if p.get('meter') in meters
    or (p.get('meter') == 'partial' and fractions.Fraction(p['lengths'][0]) in lengths)]}

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

@dataclasses.dataclass(frozen=True)
class Rhythm:
  events: tuple
  families: tuple
  contributions: tuple
  exceptions: tuple = ()
  pattern_id: str = ''
  @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), str(e.harmonic_time)) 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 flatten(inputs):
  """Harmony/timing adapter only; never concatenate separately chosen rhythms."""
  first = inputs[0][0]
  result = {**first, 'harmonies': []}
  cursor = fractions.Fraction()
  for m, prior, _ in inputs:
    chord = generator.accompaniment.active_chord(m, fractions.Fraction(), prior)
    initial = chord or {'kind': 'none'}
    result['harmonies'].append({**initial, 'onset': generator.model.rational(cursor)})
    result['harmonies'].extend({**h, 'onset': generator.model.rational(cursor + generator.model.fraction(h['onset']))} for h in m['harmonies'] if generator.model.fraction(h['onset']))
    cursor += generator.model.fraction(m['duration'])
  result['duration'] = generator.model.rational(cursor)
  return result

def validate_events(events, part, lengths=None):
  ordered = sorted(events, key=lambda e: e.start)
  for e in ordered:
    duration = e.end - e.start
    if duration <= 0:
      raise ValueError('comping event must have positive sounding duration')
    if part == 'acoustic':
      if e.lane == 'combined':
        if duration < 1 and duration != fractions.Fraction(1, 2):
          raise ValueError('combined acoustic attacks require quarter/eighth notes or sustains')
      elif e.lane not in ('bass', 'upper') or duration not in (fractions.Fraction(1, 2), 1):
        raise ValueError('ordinary acoustic attacks require separate quarter/eighth notes')
    elif e.lane != 'electric':
      raise ValueError('electric comping requires chord events')
  for i, a in enumerate(ordered):
    for b in ordered[i + 1:]:
      if a.end > b.start and (part != 'acoustic' or a.lane == b.lane or 'combined' in (a.lane, b.lane)):
        raise ValueError('comping attacks cannot overlap on the same strings')

def fit(pattern, measure, prior, part, *, anticipate=False, limit=None, initial_bass=True):
  """Bind a complete figure to written harmony, with bounded early arrivals."""
  length = generator.model.fraction(measure['duration'])
  limit = length if limit is None else min(length, limit)
  harmony = spans(measure, prior)
  events, exceptions, available = [], [], []
  for row in pattern['events']:
    start = fractions.Fraction(row['onset'])
    end = min(start + fractions.Fraction(row['duration']), limit)
    if start >= limit:
      continue
    active = next(((a, b, c) for a, b, c in harmony if a <= start < b), None)
    if active is None or active[2] is None:
      exceptions.append(('space', 'no written harmony at ' + str(start)))
      continue
    available.append((dict(row), start, end, active))
  lengths = tuple(map(fractions.Fraction, pattern.get('lengths', (length,))))
  for row, start, end, active in available:
    _, boundary, chord = active
    target = None
    if anticipate and row.get('anticipate') and row['layer'] in ('upper', 'electric'):
      future = next(((a, b, c) for a, b, c in harmony if a == start + fractions.Fraction(1, 2) and c), None)
      if future and future[0] < end and not generator.accompaniment.same_chord(chord, future[2]):
        target, boundary = future[0], future[1]
    end = min(end, boundary)
    events.append(Attack(row['layer'], start, end, target))
  if part == 'acoustic':
    for a, b, chord in harmony:
      if chord is None or (a == 0 and not initial_bass):
        continue
      local = [e for e in events if a <= (e.harmonic_time if e.harmonic_time is not None else e.start) < b]
      if not local:
        continue
      first = min(local, key=lambda e: e.start)
      if any(e.start == first.start and e.lane in ('bass', 'combined') for e in local):
        continue
      # Preserve the comping rhythm and upper voice. Add bass at the same
      # attack, allowing different releases when the next bass enters early.
      stop = min([first.end, first.start + 1, *(e.start for e in events if e.start > first.start and e.lane in ('bass', 'combined'))])
      duration = next((d for d in (fractions.Fraction(1), fractions.Fraction(1, 2)) if d <= stop - first.start), None)
      if duration is None:
        return None
      if first.end - first.start == duration:
        events[events.index(first)] = dataclasses.replace(first, lane='combined')
      else:
        events.append(Attack('bass', first.start, first.start + duration, first.harmonic_time))
      exceptions.append(('layer_assignment', f'bass included at first attack for harmony {a}: {first.start}'))
    events.sort(key=lambda e: (e.start, e.lane))
  try:
    validate_events(events, part, lengths)
  except ValueError:
    return None
  return tuple(events), tuple(exceptions)

def evaluate(events, part, melody, harmony, lengths):
  activity, support, coverage, articulation = 0., 0., 0., 0.
  cursor = fractions.Fraction()
  for length in lengths:
    local = [e for e in events if cursor <= e.start < cursor + length]
    attacks = [a for a, b in melody if cursor <= a < cursor + length]
    busy = len(attacks) >= length
    sounding = any(e.start < cursor + length and e.end > cursor for e in events)
    target = (1 if busy else 1.5) if part == 'electric' else (2 if busy else 2.5)
    # Sustained occupancy and planned space are not missing attacks.
    activity += .22 * max(0, abs(len({e.start for e in local}) - target) - .75) if not (sounding and not local) else 0
    # A bass paired with the upper chord is one rhythmic attack, regardless
    # of whether unequal releases require two internal lane events.
    attacks_by_time = {e.start: max((v for v in local if v.start == e.start), key=lambda v: v.end) for e in local}
    for e in attacks_by_time.values():
      gap = not any(a <= e.start < b for a, b in melody)
      entrance = any(a == e.start and (i == 0 or a - melody[i - 1][1] >= fractions.Fraction(1, 2)) for i, (a, _) in enumerate(melody))
      support += -.2 if entrance else -.12 if gap else .18 if busy and e.start in attacks else 0
      articulation += .08 if busy and e.end - e.start > 2 else 0
    cursor += length
  for a, b, chord in harmony:
    if chord and b - a >= 1 and not any(e.start < b and e.end > a for e in events):
      coverage += .18  # A preference, never a chord-attack quota.
  divisor = max(1, len(lengths))
  return tuple((k, v / divisor) for k, v in (('activity', activity), ('melody_support', support), ('harmonic_support', coverage), ('articulation', articulation)))

def candidates(measure, prior, part, melody, seed, *, lengths=None, limit=None, grouping=()):
  lengths = lengths or (generator.model.fraction(measure['duration']),)
  policy = generator.vocabulary.accompaniment(part)
  result = []
  for pattern in patterns_for(policy, measure, lengths, grouping):
    for anticipate in (False, True):
      fitted = fit(pattern, measure, prior, part, anticipate=anticipate, limit=limit)
      if fitted is None:
        continue
      events, exceptions = fitted
      if anticipate and not any(e.harmonic_time is not None for e in events):
        continue
      contributions = evaluate(events, part, melody, spans(measure, prior), lengths)
      result.append(Rhythm(events, (('pattern', pattern['family']),), contributions, exceptions, pattern['id'] + ('.anticipated' if anticipate else '')))
  if not result:
    raise ValueError(f'measure {measure["number"]}: no {part} pattern for lengths {lengths} and written harmonic/release constraints')
  return representatives(result, policy['pattern_candidates'], seed)

def representatives(candidates, limit, seed):
  unique = {}
  for c in candidates:
    rank = (c.cost, generator.solo.policy.stable_rank(seed, (c.pattern_id, c.signature)))
    if c.signature not in unique or rank < unique[c.signature][0]:
      unique[c.signature] = rank, c
  ordered = [c for _, c in sorted(unique.values(), key=lambda item: item[0])]
  retained, families = [], set()
  for c in ordered:
    if c.families not in families:
      retained.append(c)
      families.add(c.families)
  retained.extend(c for c in ordered if c not in retained)
  return tuple(retained[:limit])

def project(pattern, offset, lengths):
  """Take a form-position slice, rather than restarting a pattern at an entry."""
  end = offset + sum(lengths)
  rows = []
  for row in pattern['events']:
    a = fractions.Fraction(row['onset'])
    b = a + fractions.Fraction(row['duration'])
    if a < end and b > offset:
      rows.append({**row, 'onset': str(max(a, offset) - offset), 'duration': str(min(b, end) - max(a, offset))})
  return {**pattern, 'lengths': list(map(str, lengths)), 'events': rows}

def realize(pattern, inputs, part, melody, *, anticipate=False, limit=None, boundary=False):
  flat = flatten(inputs)
  lengths = tuple(generator.model.fraction(m['duration']) for m, _, _ in inputs)
  initial_bass = not generator.accompaniment.same_chord(inputs[0][1], generator.accompaniment.active_chord(inputs[0][0], fractions.Fraction(), inputs[0][1]))
  fitted = fit(pattern, flat, None, part, anticipate=anticipate, limit=limit, initial_bass=initial_bass)
  if fitted is None and boundary:
    # Keep the selected onset/layers at a protected handoff, using the longest
    # permitted short articulation that fits rather than another groove.
    clipped = project(pattern, fractions.Fraction(), (min(sum(lengths), limit) if limit is not None else sum(lengths),))
    rows = []
    for row in clipped['events']:
      a, duration = fractions.Fraction(row['onset']), fractions.Fraction(row['duration'])
      _, b, _ = next(s for s in spans(flat, None) if s[0] <= a < s[1])
      duration = min(duration, b - a)
      if part == 'acoustic' and (row['layer'] != 'combined' or duration <= 1):
        duration = next((d for d in (fractions.Fraction(1), fractions.Fraction(1, 2)) if d <= duration), fractions.Fraction())
      if duration:
        rows.append({**row, 'duration': str(duration)})
    fitted = fit({**pattern, 'events': rows}, flat, None, part, anticipate=False, limit=limit, initial_bass=initial_bass)
    if fitted is not None:
      fitted = fitted[0], (*fitted[1], ('release', 'shortened articulation at source-defined boundary'))
  if fitted is None:
    return None
  events, exceptions = fitted
  context, cursor = [], fractions.Fraction()
  for m, _, i in inputs:
    context.extend((a + cursor, b + cursor) for a, b in melody[i])
    cursor += generator.model.fraction(m['duration'])
  return Rhythm(events, (('pattern', pattern['family']),), evaluate(events, part, context, spans(flat, None), lengths), exceptions,
    pattern['id'] + ('.anticipated' if any(e.harmonic_time is not None for e in events) else ''))

def identity(pattern):
  """Layer swaps and anticipation permissions alone are not different rhythms."""
  return tuple((fractions.Fraction(r['onset']), fractions.Fraction(r['duration'])) for r in pattern['events'])

def choose_cycles(groups, part, melody, seed):
  """Select a small recurring vocabulary against every occurrence in the form.

  No local shortlist precedes whole-form evaluation. A pattern/mode is retained
  as one decision even where a chord boundary changes its legal release.
  """
  policy = generator.vocabulary.accompaniment(part)
  settings = policy['groove']
  selected, locations, reports = [], [], []
  runs = []
  for group in groups:
    m = group[0][0]
    key = (meter_context(m), generator.model.fraction(m['duration']))
    if not runs or runs[-1][0] != key:
      runs.append((key, []))
    runs[-1][1].append(group)
  for ((meter, _), length), run in runs:
    grouping, grouping_evidence = grouping_for(run, melody)
    patterns = patterns_for(policy, run[0][0][0], (length, length), grouping)
    if not patterns:
      patterns = patterns_for(policy, run[0][0][0], (length,), grouping)
    requested = settings['cycle_patterns']
    identities = {p['id']: identity(p) for p in patterns}
    available = len(set(identities.values()))
    cache = {}
    for i, group in enumerate(run):
      lengths = tuple(generator.model.fraction(m['duration']) for m, _, _ in group)
      for p in patterns:
        for anticipate in (False, True):
          cache[i, p['id'], anticipate] = realize(project(p, fractions.Fraction(), lengths), group, part, melody, anticipate=anticipate)
    winner, examined = None, 0
    for width in range(min(requested, len(run), available), 0, -1):
      layers = []
      for slot in range(width):
        unique = {}
        for p in patterns:
          for anticipate in (False, True):
            values = [cache[i, p['id'], anticipate] for i in range(slot, len(run), width)]
            if any(v is None for v in values):
              continue
            cost = sum(v.cost * len(run[i]) for i, v in zip(range(slot, len(run), width), values))
            rank = (cost, generator.solo.policy.stable_rank(seed, (meter, slot, p['id'], anticipate)))
            option = (rank, p, anticipate)
            key = identities[p['id']]
            if key not in unique or rank < unique[key][0]:
              unique[key] = option
        ordered = sorted(unique.values(), key=lambda o: o[0])
        diverse, families = [], set()
        for option in ordered:
          if option[1]['family'] not in families:
            diverse.append(option)
            families.add(option[1]['family'])
        diverse.extend(o for o in ordered if o not in diverse)
        layers.append(diverse[:settings['slot_candidates']])
      for cycle in itertools.product(*layers):
        if len({identities[o[1]['id']] for o in cycle}) != width:
          continue
        examined += 1
        rank = (sum(o[0][0] for o in cycle), tuple(o[0][1] for o in cycle))
        if winner is None or rank < winner[0]:
          winner = rank, cycle
      if winner is not None:
        break
    if winner is None:
      raise ValueError(f'source measure {run[0][0][0]["number"]}: no legal recurring {part} groove for {meter}; check harmonic releases and pattern vocabulary')
    cycle = winner[1]
    report = {'source_start': run[0][0][2], 'source_end': run[-1][-1][2], 'meter': meter,
      'grouping': list(grouping), 'grouping_evidence': grouping_evidence, 'eligible_patterns': len(patterns),
      'requested_patterns': requested, 'cycle_patterns': len(cycle), 'patterns': [o[1]['id'] for o in cycle],
      'anticipation': [o[2] for o in cycle], 'cost': winner[0][0], 'complete_cycles_compared': examined,
      'reduction': 'available distinct, legal patterns or source length' if len(cycle) < requested else None}
    reports.append(report)
    for i, group in enumerate(run):
      _, p, anticipate = cycle[i % len(cycle)]
      selected.append((p, anticipate, cache[i, p['id'], anticipate]))
      locations.append({'cycle_start': run[0][0][2], 'cycle_slot': i % len(cycle), 'base_pattern': p['id'], 'fill': None})
  add_fills(groups, selected, locations, part, melody, seed)
  return selected, locations, reports

def add_fills(groups, selected, locations, part, melody, seed):
  """Fixed source-position answers; no iteration seed or generated lead input."""
  policy = generator.vocabulary.accompaniment(part)
  spacing = policy['groove']['fill_spacing_measures']
  if not groups:
    return
  body_start, body_end = groups[0][0][2], groups[-1][-1][2] + 1
  body_length = body_end - body_start
  options = []
  if body_length < spacing:
    return
  for i, (group, (base, anticipate, current)) in enumerate(zip(groups, selected)):
    if len(group) != 2:
      continue
    last, _, index = group[-1]
    length = generator.model.fraction(last['duration'])
    reason = 'form_end' if i == len(groups) - 1 else 'section_end' if groups[i + 1][0][0].get('section_label') else None
    if reason is None and max((min(b, length) for a, b in melody[index]), default=fractions.Fraction()) <= length - 1:
      reason = 'melody_breath'
    if reason is None:
      continue
    middle = generator.model.fraction(group[0][0]['duration'])
    if any(e.start < middle < e.end for e in current.events):
      continue
    for variation in policy['patterns']:
      if variation['family'] != base['family'] or variation['lengths'] != base['lengths'] or variation.get('meter') != base.get('meter') or variation.get('grouping') != base.get('grouping'):
        continue
      hybrid = {**base, 'id': base['id'] + '.fill.' + variation['id'], 'events':
        [r for r in base['events'] if fractions.Fraction(r['onset']) < middle] +
        [r for r in variation['events'] if fractions.Fraction(r['onset']) >= middle]}
      candidate = realize(hybrid, group, part, melody, anticipate=anticipate)
      if candidate is None or candidate.cost >= current.cost - 1e-9:
        continue
      if tuple(e for e in candidate.events if e.start < middle) != tuple(e for e in current.events if e.start < middle):
        continue
      gain = current.cost - candidate.cost
      options.append((-gain, generator.solo.policy.stable_rank(seed, (index, hybrid['id'])), i, index, hybrid, anticipate, candidate, reason))
  chosen = []
  for _, _, i, index, pattern, anticipate, candidate, reason in sorted(options):
    if any(min(abs(index - other), body_length - abs(index - other)) < spacing for other in chosen):
      continue
    selected[i] = pattern, anticipate, candidate
    locations[i]['fill'] = {'pattern': pattern['id'], 'source_index': index, 'reason': reason}
    chosen.append(index)