segmentmodel.py

import pdb
import numpy as np
import random
from .model.arbdmodel import PointParticle, ParticleType, Group, ArbdModel
from .coords import rotationAboutAxis, quaternion_from_matrix, quaternion_to_matrix
from .model.nonbonded import *
from copy import copy, deepcopy
from .model.nbPot import nbDnaScheme

from scipy.special import erf
import scipy.optimize as opt
from scipy import interpolate

from .model.CanonicalNucleotideAtoms import canonicalNtFwd, canonicalNtRev, seqComplement
from .model.CanonicalNucleotideAtoms import enmTemplateHC, enmTemplateSQ, enmCorrectionsHC

# import pdb
"""
TODO:
 + fix handling of crossovers for atomic representation
 + map to atomic representation
    + add nicks
    + transform ssDNA nucleotides 
    - shrink ssDNA
    + shrink dsDNA backbone
    + make orientation continuous
    + sequence
    + handle circular dna
 + ensure crossover bead potentials aren't applied twice 
 + remove performance bottlenecks
 - test for large systems
 + assign sequence
 + ENM
 - rework Location class 
 - remove recursive calls
 - document
 - develop unit test suite
"""
class CircularDnaError(Exception):
    pass

class ParticleNotConnectedError(Exception):
    pass

class Location():
    """ Site for connection within an object """
    def __init__(self, container, address, type_, on_fwd_strand = True):
        ## TODO: remove cyclic references(?)
        self.container = container
        self.address = address  # represents position along contour length in segment
        # assert( type_ in ("end3","end5") ) # TODO remove or make conditional
        self.on_fwd_strand = on_fwd_strand
        self.type_ = type_
        self.particle = None
        self.connection = None
        self.is_3prime_side_of_connection = None

        self.prev_in_strand = None
        self.next_in_strand = None
        
        self.combine = None     # some locations might be combined in bead model 

    def get_connected_location(self):
        if self.connection is None:
            return None
        else:
            return self.connection.other(self)

    def set_connection(self, connection, is_3prime_side_of_connection):
        self.connection = connection # TODO weakref? 
        self.is_3prime_side_of_connection = is_3prime_side_of_connection

    def get_nt_pos(self):
        try:
            pos = self.container.contour_to_nt_pos(self.address, round_nt=True)
        except:
            if self.address == 0:
                pos = 0
            elif self.address == 1:
                pos = self.container.num_nt-1
            else:
                raise
        return pos

    def __repr__(self):
        if self.on_fwd_strand:
            on_fwd = "on_fwd_strand"
        else:
            on_fwd = "on_rev_strand"
        return "<Location {}.{}[{:.2f},{:d}]>".format( self.container.name, self.type_, self.address, self.on_fwd_strand)
        
class Connection():
    """ Abstract base class for connection between two elements """
    def __init__(self, A, B, type_ = None):
        assert( isinstance(A,Location) )
        assert( isinstance(B,Location) )
        self.A = A
        self.B = B
        self.type_ = type_
        
    def other(self, location):
        if location is self.A:
            return self.B
        elif location is self.B:
            return self.A
        else:
            raise Exception("OutOfBoundsError")

    def delete(self):
        self.A.container.connections.remove(self)
        self.B.container.connections.remove(self)
        self.A.connection = None
        self.B.connection = None

    def __repr__(self):
        return "<Connection {}--{}--{}]>".format( self.A, self.type_, self.B )
        

        
