import os
import time
import numpy as np
import random
from copy import deepcopy
import torch
import torch._dynamo
torch._dynamo.config.cache_size_limit = 64

from rl_games.common import object_factory
from rl_games.common import tr_helpers

from rl_games.algos_torch import a2c_continuous
from rl_games.algos_torch import a2c_discrete
from rl_games.algos_torch import players
from rl_games.common.algo_observer import DefaultAlgoObserver
from rl_games.algos_torch import sac_agent

# Limit tensor printouts to 3 decimal places globally
torch.set_printoptions(precision=3, sci_mode=False)


def _restore(agent, args):
    if 'checkpoint' in args and args['checkpoint'] is not None and args['checkpoint'] !='':
        if args['train'] and args.get('load_critic_only', False):
            if not getattr(agent, 'has_central_value', False):
                raise ValueError('Loading critic only works only for asymmetric actor critic')
            agent.restore_central_value_function(args['checkpoint'])
            return
        agent.restore(args['checkpoint'])

def _override_sigma(agent, args):
    if 'sigma' in args and args['sigma'] is not None:
        net = agent.model.a2c_network
        if hasattr(net, 'sigma') and hasattr(net, 'fixed_sigma'):
            if net.fixed_sigma:
                with torch.no_grad():
                    net.sigma.fill_(float(args['sigma']))
            else:
                print('Cannot set new sigma because fixed_sigma is False')


