from typing import Tuple
import numpy as np
import noise
import pygame

from utils import clamp

SEED = 42
np.random.seed(SEED)


class WindField:
    def __init__(self, width, height, downscale=10):
        self.width = int(width / downscale)
        self.height = int(height / downscale)
        self.downscale = downscale
        self.field_x = np.zeros((self.width, self.height))
        self.field_y = np.zeros((self.width, self.height))

    def update(
        self,
        time=0,
        scale=15.0,
        time_scale=2.0,
        octaves=5,
        persistence=0.3,
        lacunarity=4.0,
    ):
        for i in range(self.width):
            for j in range(self.height):
                self.field_x[i][j] = noise.pnoise3(
                    i / scale,
                    j / scale,
                    time / time_scale,
                    octaves=octaves,
                    persistence=persistence,
                    lacunarity=lacunarity,
                    base=SEED,
                )
                self.field_y[i][j] = noise.pnoise3(
                    i / scale,
                    j / scale,
                    time / time_scale,
                    octaves=octaves,
                    persistence=persistence,
                    lacunarity=lacunarity,
                    base=SEED + 1,
                )

    def draw(self, screen):
        for i in range(self.width):
            for j in range(self.height):
                color = (
                    abs(self.field_x[i][j] * 255),
                    abs(self.field_y[i][j] * 255),
                    0,
                )
                pygame.draw.rect(
                    screen,
                    color,
                    (
                        i * self.downscale,
                        j * self.downscale,
                        self.downscale,
                        self.downscale,
                    ),
                )

    # draw with get_wind
    # def draw(self, screen):
    #     for i in range(self.width * self.downscale):
    #         for j in range(self.height * self.downscale):
    #             wind = self.get_wind((i, j))
    #             color = (
    #                 abs(wind[0] * 255),
    #                 abs(wind[1] * 255),
    #                 0,
    #             )
    #             pygame.draw.rect(
    #                 screen,
    #                 color,
    #                 (
    #                     i,
    #                     j,
    #                     1,
    #                     1,
    #                 ),
    #             )

    def get_wind(self, position: Tuple[int, int]) -> Tuple[float, float]:
        x, y = position
        x /= self.downscale
        y /= self.downscale

        x0 = int(x)
        y0 = int(y)
        x1 = x0 + 1
        y1 = y0 + 1

        x0 = clamp(x0, 0, self.width - 1)
        y0 = clamp(y0, 0, self.height - 1)
        x1 = clamp(x1, 0, self.width - 1)
        y1 = clamp(y1, 0, self.height - 1)

        dx = x - x0
        dy = y - y0

        wind_x = (
            (1 - dx) * (1 - dy) * self.field_x[x0][y0]
            + dx * (1 - dy) * self.field_x[x1][y0]
            + (1 - dx) * dy * self.field_x[x0][y1]
            + dx * dy * self.field_x[x1][y1]
        )
        wind_y = (
            (1 - dx) * (1 - dy) * self.field_y[x0][y0]
            + dx * (1 - dy) * self.field_y[x1][y0]
            + (1 - dx) * dy * self.field_y[x0][y1]
            + dx * dy * self.field_y[x1][y1]
        )

        return wind_x, wind_y