from pymol import cmd
import numpy as np

def distance_to_plane(atom_name, plane_object):
    """
    Measure the distance from a specific atom to a least-squares plane defined by an object.
    
    :param atom_name: Name of the atom (e.g. 'atom_name')
    :param plane_object: Name of the object that defines the plane (e.g. 'plane')
    """
    # Get coordinates of the atom of interest
    atom_coords = np.array(cmd.get_atom_coords(atom_name))
    
    # Get coordinates of the atoms that define the plane from the object
    plane_model = cmd.get_model(plane_object)
    plane_coords = np.array([atom.coord for atom in plane_model.atom])
    # check enough atoms
    if len(plane_coords) < 3:
        raise ValueError("At least 3 atoms are required to define a plane.")
    # Compute the mean of the plane atoms (centroid)
    centroid = np.mean(plane_coords, axis=0)
    
    # Compute vectors from centroid to the plane atoms
    vectors = plane_coords - centroid

    # compute the covariance matrix
    cov_matrix = np.cov(vectors, rowvar=False)

    _,_,vt = np.linalg.svd(cov_matrix)
    normal = vt[-1]
    normal /= np.linalg.norm(normal)
    
    # Calculate the normal vector using the cross product
    distance = np.abs(np.dot(normal, atom_coords - centroid))
    print(round(distance,3))
    return distance