B
    b                 @   s  d dl mZ d dlZd dlZd dlZd dlmZ d dlZd dl	Z	d dl
mZ d dlmZ d dlmZ d dlmZ d dlmZ d d	lmZ d d
lmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ dZ dd Z!dd Z"G dd de#Z$G dd de#Z%G d d! d!e#Z&G d"d# d#e#Z'G d$d% d%e(Z)G d&d' d'e#Z*G d(d) d)Z+G d*d+ d+e#Z,d,d- Z-d.d/ Z.d0d1 Z/d:d3d4Z0d5d6 Z1d;d8d9Z2dS )<    )print_functionN)defaultdict)Atom)PO4)PO4_terminal_C3)PO4_terminal_C5)deoxyribose_purine)deoxyribose_purine_terminal_C3)deoxyribose_purine_terminal_C5)deoxyribose_pyrimidine)"deoxyribose_pyrimidine_terminal_C3)"deoxyribose_pyrimidine_terminal_C5)nucleic_acid_bases)nucleic_acid_isobases)ribose_purine)ribose_purine_terminal_C3)ribose_purine_terminal_C5)ribose_pyrimidine)ribose_pyrimidine_terminal_C3)ribose_pyrimidine_terminal_C5)flexz2022.8.1c             C   s"   t jt jt jtdd| S )NlibZ
regressors)ospathjoinabspathdirname__file__)filename r   6lib/python3.7/site-packages/restraintlib/restraints.pyregressor_absolute_path#   s    r!   c          	   C   s$   t t| d}t|S Q R X d S )Nrb)openr!   pickleload)Zfunction_nameZp_filer   r   r    load_function'   s    r&   c               @   sp   e Zd ZdZdd Zdd Zedd Zedd	 Zed
d Z	edd Z
edd Zedd Zdd ZdS )DistanceMeasure)measure_namemeasurerestraint_namesc             C   s   || _ t| || _|| _d S )N)r(   getattrr)   r*   )selfr)   r*   r   r   r    __init__2   s    zDistanceMeasure.__init__c             C   s   | j d t| j S )N )r(   strr*   )r,   r   r   r    __str__7   s    zDistanceMeasure.__str__c             C   sn   t |dkst |dkrd S t |t |kr4tdd}x*t||D ]\}}|| }||| 7 }qDW t|S )Nr   zuneven number of elementsg        )len	Exceptionzipmathsqrt)clsvector1vector2dist_sq_sumabdiffr   r   r    	euclidean:   s    zDistanceMeasure.euclideanc             C   sx   t |dkst |dkrd S t |t |kr4tdd}x4t||D ]&\}}|| }t|}||| 7 }qDW t|S )Nr   zunequal number of elementsg        )r1   r2   r3   ConditionItemfix_torsionr4   r5   )r6   r7   r8   r9   r:   r;   r<   r   r   r    euclidean_anglesG   s    
z DistanceMeasure.euclidean_anglesc             C   s&   ||j d  }||j d  }||S )Nr      )
atom_namesdist)r6   	restraintatomsatom0atom1r   r   r    
atoms_distU   s    zDistanceMeasure.atoms_distc             C   s6   ||j d  }||j d  }||j d  }|||S )Nr   rA      )rB   Zangle)r6   rD   rE   rF   rG   atom2r   r   r    atoms_angle[   s    zDistanceMeasure.atoms_anglec             C   sT   ||j d  }||j d  }||j d  }||j d  }||||}t|}|S )Nr   rA   rI      )rB   torsionr>   r?   )r6   rD   rE   rF   rG   rJ   atom3rM   r   r   r    atoms_torsionb   s    
zDistanceMeasure.atoms_torsionc             C   s   ||j d  }||j d  }||j d  }||j d  }||j d  }dddddg}t|||||d< t|||||d< t|||||d< t|||||d< t|||||d< t|\}	}
}}|	S )Nr   rA   rI   rL      g        )rB   r>   r?   rM   _pseudorotation_with_sd)r6   rD   rE   C1C2C3C4O4thetapsd_ptmsd_tmr   r   r    atoms_pseudorotationo   s    z$DistanceMeasure.atoms_pseudorotationc             C   s^   g }g }xH|D ]@}|j | jkr||| t| d|j ||}|| qW | ||S )Nzatoms_%s)namer*   appendvaluer+   typer)   )r,   conditional_restraintrE   r7   r8   rD   r_   r   r   r    distance   s    
zDistanceMeasure.distanceN)__name__
__module____qualname__	__slots__r-   r0   classmethodr=   r@   rH   rK   rO   r\   rb   r   r   r   r    r'   ,   s   r'   c               @   sL   e Zd ZdZdddZdd Zdd Zd	d
 ZdddZdd Z	dd Z
