import os
import tooba_f as tbf
###################################################
#NOTICE:    resids of head in each subdomain may differ in tail case
#           keep all atoms of group in the first occurent subdomain
#           in case of tail is the one closest to the head, hence
#           the code is a good approximation 
###################################################
TRAJ='system/eq_traj.trr'
NUM_FR=1
GRO='system/eq_final.gro'
HEAD=True
HD_GROUP={'NS':['C6', 'Na', 'P4', 'P3', 'C7'], 'CHOL':['ROH'], 'FFA':['AC']}
TAIL=False
TL_GROUP={'NS':['C3', 'C4', 'C5', 'C8', 'C9', 'C10'], 'CHOL':['R1', 'R2', 'R3', 'R4', 'R5'], 'FFA':['C1', 'C2', 'C3', 'C4']}
DISCET=[3, 3, 3]
RES_NAME='time_domain_c-normal-cg'
###################################################
###################################################
print(' ')
print('================')
print('Starting process')
print('================')
###################################################
###################################################
#Read .gro file
_,data_num,_,res_num,res_type,atom_type,atom_num,_ = tbf.read_gro(GRO)
print(' ')
###################################################
if os.path.isfile('./data.pkl'):
    print('WARNING: Preprocessing files exist.')
    print('         Erase data.pkl if the system is new.')
    print('--------------------------------------------')
    data_all=tbf.frompickle('./data.pkl')
else:
    #Read .trr file
    data_all=tbf.fr_export(trajfile=TRAJ,num_frames=NUM_FR)
    tbf.topickle(fl=data_all, sv_name='data')
###################################################
if HEAD==True:
    if os.path.isfile('./ndx_HEAD.pkl'): 
        print('WARNING: Preprocessing files exist.')
        print('         Erase ndx_HEAD.pkl if the system is new.')
        print('--------------------------------------------')
        group_ndx=tbf.frompickle('./ndx_HEAD.pkl')
    else:
        #Find atom type index in lists created above
        group_ndx=tbf.atomid_data(res_num, res_type, atom_type, atom_num, group=HD_GROUP)
        tbf.topickle(fl=group_ndx, sv_name='ndx_HEAD')
###################################################
if os.path.isfile('./box.pkl'):
    print('WARNING: Preprocessing files exist.')
    print('         Erase box.pkl if the system is new')
    print('         or new grid is applied !')
    print('--------------------------------------------')
    box_res=tbf.frompickle('./box.pkl')
else:
    #Create subdomains coordinates
    box_p=tbf.domain_decomposition(data=data_all,dx=DISCET[0],dy=DISCET[1],dz=DISCET[2])
    #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)
    tbf.topickle(fl=box_res, sv_name='box')
###################################################
###################################################
if HEAD==True:
    #Creates dictionary with coordinates per subdomain for each frame
    coord_norm,coord_vector=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.coord2norm2cg(coord_vector,img=True)
    print(surf)
    exit()
    tbf.topickle(fl=surf, sv_name=RES_NAME)
    tbf.tojson(fl=surf, sv_name=RES_NAME)
###################################################
if TAIL==True:
    #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)
###################################################
if HEAD==True and TAIL==True:
    tbf.SurfVector_angle(surf,vector)