import math
from dataclasses import dataclass
import matplotlib.pyplot as plt
import numpy as np


@dataclass
class P:
    x: float
    y: float


@dataclass
class Line:
    m: float
    n: float

    def plot(self, x_min: float, x_max: float, color='r'):
        x = np.linspace(x_min, x_max, 10**5)
        y = self.m * x + self.n
        plt.plot(x, y, color, label='') 


def line_from_two_points(p1: P, p2: P) -> Line:
    m = (p2.y - p1.y) / (p2.x - p1.x)
    n = ((p1.y * p2.x) - (p2.y * p1.x)) / (p2.x - p1.x)
    return Line(m, n)


def line_from_one_point(p: P, m: float) -> Line:
    n = -m * p.x + p.y
    return Line(m, n)


def reflected_slope(m1, m2):
    """ Math """
    alpha1 = math.atan(m1)
    alpha2 = -math.atan(1/m2)
    reflected_angle = alpha1 - 2*alpha2
    m3 = math.tan(reflected_angle)
    return -m3


def get_next_point(start_point: P, line: Line) -> P:
    # With y = m*x + n into 4x**2 + y**2 = 10 get
    # quadratic equation x**2 + 2mn * x + (n**2 - 10) = 0.

    # Solve quadratic equation
    a = 4 + line.m ** 2
    b = 2 * line.m * line.n
    c = line.n**2 - 100
    s = math.sqrt(b**2 - 4*a*c)

    # Pick correct x 
    EPSILON = 10e-9
    x1 = (-b + s) / (2*a)
    x2 = (-b - s) / (2*a)
    x = x2 if abs(x1 - start_point.x) < EPSILON else x1

    # Calculate corresponding y and return point
    y = x * line.m + line.n
    return P(x, y)


def calc_slope_ellipse(p: P) -> float:
    """ Given by Project Euler. """
    return -4 * p.x / p.y


def plot_first_ten():
    x = np.linspace(-10, 10, 10**5)
    y_positive = np.sqrt(100 - 4*x**2)
    y_negative = -np.sqrt(100 - 4*x**2)

    plt.figure(figsize=(6,6)) 
    plt.plot(x, y_positive, 'b', label='y=sqrt(10-4x^2)') 
    plt.plot(x, y_negative, 'b', label='y=-sqrt(10-4x^2)') 
    plt.xlabel('x') 
    plt.ylabel('y') 
    plt.title('Plot of the ellipse 4x^2 + y^2 = 10') 
    plt.axis('equal') 

    current_point = P(0.0, 10.1)
    next_point = P(1.4, -9.6)
    line = line_from_two_points(current_point, next_point)
    next_point_computed = get_next_point(current_point, line)
    line.plot(0, 1.4)
    assert(next_point == next_point_computed)
    current_point = next_point

    for _ in range(10):
        m1 = line.m
        m2 = calc_slope_ellipse(current_point)
        m3 = reflected_slope(m1, m2)
        line = line_from_one_point(current_point, m3)
        next_point = get_next_point(current_point, line)
        x_min = min(current_point.x, next_point.x)
        x_max = max(current_point.x, next_point.x)
        line.plot(x_min, x_max, 'g')
        current_point = next_point

    plt.legend()
    plt.show()


def euler_144():
    plot_first_ten()
    current_point = P(0.0, 10.1)
    next_point = P(1.4, -9.6)
    line = line_from_two_points(current_point, next_point)
    current_point = next_point
    counter = 0
    while True:
        counter += 1
        m1 = line.m
        m2 = calc_slope_ellipse(current_point)
        m3 = reflected_slope(m1, m2)
        line = line_from_one_point(current_point, m3)
        next_point = get_next_point(current_point, line)
        current_point = next_point
        if -0.01 <= current_point.x and current_point.x <= 0.01 and current_point.y > 0:
            break
    return counter


if __name__ == "__main__":
    solution = euler_144()
    print("e144.py: " + str(solution))
    assert(solution == 354)