from typing import Any, Dict, Tuple, Union
import os
import gymnasium
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.logger import HumanOutputFormat, KVWriter, Logger

from aim import Run
import numpy as np


class AimOutputFormat(KVWriter):
    """
    Track key/value pairs into Aim run.
    """

    def __init__(
        self,
        aim_run
    ):
        self.aim_run = aim_run

    def write(
        self,
        key_values: Dict[str, Any],
        key_excluded: Dict[str, Union[str, Tuple[str, ...]]],
        step: int = 0,
    ) -> None:
        for (key, value), (_, excluded) in zip(
            sorted(key_values.items()), sorted(key_excluded.items())
        ):
            if excluded is not None and 'aim' in excluded:
                continue

            if isinstance(value, np.ScalarType):
                if not isinstance(value, str):
                    tag, key = key.split('/')
                    if tag in ['train', 'valid']:
                        context = {'subset': tag}
                    else:
                        context = {'tag': tag}

                    self.aim_run.track(value, key, step=step, context=context)


run = Run()
loggers = Logger(
    folder=None,
    output_formats=[AimOutputFormat(run)],
)

environment_name = 'CartPole-v1'
env = gymnasium.make(environment_name)

env = DummyVecEnv([lambda: env])
model = PPO('MlpPolicy', env, verbose = 1)
model.set_logger(loggers)
model.learn(total_timesteps=10_000)