from pygame.locals import *
import pygame
from pygame.math import Vector2

import math
from math import sin, cos, pi, sqrt, hypot
import time


def screen2local(x, y=None):
    # Если аргументом передаётся Vector2
    if y is None:
        y = x.y
        x = x.x

    nx = x - WIDTH / 2
    ny = HEIGHT / 2 - y
    return (nx, ny)


def local2screen(x, y=None):
    # Если аргументом передаётся Vector2
    if y is None:
        y = x.y
        x = x.x

    nx = x + WIDTH / 2
    ny = HEIGHT / 2 - y
    return (nx, ny)


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

        self.c = sqrt(self.a * self.a - self.b * self.b)
        #  self.f1 = local2screen(self.c, 0)
        #  self.f2 = local2screen(-self.c, 0)
        self.f1 = (self.c, 0)
        self.f2 = (-self.c, 0)

    def draw(self, deltatime, screen, cx, cy):
        cnt = 0
        for i in range(0, int(2 * pi * 100) + 1):
            cnt += 1
            t = i / 100
            t1 = t + 1 / 100

            # Координаты в локальном пространстве
            x = self.a * cos(t)
            y = self.b * sin(t)
            x1 = self.a * cos(t1)
            y1 = self.b * sin(t1)

            start = local2screen(x, y)
            end = local2screen(x1, y1)

            pygame.draw.line(screen, WHITE, start, end, 2)

        pygame.draw.circle(screen, WHITE, local2screen(*self.f1), 3)
        pygame.draw.circle(screen, WHITE, local2screen(*self.f2), 3)


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

        self.time = 0
        self.timeout = 700
        self.i = 0
        self.lines = []
        
        self.c = sqrt(self.a * self.a - self.b * self.b)
        self.f1 = local2screen(self.c, 0)
        self.f2 = local2screen(-self.c, 0)

        self.lastend = None

    def draw(self, deltatime, screen, cx, cy):
        print(deltatime)
        for start, end in self.lines:
            pygame.draw.line(screen, WHITE, start, end, 2)

        pygame.draw.circle(screen, WHITE, self.f1, 3)
        pygame.draw.circle(screen, WHITE, self.f2, 3)

        if self.lastend is not None:
            pygame.draw.line(screen, RED, self.f1, self.lastend, 2)
            pygame.draw.line(screen, GREEN, self.f2, self.lastend, 2)

        if self.i >= int(2 * pi * 100) + 1:
            self.lastend = None
            return

        if self.time > 0:
            self.time -= deltatime
            return

        self.time = self.timeout

        t = self.i / 100
        t1 = t + 1 / 100
        self.i += 1

        # Координаты в локальном пространстве
        x = self.a * cos(t)
        y = self.b * sin(t)
        x1 = self.a * cos(t1)
        y1 = self.b * sin(t1)

        start = local2screen(x, y)
        end = local2screen(x1, y1)
        self.lastend = end

        self.lines.append((start, end))

        pygame.draw.line(screen, WHITE, start, end, 2)


class ScreenRay:
    def __init__(self, a: Vector2, b: Vector2, speed: float):
        #  self.a = a
        #  self.b = b
        self.pos = a
        self.dir = Vector2(b.x - a.x, b.y - a.y)
        self.dir.normalize_ip()
        self.speed = speed

    def update(self, deltatime):
        n_pos = self.pos + self.dir * (deltatime * self.speed)
        
        scr = Vector2(*local2screen(*tuple(self.pos)))
        scr_npos = Vector2(*local2screen(*tuple(n_pos)))

        if scr_npos.x < 0:
            self.dir = Vector2(-self.dir.x, self.dir.y)
            n_pos = Vector2(*screen2local(0, scr.y))
        elif scr_npos.x > WIDTH:
            self.dir = Vector2(-self.dir.x, self.dir.y)
            n_pos = Vector2(*screen2local(WIDTH, scr.y))
        elif scr_npos.y < 0:
            self.dir = Vector2(self.dir.x, -self.dir.y)
            n_pos = Vector2(*screen2local(scr.x, 0))
        elif scr_npos.y > HEIGHT:
            self.dir = Vector2(self.dir.x, -self.dir.y)
            n_pos = Vector2(*screen2local(scr.x, HEIGHT))

        self.pos = n_pos


