from dataclasses import dataclass
from math import pi, tan, atan2, sqrt

import numpy as np
import matplotlib.pyplot as plt

@dataclass
class Line:
    x: float
    y: float
    angle: float

    def __repr__(self):
        return f"Line(x={self.x}, y={self.y}, angle={self.angle*180/pi})"


def bounce_wall(line):
    if line.angle > pi:
        raise ValueError
    intercept_x = line.x - line.y / tan(line.angle)
    angle = -line.angle + 2*pi
    return Line(x=intercept_x, y=0, angle=angle)


def bounce_circle(line):
    x0 = line.x
    t = tan(line.angle)
    t2 = t*t

    a = 1 + t2
    b = 4*t - 2*t2*x0
    c = t2*x0*x0 - 4*x0*t + 3

    rad = b*b - 4*a*c
    if rad < 0:
        raise ValueError

    x1 = (-b + sqrt(rad)) / (2*a)
    y1 = (x1 - x0)*t

    x2 = (-b - sqrt(rad)) / (2*a)
    y2 = (x2 - x0)*t

    x, y = (x1, y1) if y1 > y2 else (x2, y2)

    t_tan = (atan2(y+2, x) - pi/2) % (2*pi)
    angle = (line.angle + 2*(t_tan - line.angle)) % (2*pi)
    return Line(x=x, y=y, angle=angle)


def do_bounces(angle_deg, x=None):
    if x is not None:
        angle_deg = np.arctan2(2, x + 2) * 180 / pi
    lines = [Line(-2, -2, angle_deg*pi/180)]
    try:
        while True:
            lines.append(bounce_wall(lines[-1]))
            lines.append(bounce_circle(lines[-1]))
    except ValueError:
        pass
    return len(lines)-1, lines


def draw_bounces(lines):
    xs = [line.x for line in lines]
    ys = [line.y for line in lines]
    plt.plot([-4, 4], [0, 0])
    plt.plot(xs, ys)
    plt.gca().add_patch(plt.Circle((0, -2), radius=1))

    plt.xlim((-1, 1))
    plt.ylim((-1.8, 0.2))

    plt.show()


def count_range(start, end, count=100):
    xs = np.linspace(start, end, count)
    angles = np.arctan2(2, xs + 2) * 180 / pi
    bounces = []

    for angle in angles:
        count, _ = do_bounces(angle)
        bounces.append(count)

    plt.plot(xs, bounces)
    plt.show()
    return xs, np.array(bounces)


def iterative_zoom(start, end, n_iter):

    maxes = []
    for _ in range(n_iter):
        xs, bounces = count_range(start, end, 1000)
        max_x = xs[np.argmax(bounces)]
        maxes.append(max_x)
        range_ = end - start
        start = max_x - 0.05*range_
        end = max_x + 0.05*range_
    print(maxes)
    plt.plot(maxes)
    plt.show()


def main():

    # count_range(-2, 0, 10000)
    # count_range(-0.824, -0.822, 10000)
    # count_range(-0.8225, -0.8223, 10000)
    # count_range(58, 62)
    # count_range(59, 60)
    # count_range(59.51, 59.515)
    # draw_bounces(do_bounces(None, x=-0.8224863249433897)[1])
    iterative_zoom(-2, 0, 20)


if __name__ == '__main__':
    main()