"""Small permutation-equivariant policy baseline on the same current-state facts.

This is a standalone supervised MLP, not Qwen, Jev or a language-model adapter.
"""
import torch
from torch import nn

FEATURES=['lives/3','pellets_left/244','to_junction/10','frightened_seconds_left/6',
          'ghost_distance/15','ghost_present','ghost_coming','edible_distance/15',
          'edible_present','nearby_pellets/6','food_distance/40','food_present',
          'power_distance/25','power_present','is_back']

def encode(state,keys):
    global_features=[state['lives']/3,state['pellets_left']/244,state['to_junction']/10,
                     state.get('frightened_seconds_left',0)/6]
    rows=[]
    for key in keys:
        f=state['options'][key]
        dist=lambda field,scale: 1. if f.get(field) is None else f[field]/scale
        rows.append(global_features+[dist('ghost',15),float(f.get('ghost') is not None),
          float(f.get('ghost_coming',False)),dist('edible',15),float(f.get('edible') is not None),
          f['pellets']/6,dist('food',40),float(f.get('food') is not None),dist('power',25),
          float(f.get('power') is not None),float(key=='back')])
    return rows

class StructuredPolicy(nn.Module):
    def __init__(self):
        super().__init__()
        self.option=nn.Sequential(nn.Linear(15,64),nn.SiLU(),nn.Linear(64,64),nn.SiLU())
        self.score=nn.Sequential(nn.Linear(192,64),nn.SiLU(),nn.Linear(64,1))

    def forward(self,x,valid):
        h=self.option(x)
        mean=(h*valid[...,None]).sum(1)/valid.sum(1,keepdim=True)
        peak=h.masked_fill(~valid[...,None],-1e9).max(1).values
        pooled=torch.cat([h,mean[:,None,:].expand_as(h),peak[:,None,:].expand_as(h)],-1)
        return self.score(pooled).squeeze(-1).masked_fill(~valid,-1e9)