import numpy as np
from pymol import cmd
from pymol.cgo import *
from pymol.vfont import plain
from chempy import cpv

## coordinate system
N=np.array([-0.8151992644651752, -0.14564473420862134, -0.5605691488235305])
newX = np.array([-0.14682247172355264, -0.8842844002080097, 0.4432654524613293])
newY = np.array([-0.5602618325501247, 0.44365381881153754, 0.6994840727575685])

## points for 1SX4
cmA=np.array([36.37642466838607, 55.57413055276292, -16.29294933034555])
cmB=np.array([26.467471047520217, 61.43781187037094, -27.948758862919276])
Uref=np.array([-1.3062681581147695, -7.867409812336256, 3.943698395384567])

## points 1XCK
cmA2=np.array([36.170959390338546, 55.64974584934531, -16.41281204075117])
cmB2=np.array([28.38129355921799, 55.36428870246669, -30.93529161769317])
Uxck=np.array([-4.057265740736864, -1.8311335753945454, 6.375975235130947])

obj=[]
obj.extend([SPHERE, cmA[0], cmA[1], cmA[2], 0.35])
obj.extend([SPHERE, cmB[0], cmB[1], cmB[2], 0.35])
#obj.extend([SPHERE, cmA2[0], cmA2[1], cmA2[2], 0.35])
obj.extend([SPHERE, cmB2[0], cmB2[1], cmB2[2], 0.35])
cmd.load_cgo(obj, 'CM')

radius = 10 # cylinder width 
w = 0.1
l = 2.5 # cylinder length
h = 0.15 # cone hight
d = w * 1.618 # cone base diameter
 
start=cmA
endZ=cmA+N*l
endZcone=cmA+N*l*1.1
endX=cmA+newX*l*(-1)
endXcone=cmA+newX*l*1.1*(-1)
endY=cmA+newY*l*(-1)
endYcone=cmA+newY*l*1.1*(-1)

obj = [CYLINDER, cmA[0], cmA[1], cmA[2], endZ[0], endZ[1], endZ[2], w, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0,
       CYLINDER, cmA[0], cmA[1], cmA[2], endX[0], endX[1], endX[2], w, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0,
       CYLINDER, cmA[0], cmA[1], cmA[2], endY[0], endY[1], endY[2], w, 1.0, 0.55, 0.0, 1.0, 0.55, 0.0,
       CONE, endZ[0], endZ[1], endZ[2], endZcone[0], endZcone[1], endZcone[2], d, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0,
       CONE, endX[0], endX[1], endX[2], endXcone[0], endXcone[1], endXcone[2], d, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0,
       CONE, endY[0], endY[1], endY[2], endYcone[0], endYcone[1], endYcone[2], d, 0.0, 1.0, 0.55, 0.0, 1.0, 0.55, 0.0, 1.0, 1.0 ]
cmd.load_cgo(obj, 'axes')

w = 0.08

R=cmA-cmB
Rc1a=cmB+R*0.05
Rc1b=cmB+R*0.01
endU1b=cmA+Uref*(-1)
endU1a=cmA+Uref*(-0.93)

R=cmA-cmB2
Rc2a=cmB2+R*0.05
Rc2b=cmB2+R*0.01
endU2b=cmA+Uxck*(-1)
endU2a=cmA+Uxck*(-0.93)

colA1=np.array([250.,128., 114.])/255
colA2=np.array([139.,0., 0.])/255

colB1=np.array([135., 206., 250.])/255
colB2=np.array([0.,0., 139.])/255

obj = [CYLINDER, cmA[0], cmA[1], cmA[2], Rc1a[0], Rc1a[1], Rc1a[2],  w, colA1[0], colA1[1], colA1[2], colA1[0], colA1[1], colA1[2],
       CYLINDER, cmA[0], cmA[1], cmA[2], endU1a[0], endU1a[1], endU1a[2],  w, colA2[0], colA2[1], colA2[2], colA2[0], colA2[1], colA2[2],
       CONE, Rc1a[0], Rc1a[1], Rc1a[2], Rc1b[0], Rc1b[1], Rc1b[2], d, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0,
       CONE, endU1a[0], endU1a[1], endU1a[2], endU1b[0], endU1b[1], endU1b[2], d, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0,
       
       CYLINDER, cmA[0], cmA[1], cmA[2], Rc2a[0], Rc2a[1], Rc2a[2],  w, colB1[0], colB1[1], colB1[2], colB1[0], colB1[1], colB1[2],
       CYLINDER, cmA[0], cmA[1], cmA[2], endU2a[0], endU2a[1], endU2a[2],  w, colB2[0], colB2[1], colB2[2], colB2[0], colB2[1], colB2[2],
       CONE, Rc2a[0], Rc2a[1], Rc2a[2], Rc2b[0], Rc2b[1], Rc2b[2], d, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 
       CONE, endU2a[0], endU2a[1], endU2a[2], endU2b[0], endU2b[1], endU2b[2], d, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0 ]
