# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This software may be used and distributed according to the terms of the
# GNU General Public License version 2.

import logging

from typing import Dict, Optional, Sequence

import hydra
from omegaconf import DictConfig, OmegaConf

import numpy as np

import torch
import torch.nn

import rlmeta.utils.nested_utils as nested_utils

from rlmeta.agents.ppo.ppo_agent import PPOAgent
from rlmeta.core.types import Action, TimeStep
from rlmeta.envs.env import Env
from rlmeta.utils.stats_dict import StatsDict

import model_utils

from cache_env_wrapper import CacheEnvCycloneWrapperFactory


def batch_obs(timestep: TimeStep) -> TimeStep:
    obs, reward, terminated, truncated, info = timestep
    return TimeStep(obs.unsqueeze(0), reward, terminated, truncated, info)


def unbatch_action(action: Action) -> Action:
    act, info = action
    act.squeeze_(0)
    info = nested_utils.map_nested(lambda x: x.squeeze(0), info)
    return Action(act, info)


def run_loop(env: Env,
             agent: PPOAgent,
             victim_addr: int = -1) -> Dict[str, float]:
    episode_length = 0
    episode_return = 0.0
    num_guess = 0
    num_correct = 0
    cyclone_attack = 0

    if victim_addr == -1:
        timestep = env.reset()
    else:
        timestep = env.reset(victim_address=victim_addr)

    agent.observe_init(timestep)
    while not (timestep.terminated or timestep.truncated):
        # Model server requires a batch_dim, so unsqueeze here for local runs.
        timestep = batch_obs(timestep)
        action = agent.act(timestep)
        # Unbatch the action.
        action = unbatch_action(action)

        timestep = env.step(action)
        agent.observe(action, timestep)

        episode_length += 1
        episode_return += timestep.reward

        if "guess_correct" in timestep.info:
            num_guess += 1
            if timestep.info["guess_correct"]:
                num_correct += 1

        cyclone_attack += timestep.info.get("cyclone_attack", 0)

    metrics = {
        "episode_length": episode_length,
        "episode_return": episode_return,
        "num_guess": num_guess,
        "num_correct": num_correct,
        "correct_rate": num_correct / num_guess,
        "bandwith": num_guess / episode_length,
        "cyclone_attack": cyclone_attack,
    }

    return metrics


def run_loops(env: Env,
              agent: PPOAgent,
              num_episodes: int = -1,
              seed: int = 0,
              reset_cache_state: bool = False) -> StatsDict:
    # env.seed(seed)
    env.reset(seed=seed)
    metrics = StatsDict()

    num_guess = 0
    num_correct = 0
    tot_length = 0

    if num_episodes == -1:
        start = env.env.victim_address_min
        stop = env.env.victim_address_max + 1 + int(
            env.env._env.allow_empty_victim_access)
        for victim_addr in range(start, stop):
            cur_metrics = run_loop(env, agent, victim_addr=victim_addr)
            num_guess += cur_metrics["num_guess"]
            num_correct += cur_metrics["num_correct"]
            tot_length += cur_metrics["episode_length"]
            metrics.extend(cur_metrics)
    else:
        for _ in range(num_episodes):
            cur_metrics = run_loop(env, agent, victim_addr=-1)
            num_guess += cur_metrics["num_guess"]
            num_correct += cur_metrics["num_correct"]
            tot_length += cur_metrics["episode_length"]
            metrics.extend(cur_metrics)

    metrics.add("overall_correct_rate", num_correct / num_guess)
    metrics.add("overall_bandwith", num_guess / tot_length)

    return metrics


@hydra.main(config_path="./config", config_name="sample_cchunter")
def main(cfg):
    # Create env
    cfg.env_config.verbose = 1
    env_fac = CacheEnvCycloneWrapperFactory(
        OmegaConf.to_container(cfg.env_config))
    env = env_fac(index=0)

    # Load model
    model = model_utils.get_model(cfg.model_config, cfg.env_config.window_size,
                                  env.action_space.n, cfg.checkpoint)
    model.eval()

    # Create agent
    agent = PPOAgent(model, deterministic_policy=cfg.deterministic_policy)

    # Run loops
    metrics = run_loops(env, agent, cfg.num_episodes, cfg.seed)
    logging.info("\n\n" + metrics.table(info="sample") + "\n")


if __name__ == "__main__":
    main()