# class ConnectableElement(Transformable):
class ConnectableElement():
    """ Abstract base class """
    def __init__(self, connection_locations=None, connections=None):
        if connection_locations is None: connection_locations = []
        if connections is None: connections = []

        ## TODO decide on names
        self.locations = self.connection_locations = connection_locations
        self.connections = connections

    def get_locations(self, type_=None, exclude=()):
        locs = [l for l in self.connection_locations if (type_ is None or l.type_ == type_) and l.type_ not in exclude]
        counter = dict()
        for l in locs:
            if l in counter:
                counter[l] += 1
            else:
                counter[l] = 1
        assert( np.all( [counter[l] == 1 for l in locs] ) )
        return locs

    def get_location_at(self, address, on_fwd_strand=True, new_type="crossover"):
        loc = None
        if (self.num_nt == 1):
            # import pdb
            # pdb.set_trace()
            ## Assumes that intrahelical connections have been made before crossovers
            for l in self.locations:
                if l.on_fwd_strand == on_fwd_strand and l.connection is None:
                    assert(loc is None)
                    loc = l
            # assert( loc is not None )
        else:
            for l in self.locations:
                if l.address == address and l.on_fwd_strand == on_fwd_strand:
                    assert(loc is None)
                    loc = l
        if loc is None:
            loc = Location( self, address=address, type_=new_type, on_fwd_strand=on_fwd_strand )
        return loc

    def get_connections_and_locations(self, connection_type=None, exclude=()):
        """ Returns a list with each entry of the form:
            connection, location_in_self, location_in_other """
        type_ = connection_type
        ret = []
        for c in self.connections:
            if (type_ is None or c.type_ == type_) and c.type_ not in exclude:
                if   c.A.container is self:
                    ret.append( [c, c.A, c.B] )
                elif c.B.container is self:
                    ret.append( [c, c.B, c.A] )
                else:
                    import pdb
                    pdb.set_trace()
                    raise Exception("Object contains connection that fails to refer to object")
        return ret

    def _connect(self, other, connection, in_3prime_direction=None):
        ## TODO fix circular references        
        A,B = [connection.A, connection.B]
        if in_3prime_direction is not None:
            A.is_3prime_side_of_connection = not in_3prime_direction
            B.is_3prime_side_of_connection = in_3prime_direction
            
        A.connection = B.connection = connection
        self.connections.append(connection)
        other.connections.append(connection)
        l = A.container.locations
        if A not in l: l.append(A)
        l = B.container.locations
        if B not in l: l.append(B)
        

    # def _find_connections(self, loc):
    #     return [c for c in self.connections if c.A == loc or c.B == loc]

class SegmentParticle(PointParticle):
    def __init__(self, type_, position, name="A", **kwargs):
        self.name = name
        self.contour_position = None
        PointParticle.__init__(self, type_, position, name=name, **kwargs)
        self.intrahelical_neighbors = []
        self.other_neighbors = []
        self.locations = []

    def get_intrahelical_above(self):
        """ Returns bead directly above self """
        # assert( len(self.intrahelical_neighbors) <= 2 )
        for b in self.intrahelical_neighbors:
            if b.get_contour_position(self.parent) > self.contour_position:
                return b

    def get_intrahelical_below(self):
        """ Returns bead directly below self """
        # assert( len(self.intrahelical_neighbors) <= 2 )
        for b in self.intrahelical_neighbors:
            if b.get_contour_position(self.parent) < self.contour_position:
                return b

    def _neighbor_should_be_added(self,b):
        if type(self.parent) != type(b.parent):
            return True

        c1 = self.contour_position
        c2 = b.get_contour_position(self.parent)
        if c2 < c1:
            b0 = self.get_intrahelical_below()
        else:
            b0 = self.get_intrahelical_above()

        if b0 is not None:            
            c0 = b0.get_contour_position(self.parent)
            if np.abs(c2-c1) < np.abs(c0-c1):
                ## remove b0
                self.intrahelical_neighbors.remove(b0)
                b0.intrahelical_neighbors.remove(self)
                return True
            else:
                return False
        return True
        
    def make_intrahelical_neighbor(self,b):
        add1 = self._neighbor_should_be_added(b)
        add2 = b._neighbor_should_be_added(self)
        if add1 and add2:
            # assert(len(b.intrahelical_neighbors) <= 1)
            # assert(len(self.intrahelical_neighbors) <= 1)
            self.intrahelical_neighbors.append(b)
            b.intrahelical_neighbors.append(self)

    def get_nt_position(self, seg):
        """ Returns the "address" of the nucleotide relative to seg in
        nucleotides, taking the shortest (intrahelical) contour length route to seg
        """
        if seg == self.parent:
            return seg.contour_to_nt_pos(self.contour_position)
        else:
            pos = self.get_contour_position(seg)
            return seg.contour_to_nt_pos(pos)

    def get_contour_position(self,seg):
        if seg == self.parent:
            return self.contour_position
        else:
            cutoff = 30*3
            target_seg = seg

            ## depth-first search
            ## TODO cache distances to nearby locations?
            def descend_search_tree(seg, contour_in_seg, distance=0, visited_segs=None):
                nonlocal cutoff
                if visited_segs is None: visited_segs = []

                if seg == target_seg:
                    # pdb.set_trace()
                    ## Found a segment in our target
                    sign = (contour_in_seg == 1) - (contour_in_seg == 0)
                    assert( sign in (1,-1) )
                    if distance < cutoff: # TODO: check if this does anything
                        cutoff = distance
                    return [[distance, contour_in_seg+sign*seg.nt_pos_to_contour(distance)]], [(seg, contour_in_seg, distance)]
                if distance > cutoff:
                    return None,None
                    
                ret_list = []
                hist_list = []
                ## Find intrahelical locations in seg that we might pass through
                conn_locs = seg.get_connections_and_locations("intrahelical")
                if isinstance(target_seg, SingleStrandedSegment):
                    tmp = seg.get_connections_and_locations("sscrossover")
                    conn_locs = conn_locs + list(filter(lambda x: x[2].container == target_seg, tmp))
                for c,A,B in conn_locs:
                    if B.container in visited_segs: continue
                    dx = seg.contour_to_nt_pos( A.address, round_nt=False ) - seg.contour_to_nt_pos( contour_in_seg, round_nt=False)
                    dx = np.abs(dx)
                    results,history = descend_search_tree( B.container, B.address,
                                                   distance+dx, visited_segs + [seg] )
                    if results is not None:
                        ret_list.extend( results )
                        hist_list.extend( history )
                return ret_list,hist_list

            results,history = descend_search_tree(self.parent, self.contour_position)
            if results is None or len(results) == 0:
                raise Exception("Could not find location in segment") # TODO better error
            return sorted(results,key=lambda x:x[0])[0][1]

            # nt_pos = self.get_nt_position(seg)
            # return seg.nt_pos_to_contour(nt_pos)

    def update_position(self, contour_position):
        self.contour_position = contour_position
        self.position = self.parent.contour_to_position(contour_position)
        if 'orientation_bead' in self.__dict__:
            o = self.orientation_bead
            o.contour_position = contour_position
            orientation = self.parent.contour_to_orientation(contour_position)
            if orientation is None:
                print("WARNING: local_twist is True, but orientation is None; using identity")
                orientation = np.eye(3)
            o.position = self.position + orientation.dot( np.array((Segment.orientation_bond.r0,0,0)) )
            
    def __repr__(self):
        return "<SegmentParticle {} on {}[{:.2f}]>".format( self.name, self.parent, self.contour_position)


