"""Zenith-specific permanent-trail and health-aware scoring."""
from snakes.engine.health_routes import health_route_cost
from snakes.engine.survival_search import CompactSurvivalSearch

class ZenithSurvivalMixin:
  def _choose_constrictor_move(
    self, safe_moves, my_body, my_len, my_health, other_snakes, food_set,
    hazard_set, hazard_damage, hazard_count, previous_hazard_set,
    enemy_attack_map, enemy_can_grow, total_occupancy, width, height, deadline,
  ):
    # In constrictor bodies never vacate. Remaining free cells are a survival
    # horizon, not storage for the existing body. Requiring free space >= body
    # length incorrectly marks every late-game candidate dead.
    bb = self._get_bb(width, height)
    occupied = bb.set_to_bits({
      (p['x'], p['y']) for body in [my_body, *(s['body'] for s in other_snakes)]
      for p in body
    })
    scores = {}
    for move, pos in safe_moves.items():
      target = bb.idx(pos['x'], pos['y'])
      space = bb.flood_count(target, occupied)
      liberties = bb.open_neighbor_count(target, occupied)
      danger = enemy_attack_map.get((pos['x'], pos['y']), 0) >= my_len
      # Preserve free capacity first. Enemy confinement must not outweigh our
      # own ability to survive, and tail access has no meaning in this mode.
      scores[move] = space * 10.0 + liberties * 8.0 - (10000 if danger else 0)
      if liberties == 0:
        scores[move] -= 5000

    # Give each candidate the same shallow horizon before considering deeper
    # layers. Never let the first candidate consume everyone else's budget.
    search = CompactSurvivalSearch(
      bb, set(), True, deadline, branch=3, enemy_branch=2, response_cap=8,
    )
    ranked = scores
    for depth in range(1, 5):
      layer = {}
      for move, pos in safe_moves.items():
        if self._remaining_ms(deadline) < 15:
          break
        before = search.deadline_exits
        value = search.search_selected(my_body, other_snakes, (pos['x'], pos['y']), depth)
        if search.deadline_exits != before:
          break
        layer[move] = value
      if len(layer) != len(safe_moves):
        break
      # Rollout can veto immediate/near-term death, but cannot turn temporary
      # enemy confinement into an overwhelming speculative attack reward.
      ranked = {m: scores[m] + v * 0.15 for m, v in layer.items()}
      best = max(ranked, key=ranked.get)
      if depth == 4:
        return best, ranked
    return max(ranked, key=ranked.get), ranked

  def _score_move(self, *args, **kwargs):
    score, info = super()._score_move(*args, **kwargs)
    # Apex calls this hook with keywords. Constrictor uses its dedicated path.
    if not kwargs or kwargs['is_constrictor']:
      return score, info
    point = (kwargs['pos']['x'], kwargs['pos']['y'])
    food = kwargs['food_set']
    hazards = kwargs['hazard_set']
    damage = kwargs['hazard_damage']
    stacks = kwargs['hazard_count']
    health = 100 if point in food else kwargs['my_health'] - 1 - (
      damage * stacks.get(point, 1) if point in hazards else 0
    )
    if health <= 0:
      score -= 20000
      info['is_survivable'] = False
      return score, info
    if not hazards:
      return score, info
    blocked = info['blocked'] - {point}
    cost = health_route_cost(
      point, food, hazards, stacks, damage, blocked,
      kwargs['width'], kwargs['height'], refuge=True,
    )
    # Existing hazards deal damage on entry even when newly seen. Prefer
    # leaving lethal hazard corridors before low-health urgency kicks in.
    if cost is None or cost >= health:
      score -= 12000
      info['is_survivable'] = False
    else:
      score -= cost * 8.0
    if health < 55:
      food_cost = health_route_cost(
        point, food, hazards, stacks, damage, blocked,
        kwargs['width'], kwargs['height'],
      )
      if food_cost is None or food_cost >= health:
        score -= 2500
      else:
        score += (55 - health) * 12.0 / (food_cost + 1)
    return score, info