class EllipseRay:
    def __init__(self, s: Vector2, e: Vector2, speed: float, el):
        self.a = el.a
        self.b = el.b
        self.pos = e
        self.dir = Vector2(e.x - s.x, e.y - s.y)
        self.dir.normalize_ip()
        self.speed = speed

    def eps(self, x, y):
        a = self.a
        b = self.b
        tmp = self.ellipse(x, y)
        e = 0.05
        return tmp < e and tmp > e

    def hit(self):
        eps = 0.001
        npos = self.pos + self.dir
        tmp = self.ellipse(npos.x, npos.y)

        if tmp > eps:
            return self.find_point(self.pos.x, self.pos.y, npos.x, npos.y)
        else:
            return None
    
    def ellipse(self, x, y):
        return (x * x) / (self.a * self.a) + (y * y) / (self.b * self.b) - 1

    def find_point(self, x0, y0, x1, y1, depth=0):
        eps = 0.0001
        xc = (x0 + x1) / 2
        yc = (y0 + y1) / 2
        tmp = self.ellipse(xc, yc)
        #  if abs(tmp) < eps or hypot(x1 - x0, y1 - y0) < eps:
        if abs(tmp) < eps:
            return Vector2(xc, yc)
        elif tmp >= eps:
            #  print(tmp, (x0, y0), (x1, y1), (xc, yc), (x1 - x0, y1 - y0))
            t1 = self.ellipse(xc, yc)
            t2 = self.ellipse(x0, y0)
            if t1 > 0 and t2 > 0:
                print(">", t1, t2)

            if depth > 20:
                print("> ret")
                return Vector2(xc, yc)
            return self.find_point(x0, y0, xc, yc, depth+1)
        elif tmp <= -eps:
            t1 = self.ellipse(xc, yc)
            t2 = self.ellipse(x1, y1)
            if t1 > 0 and t2 > 0:
                print("<", t1, t2)

            if depth > 20:
                print("< ret")
                return Vector2(xc, yc)
            return self.find_point(xc, yc, x1, y1, depth+1)

    def update(self, deltatime):
        npos = self.pos + self.dir * (deltatime * self.speed)

        hit_result = self.hit()
        if hit_result is None:
            self.pos = npos
            return

        npos = hit_result
        scr = Vector2(*local2screen(*tuple(self.pos)))
        scr_npos = Vector2(*local2screen(*tuple(npos)))
        p0 = npos

        ########################
        # Отражение от эллипса #
        ########################

        if npos.y == 0:
            if npos.x < 0:
                self.dir.reflect_ip(Vector2(-1, 0))
            else:
                self.dir.reflect_ip(Vector2(1, 0))
        else:
            t = 1
            m = 1
            n = -(p0.x * self.b * self.b) / (p0.y * self.a * self.a)
            p1 = Vector2(p0.x + m * t, p0.y + n * t)

            kas = Vector2(p1.x - p0.x, p1.y - p0.y)
            kas.normalize_ip()
            #  angle = kas.angle_to(self.dir)

            kas_p = Vector2(1, -kas.x / kas.y)
            kas_p.normalize_ip()

            self.dir.reflect_ip(kas_p)

        ########################

        self.pos = npos


    def draw(self, screen):
        loc = tuple(self.pos)
        scr = local2screen(*loc)
        pygame.draw.circle(screen, RED, scr, 3)


