import tooba_f as tbf

#P1=(0,1,0)
#P2=(1,1,0)
#print(tbf.angle_between3D(P1,P2))

#Read .trr file
data_all=tbf.fr_export(trajfile='traj.trr',num_frames=1)
#Read .gro file
_,data_num,_,res_num,res_type,atom_type,atom_num,_ = tbf.read_gro('initial.gro')
#Create subdomains coordinates
box_p=tbf.domain_decomposition(data=data_all,dx=2,dy=2,dz=2)

###################################################
HEAD=False
if HEAD==True:
    #Find atom type index in lists created above
    group_ndx=tbf.atomid_data(res_num, atom_type, atom_num, group={'MMA':['C1', 'C2']})
    #Assign desired atoms (from above function) to subdomains
    ##result1: {step:{res:{atom_type:{atom_num:(subX,subYsubZ)}}}}
    ##result2: {step:{res:{atom_type:{(subX,subYsubZ):[atom_num]}}}}
    _,box_res=tbf.atom2grid(data_all,box_p, group_ndx)
    #Creates dictionary with coordinates per subdomain for each frame
    coord_norm,_=tbf.sub_coord(box=box_res, data=data_all, res_num=res_num)
    #Creates dictionary with c, normal per subdomain for each frame
    surf=tbf.coord2norm(coord_norm,img=True)

###################################################
TAIL=True
if TAIL==True:
    #Find atom type index in lists created above
    group_ndx=tbf.atomid_data(res_num, atom_type, atom_num, group={'MMA':['C3', 'C4', 'C5']})
    #Assign desired atoms (from above function) to subdomains
    ##result1: {step:{res:{atom_type:{atom_num:(subX,subYsubZ)}}}}
    ##result2: {step:{res:{atom_type:{(subX,subYsubZ):[atom_num]}}}}
    _,box_res=tbf.atom2grid(data_all,box_p, group_ndx)
    #Creates dictionary with coordinates per subdomain for each frame
    _,coord_vector=tbf.sub_coord(box=box_res, data=data_all, res_num=res_num)
    vector=tbf.coord2vector(coord_vector)