from types import SimpleNamespace
from models import SimpleDQN, DuelingDQN
import torch
import warnings

import ptan
import ptan.ignite as ptan_ignite
from ignite.engine import Engine
from ignite.metrics import RunningAverage
from ignite.contrib.handlers import tensorboard_logger as tb_logger

from datetime import timedelta, datetime

SEED = 0
DEVICE = torch.device('cuda')

HYPERPARAMETERS = {
    "Pong": SimpleNamespace(**{
        "env_name": "PongNoFrameskip-v4",
        "model_fn": SimpleDQN,
        "stop_reward": 18.0,
        "run_name": "pong",
        "replay_size": 100000,
        "replay_initial": 10000,
        "target_net_sync": 1000,
        "epsilon_frames": 100000,
        "epsilon_start": 1.0,
        "epsilon_end": 0.02,
        "learning_rate": 0.0001,
        "gamma": 0.99,
        "batch_size": 32
    }),
    "PongDueling": SimpleNamespace(**{
        "env_name": "PongNoFrameskip-v4",
        "model_fn": DuelingDQN,
        "stop_reward": 18.0,
        "run_name": "pong",
        "replay_size": 100000,
        "replay_initial": 10000,
        "target_net_sync": 1000,
        "epsilon_frames": 100000,
        "epsilon_start": 1.0,
        "epsilon_end": 0.02,
        "learning_rate": 0.0001,
        "gamma": 0.99,
        "batch_size": 32
    })
}

def setup_ignite(engine, params, exp_source, run_name, extra_metrics = []):
    # get rid of missing metrics warning
    warnings.simplefilter("ignore", category=UserWarning)

    handler = ptan_ignite.EndOfEpisodeHandler(
        exp_source, bound_avg_reward=params.stop_reward)
    handler.attach(engine)
    ptan_ignite.EpisodeFPSHandler().attach(engine)

    @engine.on(ptan_ignite.EpisodeEvents.EPISODE_COMPLETED)
    def episode_completed(trainer: Engine):
        passed = trainer.state.metrics.get('time_passed', 0)
        print("Episode %d: reward=%.0f, steps=%s, "
              "speed=%.1f f/s, elapsed=%s" % (
            trainer.state.episode, trainer.state.episode_reward,
            trainer.state.episode_steps,
            trainer.state.metrics.get('avg_fps', 0),
            timedelta(seconds=int(passed))))

    @engine.on(ptan_ignite.EpisodeEvents.BOUND_REWARD_REACHED)
    def game_solved(trainer: Engine):
        passed = trainer.state.metrics['time_passed']
        print("Game solved in %s, after %d episodes "
              "and %d iterations!" % (
            timedelta(seconds=int(passed)),
            trainer.state.episode, trainer.state.iteration))
        trainer.should_terminate = True

    now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    logdir = f"runs/{now}-{params.run_name}-{run_name}"
    tb = tb_logger.TensorboardLogger(log_dir=logdir)
    run_avg = RunningAverage(output_transform=lambda v: v['loss'])
    run_avg.attach(engine, "avg_loss")

    metrics = ['reward', 'steps', 'avg_reward']
    handler = tb_logger.OutputHandler(
        tag="episodes", metric_names=metrics)
    event = ptan_ignite.EpisodeEvents.EPISODE_COMPLETED
    tb.attach(engine, log_handler=handler, event_name=event)

    # write to tensorboard every 100 iterations
    ptan_ignite.PeriodicEvents().attach(engine)
    metrics = ['avg_loss', 'avg_fps']
    metrics.extend(extra_metrics)
    handler = tb_logger.OutputHandler(
        tag="train", metric_names=metrics,
        output_transform=lambda a: a)
    event = ptan_ignite.PeriodEvents.ITERS_100_COMPLETED
    tb.attach(engine, log_handler=handler, event_name=event)
from types import SimpleNamespace
from models import SimpleDQN, DuelingDQN
import torch
import warnings

import ptan
import ptan.ignite as ptan_ignite
from ignite.engine import Engine
from ignite.metrics import RunningAverage
from ignite.contrib.handlers import tensorboard_logger as tb_logger

from datetime import timedelta, datetime

SEED = 0
DEVICE = torch.device('cuda')

HYPERPARAMETERS = {
    "Pong": SimpleNamespace(**{
        "env_name": "PongNoFrameskip-v4",
        "model_fn": SimpleDQN,
        "stop_reward": 18.0,
        "run_name": "pong",
        "replay_size": 100000,
        "replay_initial": 10000,
        "target_net_sync": 1000,
        "epsilon_frames": 100000,
        "epsilon_start": 1.0,
        "epsilon_end": 0.02,
        "learning_rate": 0.0001,
        "gamma": 0.99,
        "batch_size": 32
    }),
    "PongDueling": SimpleNamespace(**{
        "env_name": "PongNoFrameskip-v4",
        "model_fn": DuelingDQN,
        "stop_reward": 18.0,
        "run_name": "pong",
        "replay_size": 100000,
        "replay_initial": 10000,
        "target_net_sync": 1000,
        "epsilon_frames": 100000,
        "epsilon_start": 1.0,
        "epsilon_end": 0.02,
        "learning_rate": 0.0001,
        "gamma": 0.99,
        "batch_size": 32
    })
}

def setup_ignite(engine, params, exp_source, run_name, extra_metrics = []):
    # get rid of missing metrics warning
    warnings.simplefilter("ignore", category=UserWarning)

    handler = ptan_ignite.EndOfEpisodeHandler(
        exp_source, bound_avg_reward=params.stop_reward)
    handler.attach(engine)
    ptan_ignite.EpisodeFPSHandler().attach(engine)

    @engine.on(ptan_ignite.EpisodeEvents.EPISODE_COMPLETED)
    def episode_completed(trainer: Engine):
        passed = trainer.state.metrics.get('time_passed', 0)
        print("Episode %d: reward=%.0f, steps=%s, "
              "speed=%.1f f/s, elapsed=%s" % (
            trainer.state.episode, trainer.state.episode_reward,
            trainer.state.episode_steps,
            trainer.state.metrics.get('avg_fps', 0),
            timedelta(seconds=int(passed))))

    @engine.on(ptan_ignite.EpisodeEvents.BOUND_REWARD_REACHED)
    def game_solved(trainer: Engine):
        passed = trainer.state.metrics['time_passed']
        print("Game solved in %s, after %d episodes "
              "and %d iterations!" % (
            timedelta(seconds=int(passed)),
            trainer.state.episode, trainer.state.iteration))
        trainer.should_terminate = True

    now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    logdir = f"runs/{now}-{params.run_name}-{run_name}"
    tb = tb_logger.TensorboardLogger(log_dir=logdir)
    run_avg = RunningAverage(output_transform=lambda v: v['loss'])
    run_avg.attach(engine, "avg_loss")

    metrics = ['reward', 'steps', 'avg_reward']
    handler = tb_logger.OutputHandler(
        tag="episodes", metric_names=metrics)
    event = ptan_ignite.EpisodeEvents.EPISODE_COMPLETED
    tb.attach(engine, log_handler=handler, event_name=event)

    # write to tensorboard every 100 iterations
    ptan_ignite.PeriodicEvents().attach(engine)
    metrics = ['avg_loss', 'avg_fps']
    metrics.extend(extra_metrics)
    handler = tb_logger.OutputHandler(
        tag="train", metric_names=metrics,
        output_transform=lambda a: a)
    event = ptan_ignite.PeriodEvents.ITERS_100_COMPLETED
    tb.attach(engine, log_handler=handler, event_name=event)
# get rid of missing metrics warning # write to tensorboard every 100 iterations