from glob import glob
import numpy as np
from arbdmodel.ssdna_two_bead import DnaModel, PolymerSection, _P, _B
from arbdmodel.coords import rotationAboutAxis
from copy import copy


gridfile = '/data/server5/cmaffeo2/latexDocuments/2013-coarse-SSB/sources/replisome/maps/combine.dx'
restraint1 = np.array((189.5, 199.5, 209.8))
restraint2 = np.array((153.9, 161.1, 196.6))

_P.grid = (gridfile, 1)
_B.grid = (gridfile, 1)

""" Create spline-based polymer """
a = PolymerSection(name='loop', num_monomers=500,
                   monomer_length = 5,
                   start_position = restraint1,
                   end_position = restraint2,
)


""" VMD: set R [lindex [molinfo top get rotate_matrix] 0]
vectrans [transtranspose $R] {0 0 1}
"""
axis1 = np.array((-0.9432600140571594, -0.12140899896621704, -0.3090659976005554))
axis2 = np.array((-0.8285319805145264, -0.4766559898853302, -0.29383501410484314))

delta = (0.4*a.num_monomers) * a.monomer_length
midpoint1 = restraint1 +  delta*axis1
midpoint2 = restraint2 +  delta*axis2

# def get_ideal_position(contour):
#     if contour < 0.4:
#         return restraint1 + delta*axis1

dna_coordinates = np.loadtxt('0-initial/p-coords.dat')
ssb_dna_contacts = [np.loadtxt('0-initial/contacts/{}.dat'.format(i)) for i in range(1,15)]
ssb_coords = [np.loadtxt('0-initial/coords/{}.dat'.format(i)) for i in range(1,15)]
ssb_dna_contact_center = [a.mean() for a in ssb_dna_contacts]
midpoints = [int(0.5*(a[-1]+b[0])) for a,b in zip(ssb_dna_contacts[:-1],ssb_dna_contacts[1:])]
edges = [0] + midpoints + [500]

idxidx1 = 4
idxidx2 = 9
idx1 = edges[idxidx1+1]
idx2 = edges[idxidx2+1]
r0 = dna_coordinates[0]
r1 = dna_coordinates[idx1]
r2 = dna_coordinates[idx2]
r3 = dna_coordinates[-1]

def rotate_to(initial, final):
    a0,a1 = [a/np.linalg.norm(a) for a in (initial, final)]
    cross = np.cross(a0,a1)
    cross_l = np.linalg.norm(cross)
    axis = cross/cross_l
    angle = np.arcsin(cross_l)*180/np.pi
    return rotationAboutAxis(axis,angle, normalizeAxis=False)

R1 = rotate_to(r1-r0, axis1)
target1 = R1.dot(r1-r0) + restraint1

tmp_R3 = rotate_to(r2-r3, axis2)
tmp_target2 = tmp_R3.dot(r2-r3) + restraint2

R2 = rotate_to(r2-r1, tmp_target2-target1)
target2 = R2.dot(r2-r1) + target1

R3 = rotate_to(r2-r3, target2-restraint2)
target3 = R3.dot(r2-r3) + restraint2

# print(target3-restraint2)
print(restraint1)
print(target1)
print(target2)
print(target3)
print(restraint2)
print( np.linalg.norm(target3-restraint2) )

coords = [R1.dot(r-r0)+restraint1 for r in dna_coordinates[:idx1]] + \
         [R2.dot(r-r1)+target1 for r in dna_coordinates[idx1:idx2]] + \
         [R3.dot(r-r3)+restraint2 for r in dna_coordinates[idx2:]]
coords = np.array(coords)
print(coords[idx1],coords[idx1+1])
print(coords[idx2],coords[idx2+1])
print(coords[0],restraint1)
print(coords[-1],restraint2)

a.set_splines( np.linspace(0,1,len(coords)), coords )

# coords = np.array((restraint1, midpoint1, midpoint2, restraint2))
# a.set_splines( (0,0.4,0.6,1), coords )

## print RB positions
def print_rb_coords(original, R, about0, about1):
    r = original[:3]
    R0 = original[3:].reshape((3,3))
    r = R.dot(r-about0)+about1
    R = R.dot(R0)
    fh.write( '{} {} {} {} {} {} {} {} {} {} {} {}\n'.format(*r, *R.reshape(9)) )

with open("trombone-ssb.rbcoords.txt",'w') as fh:
    for rR in ssb_coords[:idxidx1]:
        print_rb_coords(rR, R1, r0, restraint1)
    for rR in ssb_coords[idxidx1:idxidx2]:
        print_rb_coords(rR, R2, r1, target1)
    for rR in ssb_coords[idxidx2:]:
        print_rb_coords(rR, R3, r3, restraint2)


# for s in np.linspace(0.25,0.75,5):
#     print( '{} {} {} 1 0 0 0 1 0 0 0 1'.format( *(restraint1 * (1-s) + midpoint1*s) ))
# for s in np.linspace(0.25,0.75,5):
#     print( '{} {} {} 1 0 0 0 1 0 0 0 1'.format( *(restraint2 * (1-s) + midpoint2*s) ))



model = DnaModel([a], dimensions=(3000,3000,3000))
model.timestep = 20e-6

end1 = model.strands[0].children[0].children[0]
end2 = model.strands[0].children[-1].children[0]

end1.add_restraint( (10, restraint1) )
end2.add_restraint( (10, restraint2) )

model.simulate(output_name='trombone-dna',
               num_steps=1e5, output_period=1e4,
               gpu=1
)