import pygame
import numpy as np

class PortalEffect:
    def __init__(self, x, y, radius=100, strength=10):
        self.x = x
        self.y = y
        self.radius = radius
        self.strength = strength

    def apply(self, surface, offset_x=0, offset_y=0):
        # Get the region around the portal
        rect = pygame.Rect(self.x - self.radius - offset_x,
                           self.y - self.radius - offset_y,
                           self.radius * 2,
                           self.radius * 2)
        sub = surface.subsurface(rect).copy()

        # Convert to array
        arr = pygame.surfarray.pixels3d(sub).copy()
        h, w = arr.shape[:2]
        y, x = np.ogrid[:h, :w]
        cx, cy = w // 2, h // 2
        dx = x - cx
        dy = y - cy
        dist = np.sqrt(dx**2 + dy**2)
        mask = dist < self.radius
        factor = np.clip((self.radius - dist) / self.radius, 0, 1)
        shift = factor * self.strength

        # Distort coordinates
        new_x = np.clip((x + dx / (dist + 1) * shift).astype(int), 0, w - 1)
        new_y = np.clip((y + dy / (dist + 1) * shift).astype(int), 0, h - 1)

        arr2 = arr[new_y, new_x]
        distorted = pygame.surfarray.make_surface(arr2)
        surface.blit(distorted, rect.topleft, special_flags=pygame.BLEND_RGBA_ADD)
