project-euler/144_reflections_of_laser.py

#!/usr/bin/env python3

from math import sqrt, pi
import matplotlib.pyplot as plt
import numpy as np

class Ellipse:
    def __init__(self, a, b):
        self.a = float(a)
        self.b = float(b)

    def normal(self, vec):
        x, y = vec
        return norm([self.b**2 * x, self.a**2 * y])

def quadratic(a, b, c):
    D = b**2 - 4*a*c
    if D < 0:
        return False
    r1 = (-b + sqrt(D)) / (2*a)
    r2 = (-b - sqrt(D)) / (2*a)
    return r1, r2

def norm(v):
    l = sqrt(sum(x**2 for x in v))
    return np.array([x/l for x in v])

def intersections(ellipse, point, dir):
    m = dir[1] / dir[0]
    yint = point[1] - point[0]*m

    a = ellipse.a**2 * m**2 + ellipse.b**2
    b = 2 * ellipse.a**2 * m * yint
    c = ellipse.a**2 * (yint**2 - ellipse.b**2)

    x1, x2 = quadratic(a, b, c)
    y1, y2 = m*x1 + yint, m*x2 + yint

    return np.array([x1, y1]), np.array([x2, y2])

def reflect(ellipse, A, B):
    point = B
    dir = B - A
    n = ellipse.normal(point)
    ref = 2 * (n.dot(dir) * n - dir) + dir
    ps = intersections(ellipse, point, ref)
    return max(ps, key=lambda v: (B-v).dot(B-v))

def plot_ellipse(ellipse):
    t = np.linspace(0, 2*pi, 100)
    plt.plot(ellipse.a*np.cos(t), ellipse.b*np.sin(t))

ellipse = Ellipse(5, 10)
A = np.array([0.0, 10.1])
B = np.array([1.4, -9.6])
x = 0.01


plt.axis('equal')
plot_ellipse(ellipse)

count = 0

while not (-x <= B[0] <= x and B[1] > 0):
    plt.plot(*zip(A, B))
    A, B = B, reflect(ellipse, A, B)
    count += 1


plt.show()
print(count)