from easy import Mapping

import numpy as np
import math

# TODO move easy.Mapping here?

class DistanceMapping(Mapping):
    '''a hooke-like distance mapping: 
    
    f(p1, p2) = norm(p2 - p1) - rest_length
    '''

    # TODO add pairs
    
    def __init__(self, node, **kwargs):
        Mapping.__init__(self, node, **kwargs)

        self.rest_length = 0

    def apply(self, out, at):
        
        p1 = at[0][0]
        p2 = at[1][0]
        
        delta = p2 - p1
        norm = math.sqrt(delta.dot(delta))
        
        out[0][0] = norm - self.rest_length

        
    def jacobian(self, js, at):

        p1 = at[0][0]
        p2 = at[1][0]
        
        delta = p2 - p1
        norm = math.sqrt(delta.dot(delta))

        if norm > 1e-14:
            u = delta / norm
            js[0][0] = -u
            js[1][0] = u
        else:
            # warning
            js[0][0] = np.zeros(3)
            js[1][0] = np.zeros(3)


    def geometric_stiffness(self, gs, at, force):
        
        p1 = at[0][0]
        p2 = at[1][0]
        
        delta = p2 - p1

        norm = math.sqrt(delta.dot(delta))
        if norm > 1e-14:
            
            u = delta / norm
            f = force[0][0]

            block = (np.identity(3) - np.outer(u, u)) * (f / norm)
            gs[:3, :3] = block
            gs[3:, 3:] = block

            gs[:3, 3:] = -block
            gs[3:, :3] = -block

            
        else:
            gs[:, :] = 0