## TODO break this class into smaller, better encapsulated pieces
class Segment(ConnectableElement, Group):

    """ Base class that describes a segment of DNA. When built from
    cadnano models, should not span helices """

    """Define basic particle types"""
    dsDNA_particle = ParticleType("D",
                                  diffusivity = 43.5,
                                  mass = 300,
                                  radius = 3,                 
                              )
    orientation_particle = ParticleType("O",
                                        diffusivity = 100,
                                        mass = 300,
                                        radius = 1,
                                    )

    # orientation_bond = HarmonicBond(10,2)
    orientation_bond = HarmonicBond(30,1.5, rRange = (0,500) )

    ssDNA_particle = ParticleType("S",
                                  diffusivity = 43.5,
                                  mass = 150,
                                  radius = 3,                 
                              )

    def __init__(self, name, num_nt, 
                 start_position = None,
                 end_position = None, 
                 segment_model = None,
                 **kwargs):

        if start_position is None: start_position = np.array((0,0,0))

        Group.__init__(self, name, children=[], **kwargs)
        ConnectableElement.__init__(self, connection_locations=[], connections=[])

        if 'segname' not in kwargs:
            self.segname = name
        # self.resname = name
        self.start_orientation = None
        self.twist_per_nt = 0

        self.beads = [c for c in self.children] # self.beads will not contain orientation beads

        self._bead_model_generation = 0    # TODO: remove?
        self.segment_model = segment_model # TODO: remove?

        self.strand_pieces = dict()
        for d in ('fwd','rev'):
            self.strand_pieces[d] = []

        self.num_nt = int(num_nt)
        if end_position is None:
            end_position = np.array((0,0,self.distance_per_nt*num_nt)) + start_position
        self.start_position = start_position
        self.end_position = end_position

        ## Used to assign cadnano names to beads
        self._generate_bead_callbacks = []
        self._generate_nucleotide_callbacks = []

        ## Set up interpolation for positions
        self._set_splines_from_ends()

        self.sequence = None

    def __repr__(self):
        return "<{} {}[{:d}]>".format( type(self), self.name, self.num_nt )

    def set_splines(self, contours, coords):
        tck, u = interpolate.splprep( coords.T, u=contours, s=0, k=1)
        self.position_spline_params = (tck,u)

    def set_orientation_splines(self, contours, quaternions):
        tck, u = interpolate.splprep( quaternions.T, u=contours, s=0, k=1)
        self.quaternion_spline_params = (tck,u)

    def get_center(self):
        tck, u = self.position_spline_params
        return np.mean(self.contour_to_position(u), axis=0)

    def __filter_contours(contours, positions, position_filter, contour_filter):
        u = contours
        r = positions

        ## Filter
        ids = list(range(len(u)))
        if contour_filter is not None:
            ids = list(filter(lambda i: contour_filter(u[i]), ids))
        if position_filter is not None:
            ids = list(filter(lambda i: position_filter(r[i,:]), ids))
        return ids

    def translate(self, translation_vector, position_filter=None, contour_filter=None):
        dr = np.array(translation_vector)
        tck, u = self.position_spline_params
        r = self.contour_to_position(u)

        ids = Segment.__filter_contours(u, r, position_filter, contour_filter)

        ## Translate
        r[ids,:] = r[ids,:] + dr[np.newaxis,:]
        self.set_splines(u,r)

    def rotate(self, rotation_matrix, about=None, position_filter=None, contour_filter=None):
        tck, u = self.position_spline_params
        r = self.contour_to_position(u)

        ids = Segment.__filter_contours(u, r, position_filter, contour_filter)

        if about is None:
            ## TODO: do this more efficiently
            r[ids,:] = np.array([rotation_matrix.dot(r[i,:]) for i in range(r.shape[0])])
        else:
            dr = np.array(about)
            ## TODO: do this more efficiently
            r[ids,:] = np.array([rotation_matrix.dot(r[i,:]-dr) + dr for  i in range(r.shape[0])])

        self.set_splines(u,r)

        if self.quaternion_spline_params is not None:
            ## TODO: performance: don't shift between quaternion and matrix representations so much
            tck, u = self.quaternion_spline_params
            orientations = [self.contour_to_orientation(v) for v in u]
            for i in ids:
                orientations[i,:] = rotation_matrix.dot(orientations[i])
            quats = [quaternion_from_matrix(o) for o in orientations]
            self.set_orientation_splines(u, quats)

    def _set_splines_from_ends(self):
        self.quaternion_spline_params = None
        coords = np.array([self.start_position, self.end_position])
        self.set_splines([0,1], coords)

    def clear_all(self):
        Group.clear_all(self)  # TODO: use super?
        self.beads = []
        for c,loc,other in self.get_connections_and_locations():
            loc.particle = None

    def contour_to_nt_pos(self, contour_pos, round_nt=False):
        nt = contour_pos*(self.num_nt) - 0.5
        if round_nt:
            assert( np.isclose(np.around(nt),nt) )
            nt = np.around(nt)
        return nt

    def nt_pos_to_contour(self,nt_pos):
        return (nt_pos+0.5)/(self.num_nt)

    def contour_to_position(self,s):
        p = interpolate.splev( s, self.position_spline_params[0] )
        if len(p) > 1: p = np.array(p).T
        return p

    def contour_to_tangent(self,s):
        t = interpolate.splev( s, self.position_spline_params[0], der=1 )
        t = (t / np.linalg.norm(t,axis=0))
        return t.T
        

    def contour_to_orientation(self,s):
        assert( isinstance(s,float) or isinstance(s,int) or len(s) == 1 )   # TODO make vectorized version

        if self.quaternion_spline_params is None:
            axis = self.contour_to_tangent(s)
            axis = axis / np.linalg.norm(axis)
            rotAxis = np.cross(axis,np.array((0,0,1)))
            rotAxisL = np.linalg.norm(rotAxis)
            zAxis = np.array((0,0,1))

            if rotAxisL > 0.001:
                theta = np.arcsin(rotAxisL) * 180/np.pi
                if axis.dot(zAxis) < 0: theta = 180-theta
                orientation0 = rotationAboutAxis( rotAxis/rotAxisL, theta, normalizeAxis=False ).T
            else:
                orientation0 = np.eye(3) if axis.dot(zAxis) > 0 else \
                               rotationAboutAxis( np.array((1,0,0)), 180, normalizeAxis=False )
            orientation = rotationAboutAxis( axis, self.twist_per_nt*self.contour_to_nt_pos(s), normalizeAxis=False )
            orientation = orientation.dot(orientation0)
        else:
            q = interpolate.splev( s, self.quaternion_spline_params[0] )
            if len(q) > 1: q = np.array(q).T # TODO: is this needed?
            orientation = quaternion_to_matrix(q)

        return orientation

    def get_contour_sorted_connections_and_locations(self,type_):
        sort_fn = lambda c: c[1].address
        cl = self.get_connections_and_locations(type_)
        return sorted(cl, key=sort_fn)
    
    def randomize_unset_sequence(self):
        bases = list(seqComplement.keys())
        # bases = ['T']        ## FOR DEBUG
        if self.sequence is None:
            self.sequence = [random.choice(bases) for i in range(self.num_nt)]
        else:
            assert(len(self.sequence) == self.num_nt) # TODO move
            for i in range(len(self.sequence)):
                if self.sequence[i] is None:
                    self.sequence[i] = random.choice(bases)

    def _get_num_beads(self, max_basepairs_per_bead, max_nucleotides_per_bead ):
        raise NotImplementedError

    def _generate_one_bead(self, contour_position, nts):
        raise NotImplementedError

    def _generate_atomic_nucleotide(self, contour_position, is_fwd, seq, scale, strand_segment):
        """ Seq should include modifications like 5T, T3 Tsinglet; direction matters too """

        # print("Generating nucleotide at {}".format(contour_position))
        
        pos = self.contour_to_position(contour_position)
        orientation = self.contour_to_orientation(contour_position)

        """ deleteme
        ## TODO: move this code (?)
        if orientation is None:
            import pdb
            pdb.set_trace()
            axis = self.contour_to_tangent(contour_position)
            angleVec = np.array([1,0,0])
            if axis.dot(angleVec) > 0.9: angleVec = np.array([0,1,0])
            angleVec = angleVec - angleVec.dot(axis)*axis
            angleVec = angleVec/np.linalg.norm(angleVec)
            y = np.cross(axis,angleVec)
            orientation = np.array([angleVec,y,axis]).T
            ## TODO: improve placement of ssDNA
            # rot = rotationAboutAxis( axis, contour_position*self.twist_per_nt*self.num_nt, normalizeAxis=True )
            # orientation = rot.dot(orientation)
        else:
            orientation = orientation                            
        """
        key = seq
        nt_dict = canonicalNtFwd if is_fwd else canonicalNtRev

        atoms = nt_dict[ key ].generate() # TODO: clone?
        atoms.orientation = orientation.dot(atoms.orientation)
        if isinstance(self, SingleStrandedSegment):
            if scale is not None and scale != 1:
                for a in atoms:
                    a.position = scale*a.position
                    a.beta = 0
            atoms.position = pos - atoms.atoms_by_name["C1'"].collapsedPosition()
        else:
            if scale is not None and scale != 1:
                if atoms.sequence in ("A","G"):
                    r0 = atoms.atoms_by_name["N9"].position
                else:
                    r0 = atoms.atoms_by_name["N1"].position
                for a in atoms:
                    if a.name[-1] in ("'","P","T"):
                        a.position = scale*(a.position-r0) + r0
                        a.beta = 0
            atoms.position = pos
        
        atoms.contour_position = contour_position
        strand_segment.add(atoms)

        for callback in self._generate_nucleotide_callbacks:
            callback(atoms)

        return atoms

    def add_location(self, nt, type_, on_fwd_strand=True):
        ## Create location if needed, add to segment
        c = self.nt_pos_to_contour(nt)
        assert(c >= 0 and c <= 1)
        # TODO? loc = self.Location( address=c, type_=type_, on_fwd_strand=is_fwd )
        loc = Location( self, address=c, type_=type_, on_fwd_strand=on_fwd_strand )
        self.locations.append(loc)

    ## TODO? Replace with abstract strand-based model?

    def add_nick(self, nt, on_fwd_strand=True):
        self.add_3prime(nt,on_fwd_strand)
        self.add_5prime(nt+1,on_fwd_strand)

    def add_5prime(self, nt, on_fwd_strand=True):
        if isinstance(self,SingleStrandedSegment):
            on_fwd_strand = True
        self.add_location(nt,"5prime",on_fwd_strand)

    def add_3prime(self, nt, on_fwd_strand=True):
        if isinstance(self,SingleStrandedSegment):
            on_fwd_strand = True
        self.add_location(nt,"3prime",on_fwd_strand)

    def get_3prime_locations(self):
        return sorted(self.get_locations("3prime"),key=lambda x: x.address)
    
    def get_5prime_locations(self):
        ## TODO? ensure that data is consistent before _build_model calls
        return sorted(self.get_locations("5prime"),key=lambda x: x.address)

    def iterate_connections_and_locations(self, reverse=False):
        ## connections to other segments
        cl = self.get_contour_sorted_connections_and_locations()
        if reverse:
            cl = cl[::-1]
            
        for c in cl:
            yield c

    ## TODO rename
    def _add_strand_piece(self, strand_piece):
        """ Registers a strand segment within this object """

        ## TODO use weakref
        d = 'fwd' if strand_piece.is_fwd else 'rev'

        ## Validate strand_piece (ensure no clashes)
        for s in self.strand_pieces[d]:
            l,h = sorted((s.start,s.end))
            for value in (strand_piece.start,strand_piece.end):
                assert( value < l or value > h )

        ## Add strand_piece in correct order
        self.strand_pieces[d].append(strand_piece)
        self.strand_pieces[d] = sorted(self.strand_pieces[d],
                                       key = lambda x: x.start)

    ## TODO rename
    def get_strand_segment(self, nt_pos, is_fwd, move_at_least=0.5):
        """ Walks through locations, checking for crossovers """
        # if self.name in ("6-1","1-1"):
        #     import pdb
        #     pdb.set_trace()
        move_at_least = 0

        ## Iterate through locations
        # locations = sorted(self.locations, key=lambda l:(l.address,not l.on_fwd_strand), reverse=(not is_fwd))
        def loc_rank(l):
            nt = l.get_nt_pos()
            ## optionally add logic about type of connection
            return (nt, not l.on_fwd_strand)
        # locations = sorted(self.locations, key=lambda l:(l.address,not l.on_fwd_strand), reverse=(not is_fwd))
        locations = sorted(self.locations, key=loc_rank, reverse=(not is_fwd))
        # print(locations)

        for l in locations:
            # TODOTODO probably okay
            if l.address == 0:
                pos = 0.0
            elif l.address == 1:
                pos = self.num_nt-1
            else:
                pos = self.contour_to_nt_pos(l.address, round_nt=True)

            ## DEBUG


            ## Skip locations encountered before our strand
            # tol = 0.1
            # if is_fwd:
            #     if pos-nt_pos <= tol: continue 
            # elif   nt_pos-pos <= tol: continue
            if (pos-nt_pos)*(2*is_fwd-1) < move_at_least: continue
            ## TODO: remove move_at_least
            if np.isclose(pos,nt_pos):
                if l.is_3prime_side_of_connection: continue

            ## Stop if we found the 3prime end
            if l.on_fwd_strand == is_fwd and l.type_ == "3prime":
                # print("  found end at",l)
                return pos, None, None, None, None

            ## Check location connections
            c = l.connection
            if c is None: continue
            B = c.other(l)            

            ## Found a location on the same strand?
            if l.on_fwd_strand == is_fwd:
                # print("  passing through",l)
                # print("from {}, connection {} to {}".format(nt_pos,l,B))
                Bpos = B.get_nt_pos()
                return pos, B.container, Bpos, B.on_fwd_strand, 0.5
                
            ## Stop at other strand crossovers so basepairs line up
            elif c.type_ == "crossover":
                if nt_pos == pos: continue
                # print("  pausing at",l)
                return pos, l.container, pos+(2*is_fwd-1), is_fwd, 0

        raise Exception("Shouldn't be here")
        # print("Shouldn't be here")
        ## Made it to the end of the segment without finding a connection
        return 1*is_fwd, None, None, None

    def get_nearest_bead(self, contour_position):
        if len(self.beads) < 1: return None
        cs = np.array([b.contour_position for b in self.beads]) # TODO: cache
        # TODO: include beads in connections?
        i = np.argmin((cs - contour_position)**2)

        return self.beads[i]

    def _get_atomic_nucleotide(self, nucleotide_idx, is_fwd=True):
        d = 'fwd' if is_fwd else 'rev'
        for s in self.strand_pieces[d]:
            try:
                return s.get_nucleotide(nucleotide_idx)
            except:
                pass
        raise Exception("Could not find nucleotide in {} at {}.{}".format( self, nucleotide_idx, d ))

    def get_all_consecutive_beads(self, number):
        assert(number >= 1)
        ## Assume that consecutive beads in self.beads are bonded
        ret = []
        for i in range(len(self.beads)-number+1):
            tmp = [self.beads[i+j] for j in range(0,number)]
            ret.append( tmp )
        return ret   

    def _add_bead(self,b,set_contour=False):
        if set_contour:
            b.contour_position = b.get_contour_position(self)
        
        # assert(b.parent is None)
        if b.parent is not None:
            b.parent.children.remove(b)
        self.add(b)
        self.beads.append(b) # don't add orientation bead
        if "orientation_bead" in b.__dict__: # TODO: think of a cleaner approach
            o = b.orientation_bead
            o.contour_position = b.contour_position
            if o.parent is not None:
                o.parent.children.remove(o)
            self.add(o)
            self.add_bond(b,o, Segment.orientation_bond, exclude=True)

    def _rebuild_children(self, new_children):
        # print("_rebuild_children on %s" % self.name)
        old_children = self.children
        old_beads = self.beads
        self.children = []
        self.beads = []

        if True:
            ## TODO: remove this if duplicates are never found 
            # print("Searching for duplicate particles...")
            ## Remove duplicates, preserving order
            tmp = []
            for c in new_children:
                if c not in tmp:
                    tmp.append(c)
                else:
                    print("  DUPLICATE PARTICLE FOUND!")
            new_children = tmp

        for b in new_children:
            self.beads.append(b)
            self.children.append(b)
            if "orientation_bead" in b.__dict__: # TODO: think of a cleaner approach
                self.children.append(b.orientation_bead)
            
        # tmp = [c for c in self.children if c not in old_children]
        # assert(len(tmp) == 0)
        # tmp = [c for c in old_children if c not in self.children]
        # assert(len(tmp) == 0)
        assert(len(old_children) == len(self.children))
        assert(len(old_beads) == len(self.beads))


    def _generate_beads(self, bead_model, max_basepairs_per_bead, max_nucleotides_per_bead):

        """ Generate beads (positions, types, etc) and bonds, angles, dihedrals, exclusions """
        ## TODO: decide whether to remove bead_model argument
        ##       (currently unused)

        ## First find points between-which beads must be generated
        # conn_locs = self.get_contour_sorted_connections_and_locations()
        # locs = [A for c,A,B in conn_locs]
        # existing_beads = [l.particle for l in locs if l.particle is not None]
        # if self.name == "S001":
        #     pdb.set_trace()

        existing_beads0 = {l.particle for l in self.locations if l.particle is not None}
        existing_beads = sorted( list(existing_beads0), key=lambda b: b.get_contour_position(self) )

        # if self.num_nt == 1 and all([l.particle is not None for l in self.locations]):
        #     pdb.set_trace()
        #     return

        for b in existing_beads:
            assert(b.parent is not None)

        ## Add ends if they don't exist yet
        ## TODOTODO: test 1 nt segments?
        if len(existing_beads) == 0 or existing_beads[0].get_nt_position(self) >= 0.5:
            # if len(existing_beads) > 0:            
            #     assert(existing_beads[0].get_nt_position(self) >= 0.5)
            b = self._generate_one_bead( self.nt_pos_to_contour(0), 0)
            existing_beads = [b] + existing_beads

        if existing_beads[-1].get_nt_position(self)-(self.num_nt-1) < -0.5 or len(existing_beads)==1:
            b = self._generate_one_bead( self.nt_pos_to_contour(self.num_nt-1), 0)
            existing_beads.append(b)
        assert(len(existing_beads) > 1)

        ## Walk through existing_beads, add beads between
        tmp_children = []       # build list of children in nice order
        last = None

        for I in range(len(existing_beads)-1):
            eb1,eb2 = [existing_beads[i] for i in (I,I+1)]
            assert( eb1 is not eb2 )

            # if np.isclose(eb1.position[2], eb2.position[2]):
            #     import pdb
            #     pdb.set_trace()

            # print(" %s working on %d to %d" % (self.name, eb1.position[2], eb2.position[2]))
            e_ds = eb2.get_contour_position(self) - eb1.get_contour_position(self)
            num_beads = self._get_num_beads( e_ds, max_basepairs_per_bead, max_nucleotides_per_bead )

            ## Ensure there is a ssDNA bead between dsDNA beads
            if num_beads == 0 and isinstance(self,SingleStrandedSegment) and isinstance(eb1.parent,DoubleStrandedSegment) and isinstance(eb2.parent,DoubleStrandedSegment):
                num_beads = 1
            ## TODO similarly ensure there is a dsDNA bead between ssDNA beads

            ds = e_ds / (num_beads+1)
            nts = ds*self.num_nt
            eb1.num_nt += 0.5*nts
            eb2.num_nt += 0.5*nts

            ## Add beads
            if eb1.parent == self:
                tmp_children.append(eb1)

            s0 = eb1.get_contour_position(self)
            if last is not None:
                last.make_intrahelical_neighbor(eb1)
            last = eb1
            for j in range(num_beads):
                s = ds*(j+1) + s0
                # if self.name in ("51-2","51-3"):
                # if self.name in ("31-2",):
                #     print(" adding bead at {}".format(s))
                b = self._generate_one_bead(s,nts)

                last.make_intrahelical_neighbor(b)
                last = b
                tmp_children.append(b)

        last.make_intrahelical_neighbor(eb2)

        if eb2.parent == self:
            tmp_children.append(eb2)
        # if self.name in ("31-2",):
        #     pdb.set_trace()
        self._rebuild_children(tmp_children)

        for callback in self._generate_bead_callbacks:
            callback(self)

    def _regenerate_beads(self, max_nts_per_bead=4, ):
        ...
    

