# Copyright 2017 The dm_control Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ============================================================================

"""Base class for OpenGL context handlers.

`ContextBase` defines a common API that OpenGL rendering contexts should conform

to. In addition, it provides a `make_current` context manager that:

1. Makes this OpenGL context current within the appropriate rendering thread.

2. Yields an object exposing a `call` method that can be used to execute OpenGL

calls within the rendering thread.

See the docstring for `dm_control.utils.render_executor` for further details

regarding rendering threads.

"""

import abc

import atexit

import collections

import contextlib

import sys

import weakref

from absl import logging

from dm_control._render import executor

import numpy as np

_CURRENT_CONTEXT_FOR_THREAD = collections.defaultdict(lambda: None)

_CURRENT_THREAD_FOR_CONTEXT = collections.defaultdict(lambda: None)

class ContextBase(metaclass=abc.ABCMeta):

"""Base class for managing OpenGL contexts."""

def __init__(self,

max_width,

max_height,

render_executor_class=executor.RenderExecutor):

"""Initializes this context."""

logging.debug('Using render executor class: %s',

render_executor_class.__name__)

self._render_executor = render_executor_class()

self._refcount = 0

self_weakref = weakref.ref(self)

def _free_at_exit():

if self_weakref():

self_weakref()._free_unconditionally() # pylint: disable=protected-access

atexit.register(_free_at_exit)

with self._render_executor.execution_context() as ctx:

ctx.call(self._platform_init, max_width, max_height)

self._patients = []

def keep_alive(self, obj):

self._patients.append(obj)

def dont_keep_alive(self, obj):

try:

self._patients.remove(obj)

except ValueError:

pass

def increment_refcount(self):

self._refcount += 1

def decrement_refcount(self):

self._refcount -= 1

@property

def terminated(self):

return self._render_executor.terminated

@property

def thread(self):

return self._render_executor.thread

def _free_on_executor_thread(self): # pylint: disable=missing-function-docstring

current_ctx = _CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread]

if current_ctx is not None:

del _CURRENT_THREAD_FOR_CONTEXT[current_ctx]

del _CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread]

self._platform_make_current()

try:

dummy = []

while self._patients:

patient = self._patients.pop()

assert sys.getrefcount(patient) == sys.getrefcount(dummy)

if hasattr(patient, 'free'):

patient.free()

del patient

finally:

self._platform_free()

def free(self):

"""Frees resources associated with this context if its refcount is zero."""

if self._refcount == 0:

self._free_unconditionally()

def _free_unconditionally(self):

self._render_executor.terminate(self._free_on_executor_thread)

def __del__(self):

self._free_unconditionally()

@contextlib.contextmanager

def make_current(self):

"""Context manager that makes this Renderer's OpenGL context current.

Yields:

An object that exposes a `call` method that can be used to call a

function on the dedicated rendering thread.

Raises:

RuntimeError: If this context is already current on another thread.

"""

with self._render_executor.execution_context() as ctx:

if _CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread] != id(self):

if _CURRENT_THREAD_FOR_CONTEXT[id(self)]:

raise RuntimeError(

'Cannot make context {!r} current on thread {!r}: '

'this context is already current on another thread {!r}.'

.format(self, self._render_executor.thread,

_CURRENT_THREAD_FOR_CONTEXT[id(self)]))

else:

current_context = (

_CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread])

if current_context:

del _CURRENT_THREAD_FOR_CONTEXT[current_context]

_CURRENT_THREAD_FOR_CONTEXT[id(self)] = self._render_executor.thread

_CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread] = id(self)

ctx.call(self._platform_make_current)

yield ctx

def to_pixels(self, buffer):

"""Converts the buffer to pixels."""

return np.flipud(buffer)

@abc.abstractmethod

def _platform_init(self, max_width, max_height):

"""Performs an implementation-specific context initialization."""

@abc.abstractmethod

def _platform_make_current(self):

"""Make the OpenGL context current on the executing thread."""

@abc.abstractmethod

def _platform_free(self):

"""Performs an implementation-specific context cleanup."""