dS )ConditionalRestraintItem)
r`   r]   rB   _value_sigma
_regressorvalue_param_namevalue_param_atoms
value_dist
sigma_distNc
       
      C   sx   || _ || _|| _|| _|| _d | _|r2t|| _d | _d | _|	rh|	d | _|	d | _| jdkrht	d|| _
|| _d S )Nr   rA   )tau_maxtorsion_chiz-Unknown parameter for functional relationship)r`   r]   rB   ri   rj   rk   r&   rl   rm   r2   rn   ro   )
r,   restraint_typer]   rB   r_   sigmarn   ro   Zvalue_function_nameZvalue_param_defr   r   r    r-      s"    



z!ConditionalRestraintItem.__init__c             C   s  | j dkr|d }|d }|d }|d }|d }dddddg}t|||||d< t|||||d< t|||||d< t|||||d< t|||||d< t|\}}	}
}|
S | j dkr|d }|d }|d }|d }t||||}|S d S )	Nrp   r   rA   rI   rL   rP   g        rq   )rl   r>   r?   rM   rQ   )r,   param_atomsrR   rS   rT   rU   rV   rW   rX   rY   rZ   r[   rG   rJ   rN   Zatom4Zchir   r   r    
calc_param   s*    
z#ConditionalRestraintItem.calc_paramc                s    fdd|D S )Nc                s   g | ]} | qS r   r   ).0	atom_name)atom_mapr   r    
<listcomp>   s    z6ConditionalRestraintItem.map_atoms.<locals>.<listcomp>r   )r,   rx   rB   r   )rx   r    	map_atoms   s    z"ConditionalRestraintItem.map_atomsc             C   sR   | j rF| || j }| |}| jj|ggdd\}}|d |d fS | j| jfS )NT)Z
return_stdr   )rm   rz   ru   rk   Zpredictri   rj   )r,   rx   rt   Zregression_paramr_   rs   r   r   r    value_sigma   s    
z$ConditionalRestraintItem.value_sigmac          	   C   s<   |  || j}| |\}}t| j|||| j| j| j|jS )N)rz   rB   r{   	Restraintr`   rn   ro   r]   )r,   rx   ra   rE   Zrestraint_valueZrestraint_sigmar   r   r    get_restraint   s    z&ConditionalRestraintItem.get_restraintc             C   s   |  |d S )Nr   )r{   )r,   rx   r   r   r    r_      s    zConditionalRestraintItem.valuec             C   s   |  |d S )NrA   )r{   )r,   rx   r   r   r    rs      s    zConditionalRestraintItem.sigma)NNNN)N)rc   rd   re   rf   r-   ru   rz   r{   r}   r_   rs   r   r   r   r    rh      s   

rh   c               @   sX   e Zd ZdZdd ZdddZdddZed	d
 Zedd Z	edd Z
dd ZdS )r>   )
multiplierr`   r]   rB   ri   rj   c             C   s>   |dkrt d|| _|| _|| _| || _|| _d| _d S )N)rM   pseudorotationzUnknown condition typerP   )r2   r`   r]   rB   r?   ri   rj   r~   )r,   Zcondition_typer]   rB   r_   rs   r   r   r    r-      s    zConditionItem.__init__Nc             C   s   | j S )N)ri   )r,   rx   r   r   r    r_      s    zConditionItem.valuec             C   s   | j S )N)rj   )r,   rx   r   r   r    rs      s    zConditionItem.sigmac             C   s\   |dkr,|t |d d t |d  d  S |dk rX|t |d d t |d  d  S |S )z
        Normalize torsion angle to be between [-180,180]
        :param value: torsion angle in deg
        :return:  torsion angle in deg  between [-180,180]
           rI   ih  iL)int)r6   r_   r   r   r    r?      s
    $$zConditionItem.fix_torsionc             C   sL  |d |d |d |d |d g}d}d}xFt |D ]:\}}dtj | }||t| 7 }||t| 7 }q4W tt| |}|dk r|d7 }t|}	d	t|	| t|	|   }
d}dddddg}xLt |D ]@\}}|
t|	dtj |   ||< |||  }||| 7 }qW td	| d
 }|t|
 }|||
|fS )z
        Calculate pseudorotation angle
        :param theta: list of theta torsion values, for example theta[0] = torsion(C4', O4' C1', C2')
        :return: P in deg, standard deviation of P, Tm, standard deviation of Tm
        rI   rL   rP   r   rA   g        g?g     v@g?g      @)		enumerater4   Zpisincosdegreesatan2radiansr5   )r6   rW   Z_thetaZsum_sinZsum_cosZi_ttxZP_degP_radTmZSTZThcdZsd_TmZsd_Pr   r   r    rQ   
  s*    "
  z%ConditionItem._pseudorotation_with_sdc             C   sx   |d |d  |d |d   }d|d  t t dt t d  }t ||}|d t | }t ||fS )	a  
        Calculate pseudorotation angle based on
        http://dl.taq.ir/science/principles_of_Nucleic_Acid_Structure_Neidle.pdf page 29
        :param theta: list of theta torsion values, for example theta[0] = torsion(C4', O4' C1', C2')
        :return: P in deg, Tm
        rP   rA   rL   r   g       @rI   g      B@g      R@)r4   r   r   r   r   r   )r6   rW   ZP_tan_numeratorZP_tan_denominatorr   r   r   r   r    _pseudorotation0  s
     ,zConditionItem._pseudorotationc             C   sb  | j dkrt| jdkr|| jd  }|| jd  }|| jd  }|| jd  }||||}| |}|  | j|    }|  | j|    }x,dD ]$}	|||	   kr|krn qdS qW d	S td
n| j dkrVt| jdkrN| jdddddgkrtd|| jd  }
|| jd  }|| jd  }|| jd  }|| jd  }dddddg}| |||
||d< | ||
|||d< | |
||||d< | |||||d< | ||||
|d< | 	|\}}}}|  | j|    }|  | j|    }x2dD ]*}	|||	   kr<|krn ndS qW d	S tdtdd S )NrM   rP   r   rA   rI   rL   )r   ih  iTFz+Wrong number of atoms for torsion conditionr      zC1'zC2'zC3'zC4'zO4'z1Wrong order of atoms for pseudorotation conditiong        z2Wrong number of atoms for pseudorotation conditionzUnknown condition type)
r`   r1   rB   rM   r?   r_   r~   rs   r2   rQ   )r,   rE   rF   rG   rJ   rN   rM   Zlow_bondaryZhigh_bondaryr   rR   rS   rT   rU   rV   rW   rX   rY   rZ   r[   r   r   r    check_condition>  sN    


	
 
zConditionItem.check_condition)N)N)rc   rd   re   rf   r-   r_   rs   rg   r?   rQ   r   r   r   r   r   r    r>      s   