class LaserEllipse:
    def __init__(self, a, b, ray_t=0, ray_end=None):
        self.a = a
        self.b = b

        if ray_end is None:
            ray_end = Vector2(0, 0)

        self.ray_end = ray_end
        self.ray_t = ray_t

        self.c = sqrt(self.a * self.a - self.b * self.b)
        self.f1 = Vector2(self.c, 0)
        self.f2 = Vector2(-self.c, 0)

    def ellipse(self, x, y):
        return (x * x) / (self.a * self.a) + (y * y) / (self.b * self.b) - 1

    def get_normal(self, p0):
        m = 1
        n = -(p0.x * self.b * self.b) / (p0.y * self.a * self.a)
        p1 = Vector2(p0.x + m, p0.y + n)

        kas = Vector2(p1.x - p0.x, p1.y - p0.y)
        kas.normalize_ip()

        normal = Vector2(1, -kas.x / kas.y)
        return normal.normalize()

    def find_point(self, x0, y0, x1, y1, depth=0):
        eps = 0.0001
        xc = (x0 + x1) / 2
        yc = (y0 + y1) / 2
        tmp = self.ellipse(xc, yc)
        if abs(tmp) < eps or hypot(x1 - x0, y1 - y0) < eps:
            return Vector2(xc, yc)
        elif tmp >= eps:
            return self.find_point(x0, y0, xc, yc, depth+1)
        elif tmp <= -eps:
            return self.find_point(xc, yc, x1, y1, depth+1)

    def parametric(self, t):
        return Vector2(self.a * cos(t), self.b * sin(t))

    def draw_lasers(self, screen):
        start = self.parametric(self.ray_t)
        d = (self.ray_end - start).normalize()
        end = self.get_laser(start, d)
        for i in range(50):
            pygame.draw.line(screen, RED, local2screen(start), local2screen(end), 2)

            n = self.get_normal(end)
            start = end
            d.reflect_ip(n)
            end = self.get_laser(start, d)

    def get_laser(self, start: Vector2, direction: Vector2) -> Vector2:
        tmp = start + direction * ((self.a + self.b) ** 2)
        end = self.find_point(start.x, start.y, tmp.x, tmp.y)
        return end

    def update(self, deltatime):
        self.ray_t += deltatime / 10

    def draw(self, deltatime, screen, cx, cy):
        self.draw_lasers(screen)

        cnt = 0
        for i in range(0, int(2 * pi * 100) + 1):
            cnt += 1
            t = i / 100
            t1 = t + 1 / 100

            # Координаты в локальном пространстве
            p0 = self.parametric(t)
            p1 = self.parametric(t1)

            start = local2screen(p0)
            end = local2screen(p1)

            pygame.draw.line(screen, WHITE, start, end, 2)

        pygame.draw.circle(screen, WHITE, local2screen(self.f1), 3)
        pygame.draw.circle(screen, WHITE, local2screen(self.f2), 3)

        ray_start = self.parametric(self.ray_t)
        start = local2screen(ray_start)
        d = (self.ray_end - ray_start).normalize() * 30
        end = local2screen(ray_start + d)

        pygame.draw.line(screen, WHITE, start, end, 2)


pygame.init()

WIDTH, HEIGHT = SIZE = 800, 600

screen = pygame.display.set_mode(SIZE)

clock = pygame.time.Clock()
running = True

BLACK = Color("black")
WHITE = Color("white")
GREEN = Color("green")
RED = Color("red")

el = LaserEllipse(300, 200, ray_t=200, ray_end=Vector2(0, 100))

rays = []
#  nn = 10
#  for i in range(1, nn + 1):
    #  x0, y0 = el.f2
    #  xr = cos(2 * pi * i / nn) + x0
    #  yr = sin(2 * pi * i / nn) + y0
    #  r = EllipseRay(Vector2(x0, y0), Vector2(xr, yr), 70, el)
    #  rays.append(r)

#  deltatime = pygame.time.get_ticks()

while running:
    screen.fill(BLACK)

    deltatime = clock.get_time() / 1000

    for r in rays:
        r.update(deltatime)
        r.draw(screen)
    el.draw(deltatime, screen, 0, 0)
    el.update(deltatime)

    # ...

    pygame.display.flip()
    
    for event in pygame.event.get():
        if event.type == QUIT:
            pygame.quit()
            exit()
        elif event.type == KEYDOWN:
            if event.key == K_ESCAPE:
                pygame.quit()
                exit()
    clock.tick(60)
    #  deltatime = pygame.time.get_ticks() - deltatime
    #  print(clock.get_fps())