import unittest
from pathlib import Path
import torch
from baim.policy import LearnedPolicy
from baim.features import encode, candidates, lexical
from baim.state import State, Element
from baim.actions import Ticket, Kind

class BatchTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        torch.set_num_threads(2)
        cls.policy=LearnedPolicy(Path(__file__).resolve().parents[1]/'models/v000-mean')

    def test_batch_isolates_tickets_refs_and_sensitive_targets(self):
        requests=[]
        for i in range(4):
            elements=() if i==3 else (Element(f'{i}-safe','button','amber birch','0'),
                Element(f'{i}-secret','textbox','password','1',sensitive=True))
            state=State(str(i),i,'fixture',elements)
            ticket=Ticket('session',str(i),0,i,str(i),i,state.state_hash)
            requests.append(('Click amber birch.',state,ticket))
        batch=self.policy.predict_batch(requests)
        for request,decision in zip(requests,batch):
            self.assertEqual(decision.action,self.policy.predict(*request).action)
            self.assertEqual(decision.ticket,request[2])
            if decision.action.kind!=Kind.ASK_USER:
                self.assertEqual(decision.action.args[0],request[1].elements[0].ref)
        self.assertEqual(batch[-1].action.kind,Kind.ASK_USER)
        self.assertEqual(self.policy.predict_batch([]),[])
        with self.assertRaises(ValueError):self.policy.predict_batch(requests*17)

    def test_encoding_preserves_retrieval_and_lexical_features(self):
        row={'goal':'Click café', 'elements':[{'role':'button','name':'café'},
            {'role':'button','name':'café','sensitive':True},
            {'role':'button','name':'other'}, {'role':'button','name':'café','visible':False}]}
        inputs,_,_,maps=encode([row],self.policy.vocab)
        self.assertEqual(maps[0],candidates(row['goal'],row['elements']))
        for i,index in enumerate(maps[0]):
            self.assertTrue(torch.equal(inputs[2][0,i],torch.tensor(lexical(row['goal'],row['elements'][index]))))
        self.assertEqual(encode([],self.policy.vocab)[0][0].shape[0],0)

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