&r>   c               @   s0   e Zd ZdZdd Zdd Zdd Zdd	 Zd
S )ConditionalRestraint)r]   
conditions
restraintsc                sB   dd  dd || _  fdd|D | _fdd|D | _d S )Nc             S   s   t | tr| S t|  S )N)
isinstancer>   )	conditionr   r   r    <lambda>      z/ConditionalRestraint.__init__.<locals>.<lambda>c             S   s   t | tr| S t|  S )N)r   rh   )rD   r   r   r    r     r   c                s   g | ]} |qS r   r   )rv   Zcon)create_conditionr   r    ry     s    z1ConditionalRestraint.__init__.<locals>.<listcomp>c                s   g | ]} |qS r   r   )rv   res)create_restraintr   r    ry     s    )r]   r   r   )r,   r]   r   r   r   )r   r   r    r-     s
    zConditionalRestraint.__init__c             C   s    t | jdks| jdkrdS dS )Nr   defaultTF)r1   r   r]   )r,   r   r   r    
is_default  s    zConditionalRestraint.is_defaultc             C   s2   |   rdS x | jD ]}||dkrdS qW dS )NTF)r   r   r   )r,   rE   r   r   r   r    check_conditions  s    z%ConditionalRestraint.check_conditionsc             C   s*   g }x | j D ]}||||  qW |S )N)r   r^   r}   )r,   rE   Zprintable_restraintsrD   r   r   r    get_restraints  s    z#ConditionalRestraint.get_restraintsN)rc   rd   re   rf   r-   r   r   r   r   r   r   r    r   y  s
   r   c                   s^   e Zd Zd fdd	Zdd Zdd Zdd	 Zd
