import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import time
import os

print(f"🚀 Starting Julia Set Generation on {os.uname().nodename}...")

def julia_set(h, w, c, max_iter=100):
    """Generate Julia set for a given complex parameter c."""
    y, x = np.ogrid[-1.5:1.5:h*1j, -1.5:1.5:w*1j]
    z = x + y*1j
    divtime = max_iter + np.zeros(z.shape, dtype=int)

    for i in range(max_iter):
        z = z**2 + c
        diverge = z*np.conj(z) > 2**2 
        div_now = diverge & (divtime == max_iter)
        divtime[div_now] = i
        z[diverge] = 2
    
    return divtime

# Configuration
H, W = 1000, 1500
MAX_ITER = 200
# Interesting Julia set parameter
C = complex(-0.7, 0.27015) 

print(f"Computing Julia set ({W}x{H}) for c={C}...")
start_time = time.time()

divtime = julia_set(H, W, C, MAX_ITER)

compute_time = time.time() - start_time
print(f"✅ Computation done in {compute_time:.4f}s")

# Plot
print("Rendering image...")
plt.figure(figsize=(15, 10))
plt.imshow(divtime, cmap='twilight', extent=[-1.5, 1.5, -1.5, 1.5])
plt.axis('off')
plt.title(f'Julia Set c={C} (Compute: {compute_time:.2f}s)')

filename = "julia_gpu.png"
plt.savefig(filename, bbox_inches='tight', pad_inches=0, dpi=150)
print(f"💾 Saved to {filename}")