import unittest
import torch
from lightnest import rv32
from lightnest.learned_patch import MAX_WINDOW
from lightnest.sequential_stop_rl import rollout, terminal_reward

class Sequence(torch.nn.Module):
    def __init__(self, positions):
        super().__init__(); self.positions=positions; self.calls=0
    def forward(self,c,b):
        scores=torch.full((len(c),257),-100.)
        scores[:,self.positions[self.calls]]=100.; self.calls+=1
        return scores,torch.ones((len(c),256,8))*100

class StopTests(unittest.TestCase):
    def test_three_edits_then_none(self):
        p=Sequence([0,1,2,MAX_WINDOW]); out,e,h,lp=rollout(p,['test'],[b'1234'])
        self.assertEqual(out,[bytes([255,255,255,52])]);self.assertEqual(e,[3]);self.assertEqual(len(h[0]),4)
    def test_fourth_edit_is_not_applied(self):
        p=Sequence([0,1,2,3]);out,e,h,_=rollout(p,['test'],[b'1234'])
        self.assertEqual(out,[bytes([255,255,255,52])]);self.assertEqual(e,[3]);self.assertEqual(h[0][-1][0],3)
    def test_reward_requires_stop_and_charges_proposals(self):
        binary=b''.join(w.to_bytes(4,'little') for w in (rv32.addi('x10','x0',7),rv32.EBREAK)); cases=[([1,2],7)]
        self.assertAlmostEqual(terminal_reward(binary,binary,cases,0,[(MAX_WINDOW,0)]),.98)
        self.assertLess(terminal_reward(binary,binary,cases,0,[(0,0)]),0)
        self.assertLess(terminal_reward(binary,binary,cases,1,[(MAX_WINDOW,0)]),0)
        self.assertLess(terminal_reward(binary,binary,[([1,2],9)],0,[(MAX_WINDOW,0)]),0)
    def test_early_none(self):
        p=Sequence([MAX_WINDOW]);out,e,h,_=rollout(p,['test'],[b'1234'])
        self.assertEqual(out,[b'1234']);self.assertEqual(e,[0]);self.assertEqual(len(h[0]),1)

if __name__=='__main__':unittest.main()