import gym
from gym.wrappers import RecordVideo
from matplotlib import pyplot as plt

from interpretable.interpretable import InterpretablePolicyExtractor
from interpretable.utils import generate_dataset_from_expert, rollouts

if __name__ == "__main__":
    env_name = "CartPole-v1"
    dataset_path = generate_dataset_from_expert("ppo", env_name, force=True)
    ipe = InterpretablePolicyExtractor(env_name)
    results = ipe.train_from_dataset(dataset_path)
    ipe.policy.prune()
    ipe.policy.plot(mask=True)
    plt.savefig("kan-policy.png")

    env = gym.make(env_name, render_mode="rgb_array")
    env = RecordVideo(env, video_folder="videos", episode_trigger=lambda x: True, name_prefix=f"kan-{env_name}")

    ipe.policy.auto_symbolic()

    ipe.policy.plot(mask=True)
    plt.savefig("sym-policy.png")
    print(ipe.policy.symbolic_formula())

    rollouts(env, ipe.forward, 2)