punyverse/punyverse/entity.py

import random
from math import sqrt, pi

from pyglet.gl import *
# noinspection PyUnresolvedReferences
from six.moves import range

from punyverse.glgeom import compile, glRestore, belt, Disk, OrbitVBO, Matrix4f, SimpleSphere, TangentSphere
from punyverse.model import load_model, WavefrontVBO
from punyverse.orbit import KeplerOrbit
from punyverse.texture import get_best_texture, load_clouds
from punyverse.utils import cached_property

G = 6.67384e-11  # Gravitation Constant


class Entity(object):
    background = False

    def __init__(self, world, name, location, rotation=(0, 0, 0), direction=(0, 0, 0)):
        self.world = world
        self.name = name
        self.location = location
        self.rotation = rotation
        self.direction = direction

    @cached_property
    def model_matrix(self):
        return Matrix4f.from_angles(self.location, self.rotation)

    @cached_property
    def mv_matrix(self):
        return self.world.view_matrix() * self.model_matrix

    @cached_property
    def mvp_matrix(self):
        return self.world.vp_matrix * self.model_matrix

    def update(self):
        self.model_matrix = None
        self.mv_matrix = None
        self.mvp_matrix = None
        x, y, z = self.location
        dx, dy, dz = self.direction
        self.location = x + dx, y + dy, z + dz

    def collides(self, x, y, z):
        return False

    def draw(self, options):
        raise NotImplementedError()


class Asteroid(Entity):
    def __init__(self, world, model, location, direction):
        super(Asteroid, self).__init__(world, 'Asteroid', location, direction=direction)
        self.model = model

    def update(self):
        super(Asteroid, self).update()
        rx, ry, rz = self.rotation
        # Increment all axis to 'spin' 
        self.rotation = rx + 1, ry + 1, rz + 1

    def draw(self, options):
        glLoadMatrixf(self.mv_matrix)
        self.model.draw()


class AsteroidManager(object):
    def __init__(self, world):
        self.world = world
        self.asteroids = []

    def __bool__(self):
        return bool(self.asteroids)
    __nonzero__ = __bool__

    def load(self, file):
        self.asteroids.append(WavefrontVBO(load_model(file), 5, 5, 5))

    def new(self, location, direction):
        return Asteroid(self.world, random.choice(self.asteroids), location, direction)


class Belt(Entity):
    def __init__(self, name, world, info):
        x = world.evaluate(info.get('x', 0))
        y = world.evaluate(info.get('y', 0))
        z = world.evaluate(info.get('z', 0))
        radius = world.evaluate(info.get('radius', 0))
        cross = world.evaluate(info.get('cross', 0))
        count = int(world.evaluate(info.get('count', 0)))
        scale = info.get('scale', 1)
        longitude = info.get('longitude', 0)
        inclination = info.get('inclination', 0)
        argument = info.get('argument', 0)
        rotation = info.get('period', 31536000)
        models = info['model']
        if not isinstance(models, list):
            models = [models]

        objects = [WavefrontVBO(load_model(model), info.get('sx', scale), info.get('sy', scale),
                                info.get('sz', scale)) for model in models]

        self.belt_id = compile(belt, radius, cross, objects, count)
        self.rotation_angle = 360.0 / rotation if rotation else 0

        super(Belt, self).__init__(world, name, (x, y, z), (inclination, longitude, argument))

    def update(self):
        super(Belt, self).update()
        pitch, yaw, roll = self.rotation
        self.rotation = pitch, self.world.tick * self.rotation_angle % 360, roll

    def draw(self, options):
        glLoadMatrixf(self.mv_matrix)
        glCallList(self.belt_id)