class Runner:
    """Runs training/inference (playing) procedures as per a given configuration.

    The Runner class provides a high-level API for instantiating agents for either training or playing
    with an RL algorithm. It further logs training metrics.

    """

    def __init__(self, algo_observer=None):
        """Initialise the runner instance with algorithms and observers.

        Initialises runners and players for all algorithms available in the library using `rl_games.common.object_factory.ObjectFactory`

        Args:
            algo_observer (:obj:`rl_games.common.algo_observer.AlgoObserver`, optional): Algorithm observer that logs training metrics.
                Defaults to `rl_games.common.algo_observer.DefaultAlgoObserver`

        """

        self.algo_factory = object_factory.ObjectFactory()
        self.algo_factory.register_builder('a2c_continuous', lambda **kwargs : a2c_continuous.A2CAgent(**kwargs))
        self.algo_factory.register_builder('a2c_discrete', lambda **kwargs : a2c_discrete.DiscreteA2CAgent(**kwargs)) 
        self.algo_factory.register_builder('sac', lambda **kwargs: sac_agent.SACAgent(**kwargs))
        #self.algo_factory.register_builder('dqn', lambda **kwargs : dqnagent.DQNAgent(**kwargs))

        self.player_factory = object_factory.ObjectFactory()
        self.player_factory.register_builder('a2c_continuous', lambda **kwargs : players.PpoPlayerContinuous(**kwargs))
        self.player_factory.register_builder('a2c_discrete', lambda **kwargs : players.PpoPlayerDiscrete(**kwargs))
        self.player_factory.register_builder('sac', lambda **kwargs : players.SACPlayer(**kwargs))
        #self.player_factory.register_builder('dqn', lambda **kwargs : players.DQNPlayer(**kwargs))

        self.algo_observer = algo_observer if algo_observer else DefaultAlgoObserver()

        # Performance optimizations
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.backends.cudnn.benchmark = True

        # Enable TensorFloat32 (TF32) for faster matrix multiplications on NVIDIA GPUs
        torch.set_float32_matmul_precision('high')

    def reset(self):
        pass

    def load_config(self, params):
        """Loads passed config params.

        Args:
            params (:obj:`dict`): Parameters passed in as a dict obtained from a yaml file or some other config format.

        """

        self.seed = params.get('seed', None)
        if self.seed is None:
            self.seed = int(time.time())

        self.local_rank = 0
        self.global_rank = 0
        self.world_size = 1

        if params["config"].get('multi_gpu', False):
            # local rank of the GPU in a node
            self.local_rank = int(os.getenv("LOCAL_RANK", "0"))
            # global rank of the GPU
            self.global_rank = int(os.getenv("RANK", "0"))
            # total number of GPUs across all nodes
            self.world_size = int(os.getenv("WORLD_SIZE", "1"))

            # set different random seed for each GPU
            self.seed += self.global_rank

            print(f"global_rank = {self.global_rank} local_rank = {self.local_rank} world_size = {self.world_size}")

        print(f"self.seed = {self.seed}")

        self.algo_params = params['algo']
        self.algo_name = self.algo_params['name']
        self.exp_config = None

        if self.seed:
            torch.manual_seed(self.seed)
            torch.cuda.manual_seed_all(self.seed)
            np.random.seed(self.seed)
            random.seed(self.seed)

            # deal with environment specific seed if applicable
            if 'env_config' in params['config']:
                if not 'seed' in params['config']['env_config']:
                    params['config']['env_config']['seed'] = self.seed
                else:
                    if params["config"].get('multi_gpu', False):
                        params['config']['env_config']['seed'] += self.seed

        config = params['config']
        config['reward_shaper'] = tr_helpers.DefaultRewardsShaper(**config['reward_shaper'])
        if 'features' not in config:
            config['features'] = {}
        config['features']['observer'] = self.algo_observer
        self.params = params

        # Set CPU thread count for intra-op parallelism (OpenMP/MKL).
        # Default (auto): min(4, cores_per_rank) — avoids oversubscription in multi-GPU.
        # Set torch_threads: 0 in config to disable (use PyTorch default).
        num_threads = params.get('torch_threads', 'auto')
        if num_threads == 'auto':
            cpu = os.cpu_count() or 1
            per_rank = cpu // max(1, self.world_size)
            num_threads = max(1, min(4, per_rank))
        if num_threads:
            torch.set_num_threads(num_threads)

    def load(self, yaml_config):
        config = deepcopy(yaml_config)
        self.default_config = deepcopy(config['params'])
        self.load_config(params=self.default_config)

    def run_train(self, args):
        """Run the training procedure from the algorithm passed in.

        Args:
            args (:obj:`dict`): Train specific args passed in as a dict obtained from a yaml file or some other config format.

        """
        print('Started to train')

        # Add profiling support
        profiler = None
        if args.get('profile', False):
            profiler = torch.profiler.profile(
                activities=[
                    torch.profiler.ProfilerActivity.CPU,
                    torch.profiler.ProfilerActivity.CUDA,
                ],
                record_shapes=True,
                profile_memory=True,
                with_stack=True
            )
            profiler.__enter__()

        agent = self.algo_factory.create(self.algo_name, base_name='run', params=self.params)

        # Restore weights (if any) BEFORE compiling the model.  Compiling first
        # wraps the model in an `OptimizedModule`, which changes parameter
        # names (adds the `_orig_mod.` prefix) and breaks `load_state_dict`
        # when loading checkpoints that were saved from an *un‑compiled*
        # model.

        _restore(agent, args)
        _override_sigma(agent, args)

        # Now compile the (already restored) model. Doing it after the restore
        # keeps parameter names consistent with the checkpoint.

        # torch.compile modes:
        # - "default": kernel fusion + operator optimization (recommended default)
        # - "reduce-overhead": adds CUDA graph capture/replay on top of default
        # - "max-autotune": best runtime, longest compilation
        # WARNING: 'reduce-overhead' and 'max-autotune' use CUDA graphs which are
        # incompatible with RNN/LSTM models — graph output buffers get overwritten
        # across sequential rollout calls and LSTM internal allocations violate
        # CUDA graph memory pool constraints during backward pass.

        # Enable torch.compile for performance if requested
        compile_config = self.params.get('config', {}).get('torch_compile', True)

        if compile_config is False:
            print("torch.compile: Disabled")
        else:
            # Parse configuration
            if isinstance(compile_config, bool):
                actor_mode = "default"
                critic_mode = "default"
            elif isinstance(compile_config, str):
                actor_mode = compile_config
                critic_mode = compile_config
            elif isinstance(compile_config, dict):
                actor_mode = compile_config.get('mode', 'default')
                critic_mode = compile_config.get('critic_mode', actor_mode)
            else:
                print(f"Warning: Invalid torch_compile config {compile_config}, using default")
                actor_mode = "default"
                critic_mode = "default"

            print(f"torch.compile: Enabled for actor with mode='{actor_mode}'")
            agent.model = torch.compile(agent.model, mode=actor_mode)

            if hasattr(agent, 'central_value_net') and agent.central_value_net is not None:
                print(f"torch.compile: Enabled for central value critic with mode='{critic_mode}'")
                agent.central_value_net.model = torch.compile(
                    agent.central_value_net.model,
                    mode=critic_mode
                )
        agent.train()

        # Save profiling results if enabled
        if profiler is not None:
            profiler.__exit__(None, None, None)
            print(profiler.key_averages().table(sort_by="cuda_time_total", row_limit=10))
            profiler.export_chrome_trace("profile_trace.json")

    def run_play(self, args):
        """Run the inference procedure from the algorithm passed in.

        Args:
            args (:obj:`dict`): Playing specific args passed in as a dict obtained from a yaml file or some other config format.

        """
        print('Started to play')
        player = self.create_player()
        _restore(player, args)
        _override_sigma(player, args)
        player.run()

    def create_player(self):
        return self.player_factory.create(self.algo_name, params=self.params)

    def reset(self):
        pass

    def run(self, args):
        """Run either train/play depending on the args.

        Args:
            args (:obj:`dict`):  Args passed in as a dict obtained from a yaml file or some other config format.

        """
        if args['train']:
            self.run_train(args)
        elif args['play']:
            self.run_play(args)
        else:
            self.run_train(args)
