import unittest,sys
from pathlib import Path
sys.path.insert(0,str(Path(__file__).resolve().parent.parent))
import numpy as np
from field_orders import parallel_env,JSONPolicy
from pettingzoo.test import parallel_api_test

class EnvironmentAPI(unittest.TestCase):
 def test_parallel_contract(self):
  with parallel_env(team_size=4,max_cycles=50) as e:parallel_api_test(e,num_cycles=50)
 def test_seeded_reset_and_time_limit(self):
  with parallel_env(team_size=2,max_cycles=10) as e:
   a,_=e.reset(seed=94);e.step({k:np.zeros(5,dtype='f') for k in e.agents});b,_=e.reset(seed=94)
   for key in a:np.testing.assert_array_equal(a[key],b[key])
   for _ in range(12):
    if not e.agents:break
    o,r,t,tr,i=e.step({k:np.zeros(5,dtype='f') for k in e.agents})
   self.assertFalse(e.agents);self.assertTrue(all(tr.values()))
 def test_policy_and_render(self):
  with parallel_env(team_size=3,max_cycles=10,render_mode='rgb_array') as e:
   obs,_=e.reset(seed=15);p=JSONPolicy('neural-model/soldier-ppo-v2.json');actions=p.predict(obs)
   self.assertTrue(all(e.action_space(a).contains(v) for a,v in actions.items()));e.step(actions);self.assertEqual(e.render().shape,(650,1100,3))
 def test_invalid_actions(self):
  with parallel_env(team_size=1,max_cycles=2) as e:
   e.reset(seed=3)
   with self.assertRaises(ValueError):e.step({})

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