#!/usr/bin/env python3

# "Plasma" demo / screensaver
# https://eax.me/pygame/

import pygame
import math
import argparse
import os

time = 0
white = (255, 255, 255)
black = (0, 0, 0)

def generate_plasma(surface, width, height, time):
    for x in range(width):
        for y in range(height):
            # Normalized coordinates
            nx = x / width
            ny = y / height

            # Several layers of sines with different frequencies and phases
            layer1 = math.sin(nx * 8 + time * 0.5)
            layer2 = math.sin(ny * 6 + time * 0.3)
            layer3 = math.sin((nx + ny) * 5 + time * 0.7)
            layer4 = math.sin(math.sqrt(nx*nx + ny*ny) * 12 + time * 0.4)
            layer5 = math.sin((nx - ny) * 4 + time * 0.6)

            # Combine layers with different weights
            plasma = (layer1 + layer2 * 0.7 + layer3 * 0.5 + layer4 * 0.3 + layer5 * 0.4) / 2.9

            # Create RGB components with different phase shifts
            red = int((math.sin(plasma * 2 + time) + 1) * 127.5)
            green = int((math.sin(plasma * 2 + time + 2.1) + 1) * 127.5)
            blue = int((math.sin(plasma * 2 + time + 4.2) + 1) * 127.5)

            surface.set_at((x, y), (red, green, blue))

parser = argparse.ArgumentParser(description='Plasma Demo')
parser.add_argument('--width', type=int, default=1024, help='Window width (default: 1024)')
parser.add_argument('--height', type=int, default=768, help='Window height (default: 768)')
parser.add_argument('--scale', type=int, default=8, help='Plasma scale factor (default: 8)')
parser.add_argument('--step', type=float, default=0.1, help='Time increment per frame (default: 0.1)')
parser.add_argument('--fullscreen', action='store_true', help='Run in fullscreen mode')
parser.add_argument('--hide-fps', action='store_true', help='Hide FPS display')
args = parser.parse_args()

pygame.init()

if args.fullscreen:
    screen = pygame.display.set_mode((0, 0), pygame.FULLSCREEN)
    args.width, args.height = screen.get_size()
    print(f"Fullscreen resolution: {args.width}x{args.height}")
else:
    screen = pygame.display.set_mode((args.width, args.height))

pygame.display.set_caption("Plasma Demo")
plasma_surface = pygame.Surface((args.width // args.scale, args.height // args.scale))

font = pygame.font.Font(None, 36)
hide_fps = args.hide_fps

pygame.mouse.set_visible(False)
clock = pygame.time.Clock()
running = True
time = 0
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_q or event.key == pygame.K_ESCAPE:
                running = False
            elif event.key == pygame.K_f:
                hide_fps = not hide_fps

    generate_plasma(plasma_surface, args.width // args.scale, args.height // args.scale, time)
    time += args.step

    scaled_surface = pygame.transform.scale(plasma_surface, (args.width, args.height))
    screen.blit(scaled_surface, (0, 0))

    # Show FPS unless turned off
    if not hide_fps:
        fps = clock.get_fps()
        fps_text = font.render(f"FPS: {fps:.1f}", True, white)
        # Black background for better readability
        text_rect = fps_text.get_rect()
        background_rect = pygame.Rect(10, 10, text_rect.width + 10, text_rect.height + 10)
        pygame.draw.rect(screen, black, background_rect)
        screen.blit(fps_text, (15, 15))

    pygame.display.flip()
    clock.tick(60)

# Workaround for https://github.com/pygame/pygame/issues/329
os._exit(0)