import os
import numpy as np
from tensorboardX import SummaryWriter
from envs.macf.env_wrappers import DummyVecEnv


class Runner(object):
    """Base class for non-learning PCFA runner.

    Only keeps environment handles, logging, and directory infrastructure.
    No neural network, trainer, or replay buffer is instantiated.
    """

    def __init__(self, config):
        self.config = config
        self.envs: DummyVecEnv = config.envs
        self.eval_envs = config.eval_envs
        self.device = config.device
        self.num_agents = config.num_agents

        # parameters
        self.env_name = config.env_name
        self.algorithm_name = config.algorithm_name
        self.experiment_name = config.experiment_name
        self.use_centralized_V = config.use_centralized_V
        self.use_obs_instead_of_state = config.use_obs_instead_of_state
        self.num_env_steps = config.num_env_steps
        self.episode_length = config.episode_length
        self.n_rollout_threads = config.n_rollout_threads
        self.n_eval_rollout_threads = config.n_eval_rollout_threads
        self.use_linear_lr_decay = config.use_linear_lr_decay
        self.hidden_size = config.hidden_size
        self.use_render = config.use_render
        self.recurrent_N = config.recurrent_N

        # interval
        self.save_interval = config.save_interval
        self.use_eval = config.use_eval
        self.eval_interval = config.eval_interval
        self.log_interval = config.log_interval
        self.log_info_interval = config.log_info_interval

        # dir
        self.model_dir = config.model_dir
        self.run_dir = config.run_dir
        self.log_dir = str(self.run_dir / 'logs')
        if not os.path.exists(self.log_dir):
            os.makedirs(self.log_dir)
        self.writter = SummaryWriter(self.log_dir)
        self.save_dir = str(self.run_dir / 'models')
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)

    def run(self):
        """Run rollouts."""
        raise NotImplementedError

    def warmup(self):
        """Optional warmup."""
        raise NotImplementedError

    def collect(self, step):
        """Collect actions."""
        raise NotImplementedError

    def insert(self, data):
        """Insert data into buffer (not used by PCFA)."""
        raise NotImplementedError

    def compute(self):
        """Compute returns (not used by PCFA)."""
        pass

    def train(self):
        """Train policy (not used by PCFA)."""
        return {}

    def save(self):
        """Save (no-op for PCFA)."""
        pass

    def restore(self):
        """Restore (no-op for PCFA)."""
        pass

    def log_train(self, train_infos, total_num_steps):
        """Log training info."""
        for k, v in train_infos.items():
            self.writter.add_scalar(k, v, total_num_steps)

    def log_env(self, env_infos, total_num_steps):
        """Log env info."""
        for k, v in env_infos.items():
            if len(v) > 0:
                self.writter.add_scalar(k, np.mean(v), total_num_steps)