d Zdd Zdd Z	dd Z
dd Z  ZS )ConditionalRestraintListr   c                s6   t t|   dd }x|D ]}| || qW d S )Nc             S   s   t | tr| S tf | S )N)r   r   )objr   r   r    r     s   z3ConditionalRestraintList.__init__.<locals>.<lambda>)superr   r-   r^   )r,   dataZcreate_conditional_restraintitem)	__class__r   r    r-     s    
z!ConditionalRestraintList.__init__c             C   s   t t| ||S )N)r   list__getslice__)r,   ijr   r   r    r     s    z%ConditionalRestraintList.__getslice__c             C   s6   t g }x(t| D ]}||dkr|| qW |S )NT)r   iterr   r^   )r,   rE   Zfeasiblera   r   r   r    get_feasible  s
    z%ConditionalRestraintList.get_feasiblec             C   s   t dd | D S )Nc             S   s   g | ]}|  qS r   )r   )rv   ra   r   r   r    ry     s    z8ConditionalRestraintList.any_default.<locals>.<listcomp>)any)r,   r   r   r    any_default  s    z$ConditionalRestraintList.any_defaultc             C   sL   g }x&t | D ]\}}| r|| qW xt|D ]}| | q6W d S )N)r   r   r^   reversedpop)r,   Zi_to_deleter   ra   r   r   r    remove_default  s    z'ConditionalRestraintList.remove_defaultc             C   s*   d }x t | D ]\}}| r|S qW |S )N)r   r   )r,   r   r   ra   r   r   r    get_default  s
    z$ConditionalRestraintList.get_defaultc       	      C   s   t | dkrd S |  r,t | dkr,|   d}d }x:t| D ].\}}|t|||}|r>||k r>|}|}q>W |d k	r| | S d S )Nr   rA   g  dA)r1   r   r   r   rb   r+   )	r,   rE   distance_measurevariableZmin_distanceZmin_distance_ir   ra   rb   r   r   r    find_closest  s    z%ConditionalRestraintList.find_closestc             C   s   |  ||dS )Nr   )r   )r,   rE   r   r   r   r    find_restraint_closest  s    z/ConditionalRestraintList.find_restraint_closestc             C   s   |  ||dS )Nr   )r   )r,   rE   r   r   r   r    find_condition_closest  s    z/ConditionalRestraintList.find_condition_closest)r   )rc   rd   re   r-   r   r   r   r   r   r   r   r   __classcell__r   r   )r   r    r     s   r   c               @   s   e Zd ZdZdddZdS )r|   )r`   rE   r_   rs   rn   ro   r]   condition_nameNc	       	      C   s\   || _ || _|| _t| jtr&td|| _t| jtr@td|| _|| _|| _	|| _
d S )NzValue should be float type)r`   rE   r_   r   r/   r2   rs   rn   ro   r]   r   )	r,   rr   rE   r_   rs   rn   ro   r]   r   r   r   r    r-     s    zRestraint.__init__)NNNN)rc   rd   re   rf   r-   r   r   r   r    r|     s   
r|   c               @   sH   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d Zdd Z	dd Z
dS )AtomGroupCache)_atoms_neighboursc             C   s   g | _ g | _d S )N)r   r   )r,   r   r   r    r-     s    zAtomGroupCache.__init__c             C   s   | j | d S )N)r   r^   )r,   Zatom_idxr   r   r    add_atom  s    zAtomGroupCache.add_atomc             C   s   | j S )N)r   )r,   r   r   r    
iter_atoms	  s    zAtomGroupCache.iter_atomsc             C   s6   x0| j D ]&}|j|kr|j|kr|j|kr|S qW d S )N)r   rw   res_idchain_id)r,   r   rw   r   atomr   r   r    	find_atom  s    zAtomGroupCache.find_atomc             C   s   | j | d S )N)r   r^   )r,   Z	neighbourr   r   r    add_neighbour  s    zAtomGroupCache.add_neighbourc             C   s   | j S )N)r   )r,   r   r   r    iter_neighbour  s    zAtomGroupCache.iter_neighbourc             C   sd   x^| j D ]T}|j|kr|j|kr|j|kr|j|ksX|jdkrF|dksX|dkr|jdkr|S qW d S )N )r   rw   r   r   alt_loc)r,   r   rw   r   r   r   r   r   r    find_neighbour  s    


