import bpy
import os
os.system('source ~/Documents2/python/myenv/bin/activate')
import numpy as np

# ──────────────────────────────────────────────────────────────
# CONFIGURATION
# ──────────────────────────────────────────────────────────────
RES = 256              # Mesh resolution (RES x RES). Increase to 512 for ultra-high.
FRACTIONAL_DEPTH = 2.0 # Height of the "glass" peaks
MAX_ITER = 64          # Fractal detail depth
SAMPLES = 100          # Scaling factor for the mesh

# Fractal boundaries
X_MIN, X_MAX = -2.0, 0.5
Y_MIN, Y_MAX = -1.25, 1.25

def get_mandelbrot_value(cx, cy):
    """
    Implements the Mandelbrot escape-time algorithm.
    Theory: The state (z) is an uncollapsed particle moving through 
    the complex plane. We only 'measure' (collapse) when |z| > 2.
    """
    z = complex(0, 0)
    c = complex(cx, cy)
    for i in range(MAX_ITER):
        if abs(z) > 2.0:
            return i / MAX_ITER  # Collapsed value (normalized height)
        z = z*z + c
    return 0.0

def setup_glass_material():
    """Creates a high-end glass material for the fractal."""
    mat = bpy.data.materials.new(name="PNS_Glass_Fractal")
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    links = mat.node_tree.links
    
    # Clear default nodes
    for n in nodes:
        nodes.remove(n)
        
    # Create Nodes
    node_out = nodes.new(type='ShaderNodeOutputMaterial')
    node_bsdf = nodes.new(type='ShaderNodeBsdfPrincipled')
    
    # Glass Properties
    node_bsdf.inputs['Base Color'].default_value = (0.8, 0.9, 1.0, 1.0)
    node_bsdf.inputs['Roughness'].default_value = 0.05
    node_bsdf.inputs['IOR'].default_value = 1.45
    node_bsdf.inputs['Transmission Weight'].default_value = 1.0 # Glass mode
    
    # Add some color modulation based on height (simulated via vertex colors or simple mix)
    # In a real render, we'd use a Coordinate node + ColorRamp
    node_coord = nodes.new(type='ShaderNodeTexCoord')
    node_ramp = nodes.new(type='ShaderNodeValToRGB')
    node_sep = nodes.new(type='ShaderNodeSeparateXYZ')
    
    node_ramp.color_ramp.elements[0].color = (0.1, 0.2, 0.8, 1.0) # Deep blue
    node_ramp.color_ramp.elements[1].color = (0.8, 1.0, 0.9, 1.0) # Pale Cyan
    
    links.new(node_coord.outputs['Generated'], node_sep.inputs['Vector'])
    links.new(node_sep.outputs['Z'], node_ramp.inputs['Fac'])
    links.new(node_ramp.outputs['Color'], node_bsdf.inputs['Base Color'])
    links.new(node_bsdf.outputs['BSDF'], node_out.inputs['Surface'])
    
    return mat

def create_mandelbrot_mesh():
    """Generates a mesh where Z = Mandelbrot height."""
    mesh = bpy.data.meshes.new("MandelbrotMesh")
    obj = bpy.data.objects.new("MandelbrotGlass", mesh)
    bpy.context.collection.objects.link(obj)
    
    verts = []
    faces = []
    
    # 1. Calculate vertices using the PNS collapse logic
    for i in range(RES):
        for j in range(RES):
            # Map i, j to complex plane
            cx = X_MIN + (i / RES) * (X_MAX - X_MIN)
            cy = Y_MIN + (j / RES) * (Y_MAX - Y_MIN)
            
            # Particle measurement (collapse to height)
            z_val = get_mandelbrot_value(cx, cy)
            
            # Position in 3D space
            x = (i - RES/2) * 0.1
            y = (j - RES/2) * 0.1
            z = z_val * FRACTIONAL_DEPTH
            verts.append((x, y, z))
            
    # 2. Create faces (Grid topology)
    for i in range(RES - 1):
        for j in range(RES - 1):
            v1 = i * RES + j
            v2 = (i + 1) * RES + j
            v3 = (i + 1) * RES + (j + 1)
            v4 = i * RES + (j + 1)
            faces.append((v1, v2, v3, v4))
            
    mesh.from_pydata(verts, [], faces)
    mesh.update()
    
    # Apply Smooth Shading
    for poly in mesh.polygons:
        poly.use_smooth = True
        
    return obj

def main():
    # Clean scene
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    
    print("Generating PNS Glass Mandelbrot... Please wait.")
    
    # Create Geometry
    mandel_obj = create_mandelbrot_mesh()
    
    # Create and Assign Material
    glass_mat = setup_glass_material()
    mandel_obj.data.materials.append(glass_mat)
    
    # Setup Lighting for better Glass rendering
    bpy.ops.object.light_add(type='AREA', location=(5, 5, 10))
    bpy.context.object.data.energy = 1000
    
    # Add a HDRI-like environment color
    bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[0].default_value = (0.02, 0.02, 0.05, 1.0)
    
    # Set renderer to Cycles for actual glass refraction
    bpy.context.scene.render.engine = 'CYCLES'
    if bpy.app.version >= (3, 0, 0):
        bpy.context.scene.cycles.device = 'GPU' 
        
    print("Done! Switch to Rendered View (Cycles) to see the glass effect.")

if __name__ == "__main__":
    main()