cmd.load_cgo(obj, 'vector')


obj = [LINEWIDTH, 3, 
       BEGIN, LINES, 
       COLOR, 0,0,0,
       VERTEX, Rc2b[0], Rc2b[1], Rc2b[2],
       VERTEX, endU2b[0], endU2b[1], endU2b[2],
       VERTEX, Rc1b[0], Rc1b[1], Rc1b[2],
       VERTEX, endU1b[0], endU1b[1], endU1b[2] ]

cmd.load_cgo(obj, 'dash')

obj=[]
obj.extend([SPHERE, endU1b[0], endU1b[1], endU1b[2], 0.2])
obj.extend([SPHERE, endU2b[0], endU2b[1], endU2b[2], 0.2])
cmd.load_cgo(obj, 'endU')


def makePrimitive(cgo, name):
    cmd.set('auto_zoom', 0, quiet=1)
    cmd.load_cgo(cgo, name)
    cmd.set('auto_zoom', 1, quiet=1)
 
def point(p):
    x, y, z = p
    return [COLOR, 1, 1, 1, SPHERE, float(x), float(y), float(z), 0.5]
 
def line(p1, p2):
    x1, y1, z1 = p1
    x2, y2, z2 = p2
    return [CYLINDER, float(x1), float(y1), float(z1), float(x2), float(y2), float(z2), 0.25, 1, 1, 1, 1, 1, 1]
 
def plane(corner1, corner2, corner3, corner4, normal):
    planeObj = []
    planeObj.extend(point(corner1))
    planeObj.extend(point(corner2))
    planeObj.extend(point(corner3))
    planeObj.extend(point(corner4))
    planeObj.extend(line(corner1, corner2))
    planeObj.extend(line(corner2, corner3))
    planeObj.extend(line(corner3, corner4))
    planeObj.extend(line(corner4, corner1))
 
    planeObj.extend([COLOR, 0.8, 0.8, 0.8])
    planeObj.extend([BEGIN, TRIANGLE_STRIP])
    planeObj.append(NORMAL)
    planeObj.extend(normal)
    for corner in [corner1, corner2, corner3, corner4, corner1]:
        planeObj.append(VERTEX)
        planeObj.extend(corner)
    planeObj.append(END)
    return planeObj, corner4

def planeFromPoints(point1, point2, point3, facetSize):
    v1 = cpv.normalize(cpv.sub(point2, point1))
    v2 = cpv.normalize(cpv.sub(point3, point1))
    normal = cpv.cross_product(v1, v2)
    v2 = cpv.cross_product(normal, v1)
    x = cpv.scale(v1, facetSize)
    y = cpv.scale(v2, facetSize)
    center = point2
    corner1 = cpv.add(cpv.add(center, x), y)
    corner2 = cpv.sub(cpv.add(center, x), y)
    corner3 = cpv.sub(cpv.sub(center, x), y)
    corner4 = cpv.add(cpv.sub(center, x), y)
    return plane(corner1, corner2, corner3, corner4, normal)


obj, corner1 = planeFromPoints(cmA, endU1b, endU2b, 10)
cmd.load_cgo(obj, 'plane')


#endZ=corner1+N*l
#endZcone=corner1+N*l*1.05
#endX=corner1+newX*l
#endXcone=corner1+newX*l*1.05
#endY=corner1+newY*l
#endYcone=corner1+newY*l*1.05

#obj = [CYLINDER, corner1[0], corner1[1], corner1[2], endZ[0], endZ[1], endZ[2], w, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0,
#       CYLINDER, corner1[0], corner1[1], corner1[2], endX[0], endX[1], endX[2], w, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0,
#       CYLINDER, corner1[0], corner1[1], corner1[2], endY[0], endY[1], endY[2], w, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0,
#       CONE, endZ[0], endZ[1], endZ[2], endZcone[0], endZcone[1], endZcone[2], d, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0,
#       CONE, endX[0], endX[1], endX[2], endXcone[0], endXcone[1], endXcone[2], d, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0,
#       CONE, endY[0], endY[1], endY[2], endYcone[0], endYcone[1], endYcone[2], d, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0 ]
#cmd.load_cgo(obj, 'axes')