constraints.py

import math
import cmath

tai = math.tau * 1j


class Constraint:
    def __init__(self, size=1.0, center=0j, rotation=0):
        self.center = center
        self.size = size
        self.rotation = rotation

    @property
    def size(self):
        return self._size

    @property
    def center(self):
        return self._center

    @property
    def rotation(self):
        return self._rotation

    @size.setter
    def size(self, val):
        size = complex(val)
        if size.imag < 0 or size.real <= 0:
            raise ValueError(
                f"size error: real part must be strictly positive, imaginary must be positive or 0. provided value: {size}"
            )
        if not size.imag:
            size = complex(size.real, size.real)
        self._size = size
        self._resize = size
        self._desize = complex(1 / size.real, 1 / size.imag)

    @center.setter
    def center(self, val):
        self._center = complex(val)

    @rotation.setter
    def rotation(self, val):
        self._rotation = val % math.tau
        self._rerot = cmath.exp(tai * self._rotation)
        self._derot = complex(self._rerot.real, -self._rerot.imag)

    def apply_hold(self, pos):
        return pos

    def apply_free(self, pos, vel, bounce):
        return self.apply_hold(pos), vel

    def _transform(self, pos):
        return self._transform_vel(pos - self._center)

    def _detransform(self, pos):
        return self._detransform_vel(pos) + self._center

    def _transform_vel(self, vel):
        # does not apply center
        vel *= self._derot
        vel = complex(self._desize.real * vel.real, self._desize.imag * vel.imag)
        return vel

    def _detransform_vel(self, vel):
        # does not apply center
        vel = complex(self._resize.real * vel.real, self._resize.imag * vel.imag)
        vel *= self._rerot
        return vel


class Rectangle(Constraint):
    @staticmethod
    def _complex2vec(c):
        return [c.real, c.imag]

    @staticmethod
    def _vec2complex(v):
        return complex(v[0], v[1])

    def apply_hold(self, pos):
        vecpos = self._complex2vec(self._transform(pos))
        clip = False
        for x in range(2):
            if 0.5 < vecpos[x]:
                clip = True
                vecpos[x] = 0.5
            elif -0.5 > vecpos[x]:
                clip = True
                vecpos[x] = -0.5
        if clip:
            pos = self._detransform(self._vec2complex(vecpos))
        return pos

    def apply_free(self, pos, vel, bounce):
        if not bounce:
            return self.apply_hold(pos), vel
        vecpos = self._complex2vec(self._transform(pos))
        vecvel = self._complex2vec(self._transform_vel(vel))

        clip = False
        for x in range(2):
            vecpos[x] += 0.5
            if not (0 <= vecpos[x] <= 1):
                clip = True
                hits, vecpos[x] = divmod(vecpos[x], 1)
                if hits % 2:
                    vecpos[x] = 1 - vecpos[x]
                    if bounce == 1:
                        vecvel[x] = -vecvel[x]
                if bounce != 1:
                    vecvel[x] *= (-bounce) ** abs(hits)
            vecpos[x] -= 0.5
        if clip:
            pos = self._detransform(self._vec2complex(vecpos))
            vel = self._detransform_vel(self._vec2complex(vecvel))
        return pos, vel


class ModuloRectangle(Rectangle):
    def apply_hold(self, pos):
        pos_tf = self._transform(pos) + complex(0.5, 0.5)
        clip = False
        if not (1 >= pos_tf.imag >= 0):
            pos_tf.imag %= 1
            clip = True
        if not (1 >= pos_tf.real >= 0):
            pos_tf.real %= 1
            clip = True
        if clip:
            pos = self._detransform(pos_tf - complex(0.5, 0.5))
        return pos

    def apply_free(self, pos, vel, bounce):
        return self.apply_hold(pos), vel


class Ellipse(Constraint):
    def apply_hold(self, pos):
        pos_tf = self._transform(pos)
        abs_sq = pos_tf.imag * pos_tf.imag + pos_tf.real * pos_tf.real
        if abs_sq > 1:
            pos_tf *= math.sqrt(1 / abs_sq)
            pos = self._detransform(pos_tf)
        return pos

    def apply_free(self, pos, vel, bounce):
        if not bounce:
            return self.apply_hold(pos), vel
        pos_tf = self._transform(pos)
        vel_tf = self._transform_vel(vel)
        timeout = 0
        clip = False
        while True:
            abs_sq = pos_tf.imag * pos_tf.imag + pos_tf.real * pos_tf.real
            if abs_sq <= 1:
                break
            clip = True
            if not vel_tf or timeout > 100:
                # arbirary iteration depth, if it's still outside of the circle (too fast?) we simply
                # clip it and set the velocity to 0. there may be more elegant solutions to this.
                pos_tf *= math.sqrt(1 / abs_sq)
                vel_tf = 0j
                break
            pos_tf, vel_tf = self._apply_free_inner(abs_sq, pos_tf, vel_tf, bounce)
            timeout += 1
        if clip:
            pos = self._detransform(pos_tf)
            vel = self._detransform_vel(vel_tf)
        return pos, vel

    def _apply_free_inner(self, abs_sq, pos, vel, bounce):
        # https://math.stackexchange.com/questions/228841/how-do-i-calculate-the-intersections-of-a-straight-line-and-a-circle
        pos_prev = pos - vel
        A = -vel.imag
        B = vel.real
        C = pos_prev.real * pos.imag - pos_prev.imag * pos.real
        a = A * A + B * B
        if abs(vel.real) > abs(vel.imag):
            xy_is_x = True
            b = 2 * A * C
            c = C * C - B * B
        else:
            xy_is_x = False
            b = 2 * B * C
            c = C * C - A * A
        root = b * b - 4 * a * c
        if root > 0:
            root = math.sqrt(root)
            xy = [(-b + root) / (2 * a), (-b - root) / (2 * a)]
            xy_pos = pos.real if xy_is_x else pos.imag
            if abs(xy[0] - xy_pos) > abs(xy[1] - xy_pos):
                xy = xy[1]
            else:
                xy = xy[0]
        elif root == 0:
            xy = -b / (2 * a)
        else:
            # velocity vector doesn't intersect with circle,
            # shouldn't happen.
            return pos * math.sqrt(1 / abs_sq), 0j
        if xy_is_x:
            gain = (xy - pos.real) / vel.real
            y = pos.imag + gain * vel.imag
            impact = xy + y * 1j
        else:
            gain = (xy - pos.imag) / vel.imag
            x = pos.real + gain * vel.real
            impact = x + xy * 1j
        # we now know the impact point.
        impact_vel = pos - impact
        outpact_vel = -impact_vel / impact
        outpact_vel -= 2j * outpact_vel.imag
        outpact_vel *= impact
        vel = abs(vel) * outpact_vel / abs(outpact_vel)
        vel *= bounce
        pos = impact + outpact_vel
        return pos, vel