import os
import sys
import subprocess

# ──────────────────────────────────────────────────────────────
# VENV BOOTSTRAPPER
# ──────────────────────────────────────────────────────────────
# Replace 'your_username' with your actual system username
VENV_PATH = os.path.expanduser("~/Documents2/python/myenv")

def bootstrap_venv():
    """Adds the virtual environment site-packages to Blender's sys.path."""
    # Try to find the site-packages folder
    # On Linux/Mac: lib/pythonX.Y/site-packages
    # We search for any folder matching 'site-packages' inside the venv
    found_path = None
    for root, dirs, files in os.walk(VENV_PATH):
        if "site-packages" in dirs:
            found_path = os.path.join(root, "site-packages")
            break
    
    if found_path:
        if found_path not in sys.path:
            sys.path.append(found_path)
            print(f"PNS-Bootstrap: Linked venv packages from {found_path}")
    else:
        print(f"PNS-Bootstrap: Could not find site-packages in {VENV_PATH}")

# Run the bootstrapper before importing numpy
bootstrap_venv()

try:
    import numpy as np
    print("PNS-Bootstrap: Numpy imported successfully!")
except ImportError:
    print("PNS-Bootstrap: ERROR: Numpy not found. Check VENV_PATH.")
    # Fallback: we can't proceed with the optimized version, but we'll try 
    # to use standard python lists if needed.
    np = None

import bpy
import math

# ──────────────────────────────────────────────────────────────
# CONFIGURATION
# ──────────────────────────────────────────────────────────────
RES = 256              # Mesh resolution (RES x RES). 
FRACTIONAL_DEPTH = 2.0 # Height of the "glass" peaks
MAX_ITER = 64          # Fractal detail depth

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

def get_mandelbrot_value(cx, cy):
    z = complex(0, 0)
    c = complex(cx, cy)
    for i in range(MAX_ITER):
        if abs(z) > 2.0:
            return i / MAX_ITER
        z = z*z + c
    return 0.0

def setup_glass_material():
    mat = bpy.data.materials.new(name="PNS_Glass_Fractal")
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    links = mat.node_tree.links
    for n in nodes: nodes.remove(n)
    
    node_out = nodes.new(type='ShaderNodeOutputMaterial')
    node_bsdf = nodes.new(type='ShaderNodeBsdfPrincipled')
    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 
    
    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)
    node_ramp.color_ramp.elements[1].color = (0.8, 1.0, 0.9, 1.0)
    
    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():
    mesh = bpy.data.meshes.new("MandelbrotMesh")
    obj = bpy.data.objects.new("MandelbrotGlass", mesh)
    bpy.context.collection.objects.link(obj)
    
    # Optimization: Use Numpy for vertex calculation if available
    if np is not None:
        # Vectorized calculation
        x_range = np.linspace(X_MIN, X_MAX, RES)
        y_range = np.linspace(Y_MIN, Y_MAX, RES)
        
        verts = []
        for i in range(RES):
            for j in range(RES):
                val = get_mandelbrot_value(x_range[i], y_range[j])
                verts.append(((i - RES/2)*0.1, (j - RES/2)*0.1, val * FRACTIONAL_DEPTH))
    else:
        # Slow fallback to pure python
        verts = []
        for i in range(RES):
            for j in range(RES):
                cx = X_MIN + (i / RES) * (X_MAX - X_MIN)
                cy = Y_MIN + (j / RES) * (Y_MAX - Y_MIN)
                z_val = get_mandelbrot_value(cx, cy)
                verts.append(((i - RES/2)*0.1, (j - RES/2)*0.1, z_val * FRACTIONAL_DEPTH))
            
    faces = []
    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()
    for poly in mesh.polygons: poly.use_smooth = True
    return obj

def main():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    
    print("Generating PNS Glass Mandelbrot...")
    mandel_obj = create_mandelbrot_mesh()
    glass_mat = setup_glass_material()
    mandel_obj.data.materials.append(glass_mat)
    
    bpy.ops.object.light_add(type='AREA', location=(5, 5, 10))
    bpy.context.object.data.energy = 1000
    
    if "World" in bpy.data.worlds:
        bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[0].default_value = (0.02, 0.02, 0.05, 1.0)
    
    bpy.context.scene.render.engine = 'CYCLES'
    print("Done! Switch to Rendered View (Cycles) to see the glass effect.")

if __name__ == "__main__":
    main()