import gymnasium as gym
from gymnasium.wrappers import FlattenObservation
import pgtg
from stable_baselines3 import DQN
from dsmc_tool.evaluator import Evaluator
import dsmc_tool.property as prop

env = gym.make("pgtg-v3")
env = FlattenObservation(env)

model = DQN("MlpPolicy", env,verbose=1)
model.learn(total_timesteps=1000, log_interval=100)

evaluator = Evaluator(env=env, initial_episodes=100, evaluation_episodes=50)
property = prop.ActionDiversityProperty(num_actions=9)
evaluator.register_property(property)
property = prop.ActionEntropyProperty(num_actions=9)
evaluator.register_property(property)
property = prop.ActionTakenProperty(action=4)
evaluator.register_property(property)
property = prop.ActionThresholdProperty(action=7, threshold=5)
evaluator.register_property(property)
property = prop.ActionVarietyProperty(threshold=5)
evaluator.register_property(property)
property = prop.ConsecutiveSameActionProperty(action=3, threshold=5)
evaluator.register_property(property)
property = prop.EarlyTerminationProperty(step_maximum=10)
evaluator.register_property(property)
property = prop.EpisodeLengthProperty()
evaluator.register_property(property)
property = prop.GoalBeforeStepLimitProperty(goal_reward=20, step_limit=10)
evaluator.register_property(property)
property = prop.GoalReachingProbabilityProperty(goal_reward=20)
evaluator.register_property(property)
property = prop.NormalizedReturnProperty()
evaluator.register_property(property)
property = prop.PathEfficiencyProperty()
evaluator.register_property(property)
property = prop.PathLengthEfficiencyProperty(optimal_path_length=10)
evaluator.register_property(property)
property = prop.ReturnProperty()
evaluator.register_property(property)
property = prop.ReturnThresholdProperty(threshold=0)
evaluator.register_property(property)
property = prop.RewardToLengthRatioProperty()
evaluator.register_property(property)
property = prop.RewardVarianceProperty()
evaluator.register_property(property)
property = prop.StateCoverageProperty(num_states=506)
evaluator.register_property(property)
property = prop.StateTransitionSmoothnessProperty()
evaluator.register_property(property)
results = evaluator.eval(model, epsilon=0.05, kappa=0.025, save_interim_results=True)