class Sky(Entity):
    background = True

    def __init__(self, world, info):
        pitch = world.evaluate(info.get('pitch', 0))
        yaw = world.evaluate(info.get('yaw', 0))
        roll = world.evaluate(info.get('roll', 0))

        super(Sky, self).__init__(world, 'Sky', (0, 0, 0), (pitch, yaw, roll))

        self.texture = get_best_texture(info['texture'])
        division = info.get('division', 30)
        self.sphere = SimpleSphere(division, division)

    def draw(self, options):
        cam = self.world.cam
        shader = self.world.activate_shader('sky')
        shader.uniform_mat4('u_mvpMatrix', self.world.projection_matrix() *
                            Matrix4f.from_angles(rotation=(cam.pitch, cam.yaw, cam.roll)) *
                            Matrix4f.from_angles(rotation=self.rotation))

        glActiveTexture(GL_TEXTURE0)
        glBindTexture(GL_TEXTURE_2D, self.texture)
        shader.uniform_texture('u_skysphere', 0)

        glBindBuffer(GL_ARRAY_BUFFER, self.sphere.vbo)
        shader.vertex_attribute('a_direction', self.sphere.direction_size, self.sphere.type, GL_FALSE,
                                self.sphere.stride, self.sphere.direction_offset)
        shader.vertex_attribute('a_uv', self.sphere.uv_size, self.sphere.type, GL_FALSE,
                                self.sphere.stride, self.sphere.uv_offset)

        glDrawArrays(GL_TRIANGLE_STRIP, 0, self.sphere.vertex_count)

        shader.deactivate_attributes()
        glBindBuffer(GL_ARRAY_BUFFER, 0)
        self.world.activate_shader(None)


class Body(Entity):
    def __init__(self, name, world, info, parent=None):
        self.parent = parent
        self.satellites = []

        x = world.evaluate(info.get('x', 0))
        y = world.evaluate(info.get('y', 0))
        z = world.evaluate(info.get('z', 0))
        pitch = world.evaluate(info.get('pitch', 0))
        yaw = world.evaluate(info.get('yaw', 0))
        roll = world.evaluate(info.get('roll', 0))
        rotation = world.evaluate(info.get('rotation', 86400))

        self.mass = info.get('mass')

        orbit_distance = float(world.evaluate(info.get('orbit_distance', world.au)))
        self.orbit_show = orbit_distance * 1.25
        self.orbit_blend = orbit_distance / 4
        self.orbit_opaque = orbit_distance

        super(Body, self).__init__(world, name, (x, y, z), (pitch, yaw, roll))
        self.initial_roll = roll

        self.orbit = None
        self.orbit_speed = None

        if parent:
            # Semi-major axis when actually displayed in virtual space
            distance = world.evaluate(info.get('distance', 100))
            # Semi-major axis used to calculate orbital speed
            sma = world.evaluate(info.get('sma', distance))

            if hasattr(parent, 'mass') and parent.mass is not None:
                period = 2 * pi * sqrt((sma * 1000) ** 3 / (G * parent.mass))
                self.orbit_speed = 360.0 / period
                if not rotation:  # Rotation = 0 assumes tidal lock
                    rotation = period
            else:
                self.orbit_speed = info.get('orbit_speed', 1)

            self.orbit = KeplerOrbit(distance / world.length, info.get('eccentricity', 0), info.get('inclination', 0),
                                     info.get('longitude', 0), info.get('argument', 0))

        self.rotation_angle = 360.0 / rotation if rotation else 0

        # Orbit calculation
        self.orbit_vbo = None
        self.orbit_cache = None

    @cached_property
    def orbit_matrix(self):
        return self.world.view_matrix() * Matrix4f.from_angles(self.location)

    def update(self):
        super(Body, self).update()

        if self.rotation_angle:
            pitch, yaw, roll = self.rotation
            roll = (self.initial_roll + self.world.tick * self.rotation_angle) % 360
            self.rotation = pitch, yaw, roll

        if self.orbit:
            px, py, pz = self.parent.location
            x, z, y = self.orbit.orbit(self.world.tick * self.orbit_speed % 360)
            self.location = (x + px, y + py, z + pz)
            self.orbit_matrix = None

        for satellite in self.satellites:
            satellite.update()

    def get_orbit(self):
        if not self.orbit:
            return

        # Cache key is the three orbital plane parameters and eccentricity
        cache = (self.orbit.eccentricity, self.orbit.longitude, self.orbit.inclination, self.orbit.argument)
        if self.orbit_cache == cache:
            return self.orbit_vbo

        if self.orbit_vbo is not None:
            self.orbit_vbo.close()

        self.orbit_vbo = OrbitVBO(self.orbit)
        self.orbit_cache = cache
        return self.orbit_vbo

    def _draw_orbits(self, distance):
        with glRestore(GL_ENABLE_BIT | GL_LINE_BIT | GL_CURRENT_BIT):
            glLoadMatrixf(self.parent.orbit_matrix)

            glDisable(GL_LIGHTING)
            solid = distance < self.parent.orbit_opaque
            glColor4f(1, 1, 1, 1 if solid else (1 - (distance - self.parent.orbit_opaque) / self.parent.orbit_blend))
            if not solid:
                glEnable(GL_BLEND)
            glLineWidth(1)
            self.get_orbit().draw()

    def draw(self, options):
        self._draw(options)

        if options.orbit and self.orbit:
            dist = self.world.cam.distance(*self.parent.location)
            if dist < self.parent.orbit_show:
                self._draw_orbits(dist)

        for satellite in self.satellites:
            satellite.draw(options)

    def _draw(self, options):
        raise NotImplementedError()

    def collides(self, x, y, z):
        return self._collides(x, y, z) or any(satellite.collides(x, y, z) for satellite in self.satellites)

    def _collides(self, x, y, z):
        return False


