import os
import time
import numpy as np
from datetime import datetime

import torch
import torch.nn as nn
from torch import optim
from torch.utils.tensorboard import SummaryWriter

from rl_games.algos_torch import model_builder, torch_ext
from rl_games.common import vecenv, schedulers, experience
from rl_games.common.a2c_common import print_statistics
from rl_games.interfaces.base_algorithm import BaseAlgorithm


class SACAgent(BaseAlgorithm):

    def __init__(self, base_name, params):

        self.config = config = params['config']
        print(config)

        # TODO: Get obs shape and self.network
        self.load_networks(params)
        self.base_init(base_name, config)
        self.critic_tau = float(self.config.get("critic_tau", 0.005))  # Small tau for stable smoothing
        self.gamma = self.config.get("gamma", 0.99)
        self.gamma_tensor = torch.tensor(self.gamma, device=self._device, dtype=torch.float32)
        self.num_steps_per_episode = self.config.get("num_steps_per_episode", 1)
        self.num_updates_per_step = self.config.get("num_updates_per_step", 1)
        # Warmup: num_warmup_frames (total env frames) takes priority over num_warmup_steps (epochs)
        num_warmup_frames = self.config.get("num_warmup_frames", None)
        if num_warmup_frames is not None:
            frames_per_epoch = self.num_actors * self.num_steps_per_episode
            self.num_warmup_steps = int(np.ceil(num_warmup_frames / frames_per_epoch))
        else:
            self.num_warmup_steps = self.config.get("num_warmup_steps", 1000)
        self.batch_size = config["batch_size"]
        self.init_alpha = config["init_alpha"]
        self.learnable_temperature = config["learnable_temperature"]
        self.replay_buffer_size = config["replay_buffer_size"]
        self.normalize_input = config.get("normalize_input", False)
        self.enable_mixed_precision = config.get("mixed_precision", False)

        # Policy frequency: update actor every N critic updates (CleanRL uses 2)
        self.policy_frequency = config.get("policy_frequency", 2)
        self.update_counter = 0  # Track number of critic updates

        # Gradient clipping for stability (0 = disabled)
        self.critic_grad_clip = config.get("critic_grad_clip", 5.0)

        # Timeout handling - consistent with PPO
        self.value_bootstrap = config.get('value_bootstrap', True)

        print(self.batch_size, self.num_actors, self.num_agents)

        self.num_frames_per_epoch = self.num_actors * self.num_steps_per_episode

        self.log_alpha = torch.tensor(np.log(self.init_alpha)).float().to(self._device)
        self.log_alpha.requires_grad = True
        action_space = self.env_info['action_space']
        self.actions_num = action_space.shape[0]

        # Store proper action bounds per dimension
        self.action_low = torch.tensor(action_space.low, device=self._device, dtype=torch.float32)
        self.action_high = torch.tensor(action_space.high, device=self._device, dtype=torch.float32)
        self.action_scale = (self.action_high - self.action_low) / 2.0
        self.action_bias = (self.action_high + self.action_low) / 2.0
        self.action_range = [
            float(self.env_info['action_space'].low.min()),
            float(self.env_info['action_space'].high.max())
        ]

        obs_shape = torch_ext.shape_whc_to_cwh(self.obs_shape)
        net_config = {
            'obs_dim': self.env_info["observation_space"].shape[0],
            'action_dim': self.env_info["action_space"].shape[0],
            'actions_num': self.actions_num,
            'input_shape': obs_shape,
            'normalize_input': self.normalize_input,
        }
        self.model = self.network.build(net_config)
        self.model.to(self._device)

        print("Number of Agents", self.num_actors, "Batch Size", self.batch_size)

        self.actor_optimizer = optim.Adam(self.model.sac_network.actor.parameters(),
                                        lr=float(self.config['actor_lr']),
                                        betas=self.config.get("actor_betas", [0.9, 0.999]),
                                        fused=True)

        self.critic_optimizer = optim.Adam(self.model.sac_network.critic.parameters(),
                                        lr=float(self.config["critic_lr"]),
                                        betas=self.config.get("critic_betas", [0.9, 0.999]),
                                        fused=True)

        self.log_alpha_optimizer = torch.optim.Adam([self.log_alpha],
                                        lr=float(self.config["alpha_lr"]),
                                        betas=self.config.get("alphas_betas", [0.9, 0.999]),
                                        fused=True)

        self.replay_buffer = experience.VectorizedReplayBuffer(self.env_info['observation_space'].shape,
            self.env_info['action_space'].shape,
            self.replay_buffer_size,
            self._device)

        self.target_entropy_coef = config.get("target_entropy_coef", 1.0)
        self.target_entropy = self.target_entropy_coef * -self.env_info['action_space'].shape[0]
        print("Target entropy", self.target_entropy)

        self.algo_observer = config['features']['observer']

        self.amp_dtype = torch.bfloat16 if self.enable_mixed_precision else torch.float32

        # ────────────────────────────────────
        # Summary writer
        # ────────────────────────────────────
        self.global_rank = int(os.getenv("RANK", "0"))

        if self.global_rank == 0:
            self.writer = SummaryWriter(self.summaries_dir)
        else:
            class _DummyWriter:
                def add_scalar(self, *_, **__): pass
                def close(self): pass
            self.writer = _DummyWriter()

        self.gamma_tensor = torch.tensor(self.gamma, device=self._device, dtype=self.amp_dtype)

    def load_networks(self, params):
        builder = model_builder.ModelBuilder()
        self.config['network'] = builder.load(params)

    def base_init(self, base_name, config):
        self.env_config = config.get('env_config', {})
        self.num_actors = config.get('num_actors', 1)
        self.env_name = config['env_name']
        print("Env name:", self.env_name)

        self.env_info = config.get('env_info')
        if self.env_info is None:
            self.vec_env = vecenv.create_vec_env(self.env_name, self.num_actors, **self.env_config)
            self.env_info = self.vec_env.get_env_info()

        self._device = config.get('device', 'cuda:0')

        # For Isaac gym compatibility
        self.ppo_device = self._device
        print('Env info:')
        print(self.env_info)

        self.rewards_shaper = config['reward_shaper']
        self.observation_space = self.env_info['observation_space']
        self.weight_decay = config.get('weight_decay', 0.0)
        # self.use_action_masks = config.get('use_action_masks', False)
        self.is_train = config.get('is_train', True)

        # MSELoss is standard for SAC - matches CleanRL and original paper
        self.c_loss = nn.MSELoss()

        self.save_best_after = config.get('save_best_after', 500)
        self.print_stats = config.get('print_stats', True)
        self.rnn_states = None
        self.name = base_name

        self.max_epochs = self.config.get('max_epochs', -1)
        self.max_frames = self.config.get('max_frames', -1)

        self.save_freq = config.get('save_frequency', 0)

        self.network = config['network']
        self.rewards_shaper = config['reward_shaper']
        self.num_agents = self.env_info.get('agents', 1)
        self.obs_shape = self.observation_space.shape

        self.games_to_track = self.config.get('games_to_track', 100)
        self.game_rewards = torch_ext.AverageMeter(1, self.games_to_track).to(self._device)
        self.game_lengths = torch_ext.AverageMeter(1, self.games_to_track).to(self._device)
        self.obs = None

        # Min alpha in EXP space (not log space!) for comparison with self.alpha
        self.min_alpha = torch.tensor(0.01).float().to(self._device)

        self.frame = 0
        self.epoch_num = 0
        self.update_time = 0
        self.last_mean_rewards = -1000000000
        self.play_time = 0

        # TODO: put it into the separate class
        pbt_str = ''
        self.population_based_training = config.get('population_based_training', False)
        if self.population_based_training:
            # in PBT, make sure experiment name contains a unique id of the policy within a population
            pbt_str = f'_pbt_{config["pbt_idx"]:02d}'
        full_experiment_name = config.get('full_experiment_name', None)
        if full_experiment_name:
            print(f'Exact experiment name requested from command line: {full_experiment_name}')
            self.experiment_name = full_experiment_name
        else:
            self.experiment_name = config['name'] + pbt_str + datetime.now().strftime("_%d-%H-%M-%S")
        self.train_dir = config.get('train_dir', 'runs')

        # a folder inside of train_dir containing everything related to a particular experiment
        self.experiment_dir = os.path.join(self.train_dir, self.experiment_name)

        # folders inside <train_dir>/<experiment_dir> for a specific purpose
        self.nn_dir = os.path.join(self.experiment_dir, 'nn')
        self.summaries_dir = os.path.join(self.experiment_dir, 'summaries')

        os.makedirs(self.train_dir, exist_ok=True)
        os.makedirs(self.experiment_dir, exist_ok=True)
        os.makedirs(self.nn_dir, exist_ok=True)
        os.makedirs(self.summaries_dir, exist_ok=True)

        self.is_tensor_obses = False
        self.is_rnn = False
        self.last_rnn_indices = None
        self.last_state_indices = None

    def init_tensors(self):
        batch_size = self.num_agents * self.num_actors

        self.current_rewards = torch.zeros(batch_size, dtype=torch.float32, device=self._device)
        self.current_lengths = torch.zeros(batch_size, dtype=torch.float32, device=self._device)

        self.dones = torch.zeros((batch_size,), dtype=torch.uint8, device=self._device)

    @property
    def alpha(self):
        return self.log_alpha.exp()

    @property
    def device(self):
        return self._device

    def get_weights(self):
        print("Loading weights")
        state = {'actor': self.model.sac_network.actor.state_dict(),
         'critic': self.model.sac_network.critic.state_dict(),
         'critic_target': self.model.sac_network.critic_target.state_dict()}
        if self.normalize_input:
            state['running_mean_std'] = self.model.running_mean_std.state_dict()
        return state

    def save(self, fn):
        state = self.get_full_state_weights()
        torch_ext.save_checkpoint(fn, state)

    def set_weights(self, weights):
        self.model.sac_network.actor.load_state_dict(weights['actor'])
        self.model.sac_network.critic.load_state_dict(weights['critic'])
        self.model.sac_network.critic_target.load_state_dict(weights['critic_target'])

        if self.normalize_input and 'running_mean_std' in weights:
            self.model.running_mean_std.load_state_dict(weights['running_mean_std'])

    def get_full_state_weights(self):
        print("Loading full weights")
        state = self.get_weights()

        state['epoch'] = self.epoch_num
        state['frame'] = self.frame
        state['actor_optimizer'] = self.actor_optimizer.state_dict()
        state['critic_optimizer'] = self.critic_optimizer.state_dict()
        state['log_alpha_optimizer'] = self.log_alpha_optimizer.state_dict()        

        return state

    def set_full_state_weights(self, weights, set_epoch=True):
        self.set_weights(weights)

        if set_epoch:
            self.epoch_num = weights['epoch']
            self.frame = weights['frame']

        self.actor_optimizer.load_state_dict(weights['actor_optimizer'])
        self.critic_optimizer.load_state_dict(weights['critic_optimizer'])
        self.log_alpha_optimizer.load_state_dict(weights['log_alpha_optimizer'])

        self.last_mean_rewards = weights.get('last_mean_rewards', -1000000000)

        if self.vec_env is not None:
            env_state = weights.get('env_state', None)
            self.vec_env.set_env_state(env_state)

    def restore(self, fn, set_epoch=True):
        if not os.path.exists(fn):
            raise FileNotFoundError(f"Checkpoint file not found: {fn}")
        checkpoint = torch_ext.load_checkpoint(fn)
        self.set_full_state_weights(checkpoint, set_epoch=set_epoch)

    def get_param(self, param_name):
        pass

    def set_param(self, param_name, param_value):
        pass

    def get_masked_action_values(self, obs, action_masks):
        raise NotImplementedError("Masked action values are not supported in SAC agent")

    def set_eval(self):
        self.model.eval()

    def set_train(self):
        self.model.train()

    def rescale_actions(self, actions):
        """Rescale actions from [-1, 1] to actual action space bounds."""
        rescaled = actions * self.action_scale + self.action_bias
        # Clamp to handle floating point errors
        return rescaled.clamp(self.action_low, self.action_high)

    def update_critic(self, obs, action, reward, next_obs, not_done, step):
        with torch.no_grad():
            dist = self.model.actor(next_obs)
            next_action = dist.rsample()
            log_prob = dist.log_prob(next_action).sum(-1, keepdim=True)
            # Rescale next_action for critic (as it outputs [-1, 1])
            next_action_rescaled = self.rescale_actions(next_action)

            target_Q1, target_Q2 = self.model.critic_target(
                next_obs.float(),
                next_action_rescaled.float()
            )
            target_V = torch.min(target_Q1, target_Q2) - self.alpha * log_prob
            target_V = target_V.detach()

        # get current Q estimates and build target_Q inside autocast
        with torch.amp.autocast('cuda', enabled=self.enable_mixed_precision, dtype=self.amp_dtype):
            current_Q1, current_Q2 = self.model.critic(obs, action)

            # Build target_Q inside autocast to match dtypes
            target_Q = reward + (not_done * self.gamma_tensor * target_V)

            critic1_loss = self.c_loss(current_Q1, target_Q)
            critic2_loss = self.c_loss(current_Q2, target_Q)
            critic_loss = 0.5 * (critic1_loss + critic2_loss)

        self.critic_optimizer.zero_grad(set_to_none=True)
        critic_loss.backward()
        if self.critic_grad_clip > 0:
            torch.nn.utils.clip_grad_norm_(self.model.sac_network.critic.parameters(), self.critic_grad_clip)
        self.critic_optimizer.step()

        return critic_loss.detach(), critic1_loss.detach(), critic2_loss.detach()

    def update_actor_and_alpha(self, obs, step):
        for p in self.model.sac_network.critic.parameters():
            p.requires_grad = False

        with torch.amp.autocast('cuda', enabled=self.enable_mixed_precision, dtype=self.amp_dtype):
            dist = self.model.actor(obs)
            action = dist.rsample()
            # Rescale actions for critic evaluation
            action_rescaled = self.rescale_actions(action)
            log_prob = dist.log_prob(action).sum(-1, keepdim=True)
            entropy = -log_prob.mean()
            actor_Q1, actor_Q2 = self.model.critic(obs, action_rescaled)
            actor_Q = torch.min(actor_Q1, actor_Q2)

            actor_loss = (torch.max(self.alpha.detach(), self.min_alpha) * log_prob - actor_Q).mean()

            if self.learnable_temperature:
                alpha_loss = (self.alpha * (-log_prob - self.target_entropy).detach()).mean()
                total_loss = actor_loss + alpha_loss
            else:
                alpha_loss = None
                total_loss = actor_loss

        self.actor_optimizer.zero_grad(set_to_none=True)
        if self.learnable_temperature:
            self.log_alpha_optimizer.zero_grad(set_to_none=True)

        total_loss.backward()

        self.actor_optimizer.step()
        if self.learnable_temperature:
            self.log_alpha_optimizer.step()

        for p in self.model.sac_network.critic.parameters():
            p.requires_grad = True

        return actor_loss.detach(), entropy.detach(), self.alpha.detach(), alpha_loss

    def soft_update_params(self, net, target_net, tau):
        with torch.no_grad():
            for src, tgt in zip(net.parameters(), target_net.parameters()):
                tgt.data.lerp_(src.data, tau)

    def update_target_weights(self):
        tau = self.critic_tau
        self.soft_update_params(
            self.model.sac_network.critic,
            self.model.sac_network.critic_target,
            tau
        )

    def update(self, step):
        obs, action, reward, next_obs, done = self.replay_buffer.sample(self.batch_size)
        not_done = ~done

        obs = self.preproc_obs(obs)
        next_obs = self.preproc_obs(next_obs)

        critic_loss, critic1_loss, critic2_loss = self.update_critic(obs, action, reward, next_obs, not_done, step)
        self.update_counter += 1

        # Update actor and alpha only every policy_frequency critic updates (CleanRL uses 2)
        if self.update_counter % self.policy_frequency == 0:
            actor_loss, entropy, alpha, alpha_loss = self.update_actor_and_alpha(obs, step)
            actor_loss_info = actor_loss, entropy, alpha, alpha_loss
        else:
            # Return previous values (or zeros) when not updating actor
            actor_loss_info = torch.tensor(0.0, device=self._device), torch.tensor(0.0, device=self._device), self.alpha, None

        self.update_target_weights()
        return actor_loss_info, critic1_loss, critic2_loss

    def preproc_obs(self, obs):
        if isinstance(obs, dict):
            obs = obs['obs']
        obs = self.model.norm_obs(obs)

        return obs

    def cast_obs(self, obs):
        if isinstance(obs, torch.Tensor):
            self.is_tensor_obses = True
        elif isinstance(obs, np.ndarray):
            assert(self.observation_space.dtype != np.int8)
            if self.observation_space.dtype == np.uint8:
                obs = torch.ByteTensor(obs).to(self._device)
            else:
                obs = torch.FloatTensor(obs).to(self._device)

        return obs

    # TODO: move to common utils
    def obs_to_tensors(self, obs):
        obs_is_dict = isinstance(obs, dict)
        if obs_is_dict:
            upd_obs = {}
            for key, value in obs.items():
                upd_obs[key] = self._obs_to_tensors_internal(value)
        else:
            upd_obs = self.cast_obs(obs)
        if not obs_is_dict or 'obs' not in obs:
            upd_obs = {'obs': upd_obs}

        return upd_obs

    def _obs_to_tensors_internal(self, obs):
        if isinstance(obs, dict):
            upd_obs = {}
            for key, value in obs.items():
                upd_obs[key] = self._obs_to_tensors_internal(value)
        else:
            upd_obs = self.cast_obs(obs)

        return upd_obs

    def preprocess_actions(self, actions):
        if not self.is_tensor_obses:
            actions = actions.cpu().numpy()

        return actions

    def env_step(self, actions):
        actions = self.preprocess_actions(actions)
        obs, rewards, dones, infos = self.vec_env.step(actions)

        if self.is_tensor_obses:
            return self.obs_to_tensors(obs), rewards.to(self._device), dones.to(self._device), infos
        else:
            return torch.from_numpy(obs).to(self._device), torch.from_numpy(rewards).to(self._device, dtype=torch.float32), torch.from_numpy(dones).to(self._device), infos

    def env_reset(self):
        with torch.no_grad():
            obs = self.vec_env.reset()
        obs = self.obs_to_tensors(obs)

        return obs

    def act(self, obs, action_dim, sample=False):
        obs = self.preproc_obs(obs)
        dist = self.model.actor(obs)

        actions = dist.sample() if sample else dist.mean
        # Rescale from [-1, 1] to actual bounds
        actions = self.rescale_actions(actions)
        if actions.ndim != 2:
            raise ValueError(f"Actions tensor must be 2-dimensional, got shape {actions.shape}")

        return actions

    def extract_actor_stats(self, actor_losses, entropies, alphas, alpha_losses, actor_loss_info):
        actor_loss, entropy, alpha, alpha_loss = actor_loss_info

        actor_losses.append(actor_loss)
        entropies.append(entropy)
        if alpha_losses is not None:
            alphas.append(alpha)
            alpha_losses.append(alpha_loss)

    def clear_stats(self):
        self.game_rewards.clear()
        self.game_lengths.clear()
        self.mean_rewards = self.last_mean_rewards = -1000000000
        self.algo_observer.after_clear_stats()

    def play_steps(self, random_exploration=False):
        total_time_start = time.perf_counter()
        total_update_time = 0
        total_time = 0
        step_time = 0.0
        actor_losses = []
        entropies = []
        alphas = []
        alpha_losses = []
        critic1_losses = []
        critic2_losses = []

        obs = self.obs
        if isinstance(obs, dict):
            obs = self.obs['obs']

        for s in range(self.num_steps_per_episode):
            self.set_eval()
            if random_exploration:
                # Sample uniformly from the correct action bounds
                action = torch.rand((self.num_actors, *self.env_info["action_space"].shape), device=self._device)
                action = self.action_low + (self.action_high - self.action_low) * action
            else:
                with torch.no_grad():
                    action = self.act(obs.float(), self.env_info["action_space"].shape, sample=True)

            step_start = time.perf_counter()
            with torch.no_grad():
                next_obs, rewards, dones, infos = self.env_step(action)

            step_end = time.perf_counter()
            total_time += (step_end - step_start)
            step_time += (step_end - step_start)

            # Shape rewards BEFORE tracking (so logged rewards match training rewards)
            shaped_rewards = self.rewards_shaper(rewards)

            self.current_rewards.add_(shaped_rewards)
            self.current_lengths.add_(1)

            all_done_indices = dones.nonzero(as_tuple=False)
            done_indices = all_done_indices[::self.num_agents]
            self.game_rewards.update(self.current_rewards[done_indices])
            self.game_lengths.update(self.current_lengths[done_indices])

            not_dones = 1.0 - dones.float()

            self.algo_observer.process_infos(infos, done_indices)

            # Handle truncated episodes (timeouts) for auto-resetting envs.
            # Auto-resetting envs (envpool, Isaac Gym) replace next_obs with the reset obs,
            # so bootstrapping V(next_obs) would use the wrong state.
            # Fix: for truncated envs, use obs (pre-reset) as next_obs and mark done=False.
            # The critic will compute V(s_t) with fresh weights at training time.
            terminated = dones.clone()
            timeouts = None
            if self.value_bootstrap and 'time_outs' in infos:
                timeouts = self.cast_obs(infos['time_outs'])
                if timeouts.dtype != torch.bool:
                    timeouts = timeouts.bool()
                if timeouts.any():
                    terminated[timeouts] = 0.0
                else:
                    timeouts = None

            self.current_rewards.mul_(not_dones)
            self.current_lengths.mul_(not_dones)

            if isinstance(next_obs, dict):
                next_obs_processed = next_obs['obs']
                self.obs = next_obs
            else:
                next_obs_processed = next_obs
                self.obs = next_obs

            # For truncated envs, store pre-reset obs as next_obs in replay buffer
            # so the critic bootstraps V(s_t) with fresh weights at training time.
            # self.obs keeps the actual reset obs for the next step.
            replay_next_obs = next_obs_processed
            if timeouts is not None:
                replay_next_obs = next_obs_processed.clone()
                replay_next_obs[timeouts] = obs[timeouts]

            # Batch unsqueeze operations (rewards already shaped above)
            rewards_unsqueezed = shaped_rewards.unsqueeze(1)
            terminated_unsqueezed = terminated.unsqueeze(1)
            self.replay_buffer.add(obs, action, rewards_unsqueezed, replay_next_obs, terminated_unsqueezed)

            obs = next_obs_processed  # Use the already processed observation

            if not random_exploration:
                self.set_train()

                update_time_start = time.perf_counter()
                # Perform multiple gradient updates per environment step
                for _ in range(self.num_updates_per_step):
                    actor_loss_info, critic1_loss, critic2_loss = self.update(self.epoch_num)
                    self.extract_actor_stats(actor_losses, entropies, alphas, alpha_losses, actor_loss_info)
                    critic1_losses.append(critic1_loss)
                    critic2_losses.append(critic2_loss)
                update_time_end = time.perf_counter()
                update_time = update_time_end - update_time_start
            else:
                update_time = 0

            total_update_time += update_time

        total_time_end = time.perf_counter()
        total_time = total_time_end - total_time_start
        play_time = total_time - total_update_time

        return step_time, play_time, total_update_time, total_time, actor_losses, entropies, alphas, alpha_losses, critic1_losses, critic2_losses

    def train_epoch(self):
        random_exploration = self.epoch_num < self.num_warmup_steps
        return self.play_steps(random_exploration)

    def train(self):
        self.init_tensors()
        self.algo_observer.after_init(self)
        total_time = 0
        # rep_count = 0

        self.obs = self.env_reset()

        while True:
            self.epoch_num += 1
            step_time, play_time, update_time, epoch_total_time, actor_losses, entropies, alphas, alpha_losses, critic1_losses, critic2_losses = self.train_epoch()

            total_time += epoch_total_time

            curr_frames = self.num_frames_per_epoch
            self.frame += curr_frames

            fps_step = curr_frames / step_time
            fps_step_inference = curr_frames / play_time
            fps_total = curr_frames / epoch_total_time

            print_statistics(self.print_stats, curr_frames, step_time, play_time, epoch_total_time, 
                self.epoch_num, self.max_epochs, self.frame, self.max_frames)

            self.writer.add_scalar('performance/step_inference_rl_update_fps', fps_total, self.frame)
            self.writer.add_scalar('performance/step_inference_fps', fps_step_inference, self.frame)
            self.writer.add_scalar('performance/step_fps', fps_step, self.frame)
            self.writer.add_scalar('performance/rl_update_time', update_time, self.frame)
            self.writer.add_scalar('performance/step_inference_time', play_time, self.frame)
            self.writer.add_scalar('performance/step_time', step_time, self.frame)

            if self.epoch_num >= self.num_warmup_steps:
                self.writer.add_scalar('losses/a_loss', torch_ext.mean_list(actor_losses).item(), self.frame)
                self.writer.add_scalar('losses/c1_loss', torch_ext.mean_list(critic1_losses).item(), self.frame)
                self.writer.add_scalar('losses/c2_loss', torch_ext.mean_list(critic2_losses).item(), self.frame)
                self.writer.add_scalar('losses/entropy', torch_ext.mean_list(entropies).item(), self.frame)

                if alpha_losses[0] is not None:
                    self.writer.add_scalar('losses/alpha_loss', torch_ext.mean_list(alpha_losses).item(), self.frame)
                self.writer.add_scalar('info/alpha', torch_ext.mean_list(alphas).item(), self.frame)

            self.writer.add_scalar('info/epochs', self.epoch_num, self.frame)
            self.algo_observer.after_print_stats(self.frame, self.epoch_num, total_time)

            if self.game_rewards.current_size > 0:
                mean_rewards = self.game_rewards.get_mean()
                mean_lengths = self.game_lengths.get_mean()

                self.writer.add_scalar('rewards/step', mean_rewards, self.frame)
                self.writer.add_scalar('rewards/time', mean_rewards, total_time)
                self.writer.add_scalar('episode_lengths/step', mean_lengths, self.frame)
                self.writer.add_scalar('episode_lengths/time', mean_lengths, total_time)
                checkpoint_name = self.config['name'] + '_ep_' + str(self.epoch_num) + '_rew_' + str(mean_rewards)

                should_exit = False

                if self.save_freq > 0:
                    if self.epoch_num % self.save_freq == 0:
                        self.save(os.path.join(self.nn_dir, 'last_' + checkpoint_name))

                if mean_rewards > self.last_mean_rewards and self.epoch_num >= self.save_best_after:
                    print('saving next best rewards: ', mean_rewards)
                    self.last_mean_rewards = mean_rewards
                    self.save(os.path.join(self.nn_dir, self.config['name']))
                    if self.last_mean_rewards > self.config.get('score_to_win', float('inf')):
                        print('Maximum reward achieved. Network won!')
                        self.save(os.path.join(self.nn_dir, checkpoint_name))
                        should_exit = True

                if self.epoch_num >= self.max_epochs and self.max_epochs != -1:
                    if self.game_rewards.current_size == 0:
                        print('WARNING: Max epochs reached before any env terminated at least once')
                        mean_rewards = -np.inf

                    self.save(os.path.join(self.nn_dir, 'last_' + self.config['name'] + '_ep_' + str(self.epoch_num) \
                        + '_rew_' + str(mean_rewards).replace('[', '_').replace(']', '_')))
                    print('MAX EPOCHS NUM!')
                    should_exit = True

                if self.frame >= self.max_frames and self.max_frames != -1:
                    if self.game_rewards.current_size == 0:
                        print('WARNING: Max frames reached before any env terminated at least once')
                        mean_rewards = -np.inf

                    self.save(os.path.join(self.nn_dir, 'last_' + self.config['name'] + '_frame_' + str(self.frame) \
                        + '_rew_' + str(mean_rewards).replace('[', '_').replace(']', '_')))
                    print('MAX FRAMES NUM!')
                    should_exit = True

                update_time = 0

                if should_exit:
                    return self.last_mean_rewards, self.epoch_num