.zAtomGroupCache.find_neighbourN)rc   rd   re   rf   r-   r   r   r   r   r   r   r   r   r   r    r     s   r   c               @   s   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d Zdd Z	dd Z
dd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zd(d!d"Zd)d#d$Zd*d%d&Zd'S )+MonomerRestraintGroup)r]   	res_namesatom_labels_valid_atom_labelsres_numbersr   disallowed_conditionsr   condition_measurer   rE   groups_registered_res_id_res_id_pos_pos_res_idc
       
      C   s   || _ || _|| _t| | _|| _|| _|| _t	f || _
t	f |	| _t|tr\|| _n
t|| _g | _tt| _tt| _tt| _tt| _d S )N)r]   r   r   r   keysr   r   r   r   r'   r   r   r   r   r   rE   r   r   r   dictr   r   setr   )
r,   r]   r   r   r   r   r   r   r   r   r   r   r    r-   8  s"    




zMonomerRestraintGroup.__init__c             C   sB   |  }|| j| kr>t| j| | j| |< | j| | d S )N)stripr   r1   r   r^   )r,   r   r   	_chain_idr   r   r    register_res_id\  s    z%MonomerRestraintGroup.register_res_idc             C   s4   |  }| j| | }|dkr0| j| |d  S d S )Nr   rA   )r   r   r   )r,   r   r   r   posr   r   r    prev_res_idb  s
    z!MonomerRestraintGroup.prev_res_idc             C   sB   |  }| j| | }|t| j| d k r>| j| |d  S d S )NrA   )r   r   r1   r   )r,   r   r   r   r   r   r   r    next_res_idi  s
    z!MonomerRestraintGroup.next_res_idc             C   s   | j |  | d S )N)r   r   add)r,   r   r   r   r   r    register_valid_res_idp  s    z+MonomerRestraintGroup.register_valid_res_idc             C   s@   |  }| j|t }||kp>| |||kp>| |||kS )N)r   r   getr   r   r   )r,   r   r   r   registered_res_idr   r   r    !is_registered_res_id_or_neighbours  s
    z7MonomerRestraintGroup.is_registered_res_id_or_neighbourc             C   s   |  | jkS )N)r   r   )r,   res_namer   r   r    is_valid_res_name|  s    z'MonomerRestraintGroup.is_valid_res_namec             C   s   |dkS )N)ACGTUZDAZDCZDGZDTZDUZICZIGr   )r,   r   r   r   r    is_standard_res_name  s    z*MonomerRestraintGroup.is_standard_res_namec             C   s   |  | jkS )N)r   r   )r,   rw   r   r   r    is_valid_atom_name  s    z(MonomerRestraintGroup.is_valid_atom_namec             C   s|   xv| j D ]l\}}}}}	| |||}