class DoubleStrandedSegment(Segment):

    """ Class that describes a segment of ssDNA. When built from
    cadnano models, should not span helices """

    def __init__(self, name, num_bp, start_position = np.array((0,0,0)),
                 end_position = None, 
                 segment_model = None,
                 local_twist = False,
                 num_turns = None,
                 start_orientation = None,
                 twist_persistence_length = 90,
                 **kwargs):
        
        self.helical_rise = 10.44
        self.distance_per_nt = 3.4
        Segment.__init__(self, name, num_bp,
                         start_position,
                         end_position, 
                         segment_model,
                         **kwargs)
        self.num_bp = self.num_nt

        self.local_twist = local_twist
        if num_turns is None:
            num_turns = float(num_bp) / self.helical_rise
        self.twist_per_nt = float(360 * num_turns) / num_bp

        if start_orientation is None:
            start_orientation = np.eye(3) # np.array(((1,0,0),(0,1,0),(0,0,1)))
        self.start_orientation = start_orientation
        self.twist_persistence_length = twist_persistence_length

        self.nicks = []

        self.start = self.start5 = Location( self, address=0, type_= "end5" )
        self.start3 = Location( self, address=0, type_ = "end3", on_fwd_strand=False )

        self.end = self.end3 = Location( self, address=1, type_ = "end3" )
        self.end5 = Location( self, address=1, type_= "end5", on_fwd_strand=False )
        # for l in (self.start5,self.start3,self.end3,self.end5):
        #     self.locations.append(l)

        ## Set up interpolation for azimuthal angles 
        a = np.array([self.start_position,self.end_position]).T
        tck, u = interpolate.splprep( a, u=[0,1], s=0, k=1)
        self.position_spline_params = (tck,u)
        
        ## TODO: initialize sensible spline for orientation

    ## Convenience methods
    ## TODO: add errors if unrealistic connections are made
    ## TODO: make connections automatically between unconnected strands
    def connect_start5(self, end3, type_="intrahelical", force_connection=False):
        if isinstance(end3, SingleStrandedSegment):
            end3 = end3.end3
        self._connect_ends( self.start5, end3, type_, force_connection = force_connection )
    def connect_start3(self, end5, type_="intrahelical", force_connection=False):
        if isinstance(end5, SingleStrandedSegment):
            end5 = end5.start5
        self._connect_ends( self.start3, end5, type_, force_connection = force_connection )
    def connect_end3(self, end5, type_="intrahelical", force_connection=False):
        if isinstance(end5, SingleStrandedSegment):
            end5 = end5.start5
        self._connect_ends( self.end3, end5, type_, force_connection = force_connection )
    def connect_end5(self, end3, type_="intrahelical", force_connection=False):
        if isinstance(end3, SingleStrandedSegment):
            end3 = end3.end3