class SphericalBody(Body):
    _sphere_cache = {}

    @classmethod
    def _get_sphere(cls, division, tangent=True):
        if (division, tangent) in cls._sphere_cache:
            return cls._sphere_cache[division, tangent]
        cls._sphere_cache[division, tangent] = sphere = \
            (TangentSphere if tangent else SimpleSphere)(division, division)
        return sphere

    def __init__(self, name, world, info, parent=None):
        super(SphericalBody, self).__init__(name, world, info, parent)

        self.radius = world.evaluate(info.get('radius', world.length)) / world.length
        division = info.get('division', max(min(int(self.radius / 8), 60), 10))

        self.light_source = info.get('light_source', False)
        self.shininess = info.get('shininess', 0)
        self.type = info.get('type', 'planet')

        self.texture = get_best_texture(info['texture'])
        self.normal_texture = None
        self.specular_texture = None
        self.emission_texture = None
        self.sphere = self._get_sphere(division, tangent=self.type == 'planet')

        self.atmosphere = None
        self.clouds = None
        self.ring = 0

        if 'normal_map' in info:
            self.normal_texture = get_best_texture(info['normal_map'])

        if 'specular_map' in info:
            self.specular_texture = get_best_texture(info['specular_map'])

        if 'emission_map' in info:
            self.emission_texture = get_best_texture(info['emission_map'])

        if 'atmosphere' in info:
            atmosphere_data = info['atmosphere']
            atm_size = world.evaluate(atmosphere_data.get('diffuse_size', None))
            atm_texture = atmosphere_data.get('diffuse_texture', None)
            cloud_texture = atmosphere_data.get('cloud_texture', None)
            if cloud_texture is not None:
                self.cloud_transparency = get_best_texture(cloud_texture, loader=load_clouds)
                self.cloud_radius = self.radius + 2
                self.clouds = self._get_sphere(division, tangent=False)

            if atm_texture is not None:
                self.atm_texture = get_best_texture(atm_texture, clamp=True)
                self.atmosphere = Disk(self.radius, self.radius + atm_size, 30)

        if 'ring' in info:
            distance = world.evaluate(info['ring'].get('distance', self.radius * 1.2))
            size = world.evaluate(info['ring'].get('size', self.radius / 2))

            pitch, yaw, roll = self.rotation
            pitch = world.evaluate(info['ring'].get('pitch', pitch))
            yaw = world.evaluate(info['ring'].get('yaw', yaw))
            roll = world.evaluate(info['ring'].get('roll', roll))
            self.ring_rotation = pitch, yaw, roll

            self.ring_texture = get_best_texture(info['ring'].get('texture', None), clamp=True)
            self.ring = Disk(distance, distance + size, 30)

    def _draw_planet(self):
        shader = self.world.activate_shader('planet')
        shader.uniform_float('u_radius', self.radius)
        shader.uniform_mat4('u_modelMatrix', self.model_matrix)
        shader.uniform_mat4('u_mvMatrix', self.mv_matrix)
        shader.uniform_mat4('u_mvpMatrix', self.mvp_matrix)

        glActiveTexture(GL_TEXTURE0)
        glBindTexture(GL_TEXTURE_2D, self.texture)
        shader.uniform_texture('u_planet.diffuseMap', 0)

        shader.uniform_bool('u_planet.hasNormal', self.normal_texture)
        if self.normal_texture:
            glActiveTexture(GL_TEXTURE1)
            glBindTexture(GL_TEXTURE_2D, self.normal_texture)
            shader.uniform_texture('u_planet.normalMap', 1)

        shader.uniform_bool('u_planet.hasSpecular', self.specular_texture)
        if self.specular_texture:
            glActiveTexture(GL_TEXTURE2)
            glBindTexture(GL_TEXTURE_2D, self.specular_texture)
            shader.uniform_texture('u_planet.specularMap', 2)
            shader.uniform_vec3('u_planet.specular', 1, 1, 1)
            shader.uniform_float('u_planet.shininess', 10)
        else:
            shader.uniform_vec3('u_planet.specular', 0, 0, 0)
            shader.uniform_float('u_planet.shininess', 0)

        shader.uniform_bool('u_planet.hasEmission', self.emission_texture)
        if self.emission_texture:
            glActiveTexture(GL_TEXTURE3)
            glBindTexture(GL_TEXTURE_2D, self.emission_texture)
            shader.uniform_texture('u_planet.emissionMap', 3)
            shader.uniform_vec3('u_planet.ambient', 0, 0, 0)
            shader.uniform_vec3('u_planet.emission', 1, 1, 1)
        else:
            shader.uniform_vec3('u_planet.ambient', 1, 1, 1)
            shader.uniform_vec3('u_planet.emission', 0, 0, 0)

        shader.uniform_vec3('u_planet.diffuse', 1, 1, 1)

        glBindBuffer(GL_ARRAY_BUFFER, self.sphere.vbo)
        shader.vertex_attribute('a_normal', self.sphere.direction_size, self.sphere.type, GL_FALSE,
                                self.sphere.stride, self.sphere.direction_offset)
        shader.vertex_attribute('a_tangent', self.sphere.tangent_size, self.sphere.type, GL_FALSE,
                                self.sphere.stride, self.sphere.tangent_offset)
        shader.vertex_attribute('a_uv', self.sphere.uv_size, self.sphere.type, GL_FALSE,
                                self.sphere.stride, self.sphere.uv_offset)

        glDrawArrays(GL_TRIANGLE_STRIP, 0, self.sphere.vertex_count)

        shader.deactivate_attributes()
        glBindBuffer(GL_ARRAY_BUFFER, 0)
        self.world.activate_shader(None)
        glActiveTexture(GL_TEXTURE0)

    def _draw_star(self):
        shader = self.world.activate_shader('star')
        shader.uniform_float('u_radius', self.radius)
        shader.uniform_mat4('u_mvpMatrix', self.mvp_matrix)

        glActiveTexture(GL_TEXTURE0)
        glBindTexture(GL_TEXTURE_2D, self.texture)
        shader.uniform_texture('u_emission', 0)

        glBindBuffer(GL_ARRAY_BUFFER, self.sphere.vbo)
        shader.vertex_attribute('a_normal', self.sphere.direction_size, self.sphere.type, GL_FALSE,
                                self.sphere.stride, self.sphere.direction_offset)
        shader.vertex_attribute('a_uv', self.sphere.uv_size, self.sphere.type, GL_FALSE,
                                self.sphere.stride, self.sphere.uv_offset)

        glDrawArrays(GL_TRIANGLE_STRIP, 0, self.sphere.vertex_count)

        shader.deactivate_attributes()
        glBindBuffer(GL_ARRAY_BUFFER, 0)
        self.world.activate_shader(None)

    def _draw_sphere(self):
        if self.type == 'planet':
            self._draw_planet()
        elif self.type == 'star':
            self._draw_star()
        else:
            raise ValueError('Invalid type: %s' % self.type)

    def _draw_atmosphere(self):
        with glRestore(GL_ENABLE_BIT | GL_CURRENT_BIT | GL_TEXTURE_BIT):
            mv = self.mv_matrix.matrix
            matrix = Matrix4f([1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, mv[12], mv[13], mv[14], 1])
            glLoadMatrixf(matrix)

            glDisable(GL_LIGHTING)
            glEnable(GL_TEXTURE_2D)
            glEnable(GL_BLEND)
            glDisable(GL_CULL_FACE)

            glBindTexture(GL_TEXTURE_2D, self.atm_texture)
            self.atmosphere.draw()

    def _draw_clouds(self):
        glEnable(GL_BLEND)
        shader = self.world.activate_shader('clouds')
        shader.uniform_float('u_radius', self.cloud_radius)
        shader.uniform_mat4('u_modelMatrix', self.model_matrix)
        shader.uniform_mat4('u_mvpMatrix', self.mvp_matrix)

        glActiveTexture(GL_TEXTURE0)
        glBindTexture(GL_TEXTURE_2D, self.cloud_transparency)
        shader.uniform_texture('u_transparency', 0)
        shader.uniform_vec3('u_diffuse', 1, 1, 1)
        shader.uniform_vec3('u_ambient', 0.1, 0.1, 0.1)

        glBindBuffer(GL_ARRAY_BUFFER, self.clouds.vbo)
        shader.vertex_attribute('a_normal', self.clouds.direction_size, self.clouds.type, GL_FALSE,
                                self.clouds.stride, self.clouds.direction_offset)
        shader.vertex_attribute('a_uv', self.clouds.uv_size, self.clouds.type, GL_FALSE,
                                self.clouds.stride, self.clouds.uv_offset)

        glDrawArrays(GL_TRIANGLE_STRIP, 0, self.clouds.vertex_count)

        shader.deactivate_attributes()
        glBindBuffer(GL_ARRAY_BUFFER, 0)
        self.world.activate_shader(None)
        glDisable(GL_BLEND)

    def _draw_rings(self):
        glEnable(GL_BLEND)
        shader = self.world.activate_shader('ring')
        shader.uniform_mat4('u_modelMatrix', self.model_matrix)
        shader.uniform_mat4('u_mvpMatrix', self.mvp_matrix)
        shader.uniform_vec3('u_planet', *self.location)
        shader.uniform_vec3('u_sun', 0, 0, 0)
        shader.uniform_float('u_planetRadius', self.radius)
        shader.uniform_float('u_ambient', 0.1)

        glActiveTexture(GL_TEXTURE0)
        glBindTexture(GL_TEXTURE_2D, self.ring_texture)
        shader.uniform_texture('u_texture', 0)

        glBindBuffer(GL_ARRAY_BUFFER, self.ring.vbo)
        shader.vertex_attribute('a_position', self.ring.position_size, self.ring.type, GL_FALSE,
                                self.ring.stride, self.ring.position_offset)
        shader.vertex_attribute('a_u', self.ring.u_size, self.ring.type, GL_FALSE,
                                self.ring.stride, self.ring.u_offset)

        glDrawArrays(GL_TRIANGLE_STRIP, 0, self.ring.vertex_count)

        shader.deactivate_attributes()
        glBindBuffer(GL_ARRAY_BUFFER, 0)
        self.world.activate_shader(None)
        glDisable(GL_BLEND)

    def _draw(self, options):
        self._draw_sphere()

        if options.atmosphere and self.atmosphere:
            self._draw_atmosphere()

        if options.cloud and self.clouds:
            self._draw_clouds()

        if self.ring:
            self._draw_rings()

    def _collides(self, x, y, z):
        ox, oy, oz = self.location
        dx, dy, dz = x - ox, y - oy, z - oz
        distance = sqrt(dx * dx + dy * dy + dz * dz)
        return distance <= self.radius


class ModelBody(Body):
    def __init__(self, name, world, info, parent=None):
        super(ModelBody, self).__init__(name, world, info, parent)

        scale = info.get('scale', 1)
        self.vbo = WavefrontVBO(load_model(info['model']), info.get('sx', scale), info.get('sy', scale),
                                info.get('sz', scale))

    def _draw(self, options):
        glLoadMatrixf(self.mv_matrix)
        self.vbo.draw()