import maya.cmds as cmds
import random

##Created by Isabella Clements Leedeman 2024
#### Snap object to mesh surface with optional offset
def snap_to_surface(obj, mesh, offset=0.0):
    pos = cmds.xform(obj, q=True, ws=True, t=True)

    cpm = cmds.createNode("closestPointOnMesh")
    cmds.connectAttr(mesh + ".worldMesh[0]", cpm + ".inMesh")
    cmds.setAttr(cpm + ".inPosition", pos[0], pos[1], pos[2])

    point = cmds.getAttr(cpm + ".position")[0]
    normal = cmds.getAttr(cpm + ".normal")[0]

    cmds.delete(cpm)

    final_pos = [
        point[0] + normal[0] * offset,
        point[1] + normal[1] * offset,
        point[2] + normal[2] * offset
    ]

    cmds.xform(obj, ws=True, t=final_pos)


#### Scatter function

def scatter_trees(tree_group, mesh, scale_factor, random_rotation, random_placement, ground_offset, amount):

    if not cmds.objExists(tree_group):
        cmds.warning("Tree group '{}' does not exist.".format(tree_group))
        return

    if not cmds.objExists(mesh):
        cmds.warning("Mesh '{}' does not exist.".format(mesh))
        return

    #### Bounding box for random X/Z placement
    mesh_bbox = cmds.exactWorldBoundingBox(mesh)
    min_x, min_y, min_z, max_x, max_y, max_z = mesh_bbox

    for _ in range(amount):

        ####Duplicate tree group
        duplicate_tree = cmds.duplicate(
            tree_group,
            returnRootsOnly=True,
            inputConnections=True,
            name=tree_group + "_instance#"
        )[0]

        ####Random Y-axis rotation
        rotation_value = random.uniform(-random_rotation, random_rotation)
        cmds.rotate(0, rotation_value, 0, duplicate_tree)

        ####Random X/Z placement (start above mesh)
        x_pos = random.uniform(min_x, max_x)
        z_pos = random.uniform(min_z, max_z)
        y_pos = max_y

        cmds.move(x_pos, y_pos, z_pos, duplicate_tree)

        ####Snap to mesh surface + offset control
        snap_to_surface(duplicate_tree, mesh, ground_offset)

        ####Scale adjustment
        cmds.scale(scale_factor, scale_factor, scale_factor, duplicate_tree)



#### GUI

def create_tree_scatter_gui():

    window_name = "treeScatterWindow"
    if cmds.window(window_name, exists=True):
        cmds.deleteUI(window_name, window=True)

    cmds.window(window_name, title="Tree Scatter", sizeable=True)
    cmds.columnLayout(adjustableColumn=True)

    #### Tree Group dropdown
    tree_group_menu = cmds.optionMenu(label="Tree Group:")
    tree_groups = cmds.ls(type="transform", dag=True, long=True)
    for group in tree_groups:
        cmds.menuItem(label=group, parent=tree_group_menu)

    #### Mesh dropdown
    mesh_menu = cmds.optionMenu(label="Mesh:")
    meshes = cmds.ls(geometry=True, long=True)
    for mesh in meshes:
        cmds.menuItem(label=mesh, parent=mesh_menu)

    #### Scale slider
    scale_slider = cmds.floatSliderGrp(
        label="Scale Factor:",
        field=True,
        minValue=0.1,
        maxValue=2.0,
        value=1.0
    )

    #### Random Y rotation slider
    rotation_slider = cmds.floatSliderGrp(
        label="Random Rotation (Y-axis):",
        field=True,
        minValue=0,
        maxValue=180,
        value=30
    )

    #### Random placement slider (kept for compatibility)
    placement_slider = cmds.floatSliderGrp(
        label="Random Placement:",
        field=True,
        minValue=0,
        maxValue=10,
        value=5
    )

    #### Ground offset slider
    ground_offset_slider = cmds.floatSliderGrp(
        label="Ground Offset:",
        field=True,
        minValue=-1.0,
        maxValue=2.0,
        value=0.0
    )

    #### Amount slider
    amount_slider = cmds.intSliderGrp(
        label="Amount:",
        field=True,
        minValue=1,
        maxValue=100,
        value=10
    )

    #### Scatter button
    cmds.button(
        label="Scatter Trees",
        command=lambda x: scatter_trees(
            cmds.optionMenu(tree_group_menu, query=True, value=True),
            cmds.optionMenu(mesh_menu, query=True, value=True),
            cmds.floatSliderGrp(scale_slider, query=True, value=True),
            cmds.floatSliderGrp(rotation_slider, query=True, value=True),
            cmds.floatSliderGrp(placement_slider, query=True, value=True),
            cmds.floatSliderGrp(ground_offset_slider, query=True, value=True),
            cmds.intSliderGrp(amount_slider, query=True, value=True)
        )
    )

    cmds.showWindow(window_name)

create_tree_scatter_gui()