from chimerax.markers import MarkerSet
from chimerax.core.commands import ColorArg, CmdDesc, register, FloatArg
from chimerax.atomic import AtomsArg

import numpy as np
from scipy.spatial import distance_matrix
from scipy.spatial.transform import Rotation


def display_clash(session, selection, withAtoms=None, color=(255, 0, 0), scale=1):
    if not isinstance(color, tuple):
        color = color.uint8x4()

    # markersets for VDW radii
    ms1 = MarkerSet(session)
    ms2 = MarkerSet(session)
    # markerset for intersection of VDW radii
    ms3 = MarkerSet(session)
    ms1_members = set()
    ms2_members = set()

    # distance between each pair of atoms so we can easily figure out which
    # ones overlap
    dist = distance_matrix(selection.coords, withAtoms.coords)
    
    points_per_circle = 200
    angles = np.linspace(0, 2 * np.pi, num=points_per_circle)
    circle = np.array([np.cos(angles), np.sin(angles), np.zeros(points_per_circle)]).T
    
    radii_1 = scale * selection.default_radii
    radii_2 = scale * withAtoms.default_radii
    
    # add some points where the atoms interect each other
    for i, a1 in enumerate(selection):
        r1 = radii_1[i]
        for j, a2 in enumerate(withAtoms):
            r2 = radii_2[j]
            d = dist[i, j]
            if d < (r1 + r2) and d > abs(r1 - r2):
                # the intersection of two spheres is a circle
                # the plane that circle is in is perpendicular to the
                # vector from the center of one sphere to the other
                # the center of the circle is along that vector, and
                # the exact poisiton depends on the radii of the spheres
                # and the distance between their centers
                v_n = a2.coord - a1.coord
                v_n /= np.linalg.norm(v_n)
                theta = np.arccos((r2 ** 2 - r1 ** 2 - d ** 2) / (-2 * r1 * d))
                h = r1 * np.sin(theta)
                b = r1 * np.cos(theta)
                p = b * v_n
                
                rot_axis = np.cross([0, 0, 1], v_n)
                rot_axis /= np.linalg.norm(rot_axis)
                dz = v_n - [0, 0, 1]
                dz2 = np.dot(dz, dz)
                rot_angle = np.arccos(-dz2 / 2. + 1)
                R = Rotation.from_rotvec(rot_angle * rot_axis).as_matrix()
                # coordinates of points on the intersection circle
                added_points = h * np.dot(circle, R.T) + a1.coord + p
                # we will remove points that are inside of a different atom
                # this will leave us with an arc instead of a full circle
                pt_dist = distance_matrix(added_points, selection.coords)
                pt_dist -= radii_1[np.newaxis, :]
                # ensure we don't remove any points due to an apparent overlap
                # with this particular pair of atoms
                pt_dist[:, i] = 1
                mask = np.invert(np.any(pt_dist < 0, axis=1))
                if not any(mask):
                    continue
                added_points = added_points[mask]
                pt_dist = distance_matrix(added_points, withAtoms.coords)
                pt_dist -= radii_2[np.newaxis, :]
                pt_dist[:, j] = 1
                mask = np.invert(np.any(pt_dist < 0, axis=1))
                if not any(mask):
                    continue
                added_points = added_points[mask]
                # add any remaining points to ms3
                if i not in ms1_members:
                    ms1_members.add(i)
                    ms1.create_marker(a1.coord, (*a1.color[:3], 125), r1)
                if j not in ms2_members:
                    ms2_members.add(j)
                    ms2.create_marker(a2.coord, (*a2.color[:3], 125), r2)
                for pt in added_points:
                    ms3.create_marker(pt, (color[0], color[1], color[2], 255), 0.05)

    session.models.add([ms1, ms2, ms3])


desc = CmdDesc(
    required=[
        ("selection", AtomsArg),
    ],
    keyword=[
        ("color", ColorArg),
        ("withAtoms", AtomsArg),
        ("scale", FloatArg),
    ],
)


register("displayClash", desc, display_clash, logger=session.logger)