from env.environment import SchoolInterventionEnv
from env.graders import grade

env = SchoolInterventionEnv()
state = env.reset()
print(f"RESET state: {state}\n")

for step in range(1, 21):
    action = "assign_tutor" if step <= 7 else "schedule_counseling" if step <= 14 else "peer_study_group"
    state, reward, done, info = env.step(action)
    if step in [1, 5, 10, 15, 20]:
        print(f"Step {step}: {state}")
    if done:
        break

print(f"\nFINAL state: {state}")
print("\nGrades on FINAL state:")
for task in ['easy', 'medium', 'hard']:
    result = grade(task, state)
    print(f"  {task}: score={result['score']:.4f}, passed={result['passed']}")