| |||	}||||
|}|||||}|d k	r|d k	r|||krdS qW dS )NTF)r   relative_res_id_for_atomr   rC   )r,   groupr   r   alt_idatom_name_1atom_name_2rC   res_id_mod_1res_id_mod_2res_id_1res_id_2atom_1atom_2r   r   r    any_atom_in_disallowed_atoms  s    z2MonomerRestraintGroup.any_atom_in_disallowed_atomsc       	      C   s0   |  ||r,t|||||||}| j| d S )N)r   r   rE   r^   )	r,   r   r   r   rw   r   atom_xyzi_seqr   r   r   r    r     s    zMonomerRestraintGroup.add_atomc             C   s8   |dkr|S |dkr |  ||S |dkr4| ||S d S )Nr   rA   )r   r   )r,   r   r   Zres_id_deltar   r   r    r     s    z.MonomerRestraintGroup.relative_res_id_for_atomc          
   C   s  t t}xt| jD ]\}}x|D ]}xt| jD ]\}}d||}||jkr|j	| j
kr|| |j|j| j
|j	  kr|| | ||jkr4||jks|| |j|jks|| |j|jkr4| |jr4|| | q4W q$W qW xt|D ]\}}| }	tdd |	D }
|
d t|
dkrxh|
D ]T}d||}x2|	D ]*}|j}||dfkrP| j| | qPW |j| j| _q:W q|| j|< qW d S )Nz{}_{}c             S   s   g | ]
}|j qS r   )r   )rv   r   r   r   r    ry     s    z;MonomerRestraintGroup.create_res_groups.<locals>.<listcomp>r   r   )r   r   six	iteritemsr   r   rE   formatr   rw   r   r   r   r   r   r   r   r   r   r   r   discardr1   r   r   r   )r,   Zpreliminary_groupsr   r   r   Zi_atomr   keyZpreliminary_groupZp_group_atomsZlocslocZkey_altZatom_alt_locr   r   r    create_res_groups  s8    

 




z'MonomerRestraintGroup.create_res_groupsc             C   sl   t d| j xNt| jD ]>\}}x4| D ](}t | j||j|j|j	|j
|j|j q0W qW t d d S )Nz# PRINT GROUPS {}z# PRINT GROUPS finished)printr   r]   r   r   r   r   r   r   r   rw   r   r   )r,   r   r   r   r   r   r    _print_groups  s
    ,z#MonomerRestraintGroup._print_groupsc             C   s>  | d}|d  }|d }t|dk r.dn
|d  }| ||||rNdS x| jD ]\}}	}
}}| |||}| |||}|||||}|||	||}|d ks|d ks|||
krV|dkr2td	| j
|||d ks|d krd	n|||r|jnd|r
|jnd||r|jnd|r(|jnd|	
 dS qVW d
S )N_r   rA   rL   r   rI   Fz>{} {} {}: non valid dist: {}, atom1: {} {} {}, atom2: {} {} {}zat least one missingT)splitr   r1   r   r   r   r   rC   r   r   r]   r   r   )r,   	group_keyr   verboseZ	split_keyr   r   r   r   r   rC   r   r   r   r   r   r   r   r   r    is_valid_atom_group  s4    

z)MonomerRestraintGroup.is_valid_atom_groupr   c             C   s   |dkr t d t d| j xdtt| jD ]P\}}| |||sh| j|= |dkrt d| q2|dkr2t d|| j q2W |dkrt d  dS )z>
        Deletes groups that are not bonded correctly
        r   z<############################################################z # SEARCHING for {} group/monomerz#     {} ignoringz#     {} recognized as {}z# SEARCHING finishedN)r   r   r]   sortedr   r   r   r  )r,   r  r  r   r   r   r    validate_links  s    z$MonomerRestraintGroup.validate_linksc          	   C   s  g }|dkrt d| j xLtt| jD ]6\}}dd | D }|dkrpt d|| jt|	  | j
|}|dkrt d|| jdd	 |D  ||| j}|d k	r|dkrt d
|| j|j n|dkrt d| | j
 }|d kr2| j
|| j}|dkrXt d| j|j n&|dkrXt d t d| j|j ||| q0W |dkrt d t d |S )Nr   z# ANALYZING {} group/monomerc             S   s   i | ]}||j qS r   )rw   )rv   r   r   r   r    
<dictcomp>  s    z9MonomerRestraintGroup.atom_restraints.<locals>.<dictcomp>rA   z#     {} atoms for {} {}z&#     {} feasible_restraints for {} {}c             S   s   g | ]
}|j qS r   )r]   )rv   Zfcr   r   r    ry     s    z9MonomerRestraintGroup.atom_restraints.<locals>.<listcomp>z#     {} recognized as {} {}z#     {} not feasiblez#             closest to {} {}z#             default usedz#             set default {} {}z# ANALYZING finishedz<############################################################)r   r   r]   r  r   r   r   r   r   r   r   r   r   r   r   r   r   extendr   )r,   r  Zresult_restraintsZatom_group_key
atom_grouprE   Zfeasible_restraintsZclosest_restraintr   r   r    atom_restraints  s<    




z%MonomerRestraintGroup.atom_restraintsc             C   s4   |    |dkr|   | j|d | j|d}|S )Nr   )r  )r   r   r  r	  )r,   r  r   r   r   r    prepare_restraints/  s    z(MonomerRestraintGroup.prepare_restraintsN)r   )r   )r   )rc   rd   re   rf   r-   r   r   r   r   r   r   r   r   r   r   r   r   r   r  r  r	  r
  r   r   r   r    r   $  s&   $	,!

)r   c             C   s   |j | |td || | t|dkr6td| d nJ| }|dkrRt|| d td|| d | }|dkrt|| d d S )N)versionr   zH# There were no restraints to be created based on the submitted PDB file)filer   
)Z	save_infoVERSIONZsave_input_filenamer1   r   headerr   footer)streamrestraint_groupsZrestraint_text_allprinterin_filenamer  r  r   r   r    save9  s    r  c             C   sT  t t}t t}t t}x| D ]}x|jD ]~}d|j|j|j|jrLd|j ndf }|j}|| | d||f }|| |j	|jrd|j nd  || |j
 q*W qW xt| D ]}t| xt|| D ]p}	d||	f }|| }
|| }td|	  tdt|
  tdtd	d
 |D   tdtdd
 |D   qW qW d S )Nz%s_%05d_%s%sz_%sr   z%s_%sz^%sz	%sz		Atoms: %sz
		Dist: %sc             s   s   | ]}| d r|V  qdS )r   N)
startswith)rv   r   r   r   r    	<genexpr>b  s    z.print_info_about_restraints.<locals>.<genexpr>z		Angles: %sc             s   s   | ]}| d r|V  qdS )r:   N)r  )rv   r   r   r   r    r  c  s    )r   r   rE   r   r   r   r   r   r   rw   r]   r  r   r   )all_restraintsZres_id_name_mapZ
atom_cacheZrestraint_cacherD   r   r   r_   Z	cache_keyvalrE   restr   r   r    print_info_about_restraintsK  s*    
&$r  c             C   s  | |  x|  D ]x}xr| D ]f}x`| D ]T}xN| D ]B}x<|D ]4}||j|  ||j	rH|
|j|  qHW q>W q0W q"W qW x|  D ]}x| D ]}x| D ]v}xp| D ]d}x^|D ]V}||j| r|j }	x4| D ](}
||j| |j	|
j|	|
j|
j qW qW qW qW qW qW g }x|D ]}||  qBW |||}t|tkst|tjkrt|dkr|gS |S )Nr   )ZvalidateZmodelsZchainsZresidue_groupsZatom_groupsr   idZresidr   Zresnamer   r   altlocr   rE   r   r]   Zxyzr   r  r
  Zprint_restraintsr`   r/   r   	text_typer1   )pdb_hierarchyr  allowed_restraint_groupsr  ZmodelchainZresidue_groupr  rD   r  r   r  Zrestraint_groupZrestraint_textr   r   r    analyze_pdb_hierarhyf  s4    

&

8
,r"  Fc          	   C   s   g }|dkr*|  ds |  dr*tdnJ|r2|n| }|rPtjjt| |dntjj| d}	|	 }
t|
|||}t	|t
kst	|tjkrt|d}t||||| W d Q R X nt||||| d S )NFz.resz.insz Shelx format files not supported)linessource_info)	file_namew)endswithr2   iotbxZpdbinputr   Zsplit_linesZconstruct_hierarchyr"  r`   r/   r   r  r#   r  )Zin_pdbr  r   Zout_filenamer  r#  r$  Zprinted_restraints_listr  Zdata_pdbr  Zres_filer   r   r    	parse_pdb  s    
(r*  c             C   sj   t |t| d| t| d| t| d| t| d| t| d| i t| d| t| d| t| d| 	S )	Nz%s_PDB_CODESz%s_ATOM_NAMESz%s_ATOM_RESz%s_REQUIRED_CONDITIONz%s_DISALLOWED_CONDITIONz%s_RESTRAINTSz%s_DISTANCE_MEASUREz%s_CONDITION_DISTANCE_MEASURE)r   r+   )moduleprefixr]   r   r   r    create_monomer_group  s    r-  Tc             C   s  g }| r| ttdd |rB| ttdd | ttdd |rfxdD ]}| tt|| qLW |rxdD ]}| tt|| qpW |rttt	t
g}dd	d
dg}	xBt||	D ]4\}
}x*dD ]"}d||}| t|
|| qW qW |rttttg}ddddg}	xFt||	D ]8\}
}x,dD ]$}d||}| t|
|| q"W qW ttttg}ddddg}	xFt||	D ]8\}
}x,dD ]$}d||}| t|
|| qW qtW |S )Nr   ZPO4_5_TERMINALr   ZPO4_3_TERMINALr   )ZADENINEZGUANINEZURACILZTHYMINEZCYTOSINE)ZISOCYTOSINEZ
ISOGUANINEZDEOXYRIBOSE_PURINEZDEOXYRIBOSE_PYRIMIDINEZRIBOSE_PURINEZRIBOSE_PYRIMIDINE)
CHI_CONFORMATIONGAMMACONFORMATIONBASE_FUNC_OF_TORSION_CHIALL_FUNC_OF_TORSION_CHI"SUGAR_CONFORMATION_FUNC_OF_TAU_MAXCHI	CHI_GAMMAALLSUGARz{}_{}ZDEOXYRIBOSE_PURINE_TERMINAL_C3Z"DEOXYRIBOSE_PYRIMIDINE_TERMINAL_C3ZRIBOSE_PURINE_TERMINAL_C3ZRIBOSE_PYRIMIDINE_TERMINAL_C3)r.  r/  r0  r1  r2  r3  r4  r5  r6  r7  ZSUGAR_CONFORMATIONZDEOXYRIBOSE_PURINE_TERMINAL_C5Z"DEOXYRIBOSE_PYRIMIDINE_TERMINAL_C5ZRIBOSE_PURINE_TERMINAL_C5ZRIBOSE_PYRIMIDINE_TERMINAL_C5)r^   r-  r   r   r   r   r   r   r   r   r   r3   r   r	   r   r   r   r
   r   r   r   )Zpo4Zpo4terminalbasesZisobasesZribose_deoxyriboseZribose_deoxyribose_terminalZrestraint_listr,  modulesprefixesr+  r   Zprefix_groupr   r   r    load_restraints_lib  sJ    




r;  )FN)TTTTTT)3Z
__future__r   Z	iotbx.pdbr(  r4   r   collectionsr   r$   r   Zrestraintlib.atomr   Zrestraintlib.libr   r   r   r   r	   r
   r   r   r   r   r   r   r   r   r   r   r   Zcctbx.array_familyr   r  r!   r&   objectr'   rh   r>   r   r   r   r|   r   r   r  r  r"  r*  r-  r;  r   r   r   r    <module>   s\   `V $B)  &
 