import pygame
import sys
import math
from equation import Equation

# Initialize Pygame
pygame.init()

# Screen dimensions

# Screen setup
# screen = pygame.display.set_mode((800, 600))
screen = pygame.display.set_mode()
pygame.display.set_caption("Equation Plotter")
pygame.mouse.set_visible(False)
MOD_KEYS = pygame.key.get_mods()

class Hypotrochoid(Equation):

    def __init__(self, R, r, d, c=(255,255,255)):
        self.R = R
        self.r = r
        self.d = d
        super().__init__(self._h_xf, self._h_yf, self._h_cf, 5)

    def _h_xf(self, t):
        return (self.R - self.r) * math.cos(t) + self.d * math.cos(((self.R - self.r) / self.r) * t)

    def _h_yf(self, t):
        return (self.R - self.r) * math.sin(t) - self.d * math.sin(((self.R - self.r) / self.r) * t)

    def _h_cf(self, t, r):
        c = int(((t - r[0]) / (r[1] - r[0])) * 255)
        return c, c, c

    # def _h_wf(self, t, r):
    #     return int(((t - r[0]) / (r[1] - r[0])) * 5)

def main(equ):
    SCREEN_WIDTH = screen.get_width()
    SCREEN_HEIGHT = screen.get_height()
    HALF_WIDTH = SCREEN_WIDTH // 2
    HALF_HEIGHT = SCREEN_HEIGHT // 2
    # Main loop
    off = 0
    mouse_x_last, mouse_y_last = pygame.mouse.get_pos()
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        # Clear the screen
        screen.fill((0,0,0))

        for prev_x, prev_y, x, y, c, w in equ.plot((0 + off, (8 * math.pi) + off, math.pi/32), (100, 100)):
            pygame.draw.line(screen, c, (prev_x + HALF_WIDTH, HALF_HEIGHT - prev_y), (x + HALF_WIDTH, HALF_HEIGHT - y), w)

        pygame.display.flip()
        pygame.time.Clock().tick(60)
        off += math.pi/128

        # if direction:
        #     equ.R += 0.1
        #     equ.r += 0.05
        #     if equ.R >= 20:
        #         direction = False
        # else:
        #     equ.R -= 0.1
        #     equ.r -= 0.05
        #     if equ.R <= 6:
        #         direction = True

        mouse_x, mouse_y = pygame.mouse.get_pos()
        if 5 < math.sqrt(pow(mouse_x_last - mouse_x, 2) + pow(mouse_y_last - mouse_y, 2)):
            running = False
        keys = pygame.key.get_pressed()
        if any(keys[i] for i in range(len(keys))):
            running = False
        if pygame.key.get_mods() != MOD_KEYS:
            running = False

    pygame.quit()
    sys.exit()

# main(Equation(None, lambda t: math.sin(t)))
# main(Hypotrochoid(6, 4, 1))
main(Hypotrochoid(7, 4, 1))
# main(Hypotrochoid(15, 14, 1))