#***************************************************************************
#*                                                                         *
#*   Copyright (c) 2018 kbwbe (klaus)                                      * 
#*                                                                         *
#*   This program is free software; you can redistribute it and/or modify  *
#*   it under the terms of the GNU Lesser General Public License (LGPL)    *
#*   as published by the Free Software Foundation; either version 2 of     *
#*   the License, or (at your option) any later version.                   *
#*   for detail see the LICENCE text file.                                 *
#*                                                                         *
#*   This program is distributed in the hope that it will be useful,       *
#*   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
#*   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
#*   GNU Library General Public License for more details.                  *
#*                                                                         *
#*   You should have received a copy of the GNU Library General Public     *
#*   License along with this program; if not, write to the Free Software   *
#*   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  *
#*   USA                                                                   *
#*                                                                         *
#***************************************************************************

import random
import time
import traceback
import math
import copy
import FreeCAD, FreeCADGui, Part
from  FreeCAD import Base


#------------------------------------------------------------------------------
# Some small helper tools analog to hamish's assembly2
#------------------------------------------------------------------------------
def appVersionStr():
    version = int(FreeCAD.Version()[0])
    subVersion = int(FreeCAD.Version()[1])
    return "%03d.%03d" %(version,subVersion)
#------------------------------------------------------------------------------
def getObjectEdgeFromName( obj, name ):
    assert name.startswith('Edge')
    ind = int( name[4:]) -1 
    return obj.Shape.Edges[ind]
#------------------------------------------------------------------------------
def getObjectFaceFromName( obj, faceName ):
    assert faceName.startswith('Face')
    ind = int( faceName[4:]) -1 
    return obj.Shape.Faces[ind]
#------------------------------------------------------------------------------
def getObjectVertexFromName( obj, name ):
    assert name.startswith('Vertex')
    ind = int( name[6:]) -1 
    return obj.Shape.Vertexes[ind]
#------------------------------------------------------------------------------
def isLine(param):
    if hasattr(Part,"LineSegment"):
        return isinstance(param,(Part.Line,Part.LineSegment))
    else:
        return isinstance(param,Part.Line)
#------------------------------------------------------------------------------
def getPos(obj, subElementName):
    pos = None
    if subElementName.startswith('Face'):
        face = getObjectFaceFromName(obj, subElementName)
        surface = face.Surface
        if str(surface) == '<Plane object>':
            pos = getObjectFaceFromName(obj, subElementName).Faces[0].BoundBox.Center
            # axial constraint for Planes
            # pos = surface.Position
        elif all( hasattr(surface,a) for a in ['Axis','Center','Radius'] ):
            pos = surface.Center
        elif str(surface).startswith('<SurfaceOfRevolution'):
            pos = getObjectFaceFromName(obj, subElementName).Edges[0].Curve.Center
    elif subElementName.startswith('Edge'):
        edge = getObjectEdgeFromName(obj, subElementName)
        if isLine(edge.Curve):
            if appVersionStr() <= "000.016":
                pos = edge.Curve.StartPoint
            else:
                pos = edge.firstVertex(True).Point
        elif hasattr( edge.Curve, 'Center'): #circular curve
            pos = edge.Curve.Center
    elif subElementName.startswith('Vertex'):
        return  getObjectVertexFromName(obj, subElementName).Point
    return pos # maybe none !!
#------------------------------------------------------------------------------
def getAxis(obj, subElementName):
    axis = None
    if subElementName.startswith('Face'):
        face = getObjectFaceFromName(obj, subElementName)
        surface = face.Surface
        if hasattr(surface,'Axis'):
            axis = surface.Axis
        elif str(surface).startswith('<SurfaceOfRevolution'):
            axis = face.Edges[0].Curve.Axis
    elif subElementName.startswith('Edge'):
        edge = getObjectEdgeFromName(obj, subElementName)
        if isLine(edge.Curve):
            axis = edge.Curve.tangent(0)[0]
        elif hasattr( edge.Curve, 'Axis'): #circular curve
            axis =  edge.Curve.Axis
    return axis # may be none!
#------------------------------------------------------------------------------
# End of helper-tools
#------------------------------------------------------------------------------





#------------------------------------------------------------------------------
class SolverSystem():
    def __init__(self):
        self.doc = None
        self.rigids = []        # list of rigid bodies
        self.constraints = []
        self.objectNames = []
        
    def clear(self):
        for r in self.rigids:
            r.clear()
        self.rigids = []
        self.constraints = []
        self.objectNames = []
        
    def getRigid(self,objectName):
        '''get a Rigid with objectName'''
        rigs = [r for r in self.rigids if r.objectName == objectName]
        if len(rigs) > 0: return rigs[0]
        return None
        
    def loadSystem(self,doc):
        self.clear()
        self.doc = doc
        self.constraints = [ obj for obj in doc.Objects if 'ConstraintInfo' in obj.Content]
        #
        # Extract all the objectnames which are affected by constraints..
        self.objectNames = []
        for c in self.constraints:
            for attr in ['Object1','Object2']:
                objectName = getattr(c, attr, None)
                if objectName <> None and not objectName in self.objectNames:
                    self.objectNames.append( objectName )
        #
        # create a Rigid() dataStructure for each of these objectnames...
        for o in self.objectNames:
            ob1 = doc.getObject(o)
            if hasattr(ob1, "fixedPosition"):
                fx = ob1.fixedPosition
            else:
                fx = False
            rig = Rigid(
                o,
                fx,
                ob1.Placement
                )
            rig.spinCenter = ob1.Shape.BoundBox.Center
            self.rigids.append(rig)
        #
        #link contraints to rigids
        for c in self.constraints:
            rig1 = self.getRigid(c.Object1)
            dep1 = Dependency()
            dep1.Type = c.Type
            try:
                dep1.direction = c.directionConstraint
            except:
                pass
            try:
                dep1.offset = c.offset
            except:
                pass
            
            rig2 = self.getRigid(c.Object2)
            dep2 = Dependency()
            dep2.Type = c.Type
            try:
                dep2.direction = c.directionConstraint
            except:
                pass
            try:
                dep2.offset = c.offset
            except:
                pass
                
            if c.Type == "circularEdge":
                dep1.refType = "pointAxis"
                dep2.refType = "pointAxis"
                ob1 = doc.getObject(c.Object1)
                ob2 = doc.getObject(c.Object2)
                circleEdge1 = getObjectEdgeFromName(ob1, c.SubElement1)
                circleEdge2 = getObjectEdgeFromName(ob2, c.SubElement2)
                dep1.refPoint = circleEdge1.Curve.Center
                dep2.refPoint = circleEdge2.Curve.Center
                axis1 = circleEdge1.Curve.Axis
                axis2 = circleEdge2.Curve.Axis
                if dep2.direction == "opposed":
                    axis2.multiply(-1.0)
                dep1.refAxisEnd = dep1.refPoint.add(axis1)
                dep2.refAxisEnd = dep2.refPoint.add(axis2)
                #
                if abs(dep2.offset) > 1e-6:
                    offsetAdjustVec = Base.Vector(axis2.x,axis2.y,axis2.z)
                    offsetAdjustVec.multiply(dep2.offset)
                    dep2.refPoint = dep2.refPoint.add(offsetAdjustVec)
                    dep2.refAxisEnd = dep2.refAxisEnd.add(offsetAdjustVec)
                #
                dep1.foreignDependency = dep2
                dep2.foreignDependency = dep1
                rig1.dependencies.append(dep1)
                rig2.dependencies.append(dep2)
                
            if c.Type == "plane":
                dep1.refType = "pointNormal"
                dep2.refType = "pointNormal"
                ob1 = doc.getObject(c.Object1)
                ob2 = doc.getObject(c.Object2)
                plane1 = getObjectFaceFromName(ob1, c.SubElement1)
                plane2 = getObjectFaceFromName(ob2, c.SubElement2)
                dep1.refPoint = plane1.Faces[0].BoundBox.Center
                dep2.refPoint = plane2.Faces[0].BoundBox.Center
                normal1 = plane1.Surface.Axis
                normal2 = plane2.Surface.Axis
                if dep2.direction == "opposed":
                    normal2.multiply(-1.0)
                dep1.refAxisEnd = dep1.refPoint.add(normal1)
                dep2.refAxisEnd = dep2.refPoint.add(normal2)
                #
                if abs(dep2.offset) > 1e-6:
                    offsetAdjustVec = Base.Vector(normal2.x,normal2.y,normal2.z)
                    offsetAdjustVec.multiply(dep2.offset)
                    dep2.refPoint = dep2.refPoint.add(offsetAdjustVec)
                    dep2.refAxisEnd = dep2.refAxisEnd.add(offsetAdjustVec)
                #
                dep1.foreignDependency = dep2
                dep2.foreignDependency = dep1
                rig1.dependencies.append(dep1)
                rig2.dependencies.append(dep2)
                
            if c.Type == "axial":
                dep1.refType = "pointAxis"
                dep2.refType = "pointAxis"
                ob1 = doc.getObject(c.Object1)
                ob2 = doc.getObject(c.Object2)
                dep1.refPoint = getPos(ob1,c.SubElement1)
                dep2.refPoint = getPos(ob2,c.SubElement2)
                axis1 = getAxis(ob1, c.SubElement1)
                axis2 = getAxis(ob2, c.SubElement2)
                if dep2.direction == "opposed":
                    axis2.multiply(-1.0)
                dep1.refAxisEnd = dep1.refPoint.add(axis1)
                dep2.refAxisEnd = dep2.refPoint.add(axis2)
                dep1.foreignDependency = dep2
                dep2.foreignDependency = dep1
                rig1.dependencies.append(dep1)
                rig2.dependencies.append(dep2)
                
                
                
    def calcMoveData(self,doc):
        for rig in self.rigids:
            if rig.fixed: continue
            depRefPoints = [] 
            depMoveVectors = [] #collect Data to compute central movement of rigid
            #
            for dep in rig.dependencies:
                if dep.Type == "circularEdge":
                    depRefPoints.append(dep.refPoint)
                    dep.moveVector = dep.foreignDependency.refPoint.sub(dep.refPoint)
                    depMoveVectors.append(dep.moveVector)
                if dep.Type == "plane":
                    depRefPoints.append(dep.refPoint)
                    vec1 = dep.foreignDependency.refPoint.sub(dep.refPoint)
                    normal1 = dep.refAxisEnd.sub(dep.refPoint)
                    dot = vec1.dot(normal1)
                    normal1.multiply(dot)
                    dep.moveVector = normal1
                    depMoveVectors.append(dep.moveVector)
                if dep.Type == "axial":
                    depRefPoints.append(dep.refPoint)
                    vec1 = dep.foreignDependency.refPoint.sub(dep.refPoint)
                    destinationAxis = dep.foreignDependency.refAxisEnd.sub(dep.foreignDependency.refPoint)
                    dot = vec1.dot(destinationAxis)
                    parallelToAxisVec = destinationAxis.normalize().multiply(dot)
                    dep.moveVector = vec1.sub(parallelToAxisVec)
                    depMoveVectors.append(dep.moveVector)

            #
            #compute rigid.moveVectorSum
            if ( len(depMoveVectors) > 0 ):
                vec = Base.Vector(0,0,0)
                for mv in depMoveVectors:
                    vec = vec.add(mv)
                vec.multiply(1.0/len(depMoveVectors)) #the average of all movings
                rig.moveVectorSum = vec
            else:
                rig.moveVectorSum = Base.Vector(0,0,0)
            #
            #compute rotation caused by refPoint-attractions and axes mismatch
            if (
                len(depMoveVectors) > 0 and
                rig.spinCenter != None
                ):
                rig.spin = Base.Vector(0,0,0)
                for i in range(0,len(depRefPoints)):
                    vec1 = depRefPoints[i].sub(rig.spinCenter) #level
                    vec2 = depMoveVectors[i].sub(rig.moveVectorSum) # Dist = Force, sub the move!!!
                    axis = vec2.cross(vec1) #torque-vector
                    rig.spin = rig.spin.add(axis)
                    
                #adjust axis' of the dependencies //FIXME (align,opposed,none)
                for dep in rig.dependencies:
                    if (
                        dep.Type == "circularEdge" or
                        dep.Type == "plane" or
                        dep.Type == "axial"
                        ):
                        if dep.direction != "none":
                            rigAxis = dep.refAxisEnd.sub(dep.refPoint)
                            foreignAxis = dep.foreignDependency.refAxisEnd.sub(
                                dep.foreignDependency.refPoint
                                )
                            #
                            #do we have wrong alignement of axes ??
                            dot = rigAxis.dot(foreignAxis)
                            if abs(dot+1.0) < 1e-3: #both axes nearly aligned but false orientation...
                                x = random.uniform(-1e-3,1e-3)
                                y = random.uniform(-1e-3,1e-3)
                                z = random.uniform(-1e-3,1e-3)
                                disturbVector = Base.Vector(x,y,z)
                                foreignAxis = foreignAxis.add(disturbVector)
                                
                            axis = foreignAxis.cross(rigAxis)
                            try:
                                axis.normalize()
                                angle = foreignAxis.getAngle(rigAxis)
                                axis.multiply(angle*300)
                                rig.spin = rig.spin.add(axis)
                            except:
                                pass
                            
                        else: #if dep.direction... (== none)
                            rigAxis = dep.refAxisEnd.sub(dep.refPoint)
                            foreignAxis1 = dep.foreignDependency.refAxisEnd.sub(
                                dep.foreignDependency.refPoint
                                )
                            foreignAxis2 = dep.foreignDependency.refPoint.sub(
                                dep.foreignDependency.refAxisEnd
                                )
                            angle1 = abs(foreignAxis1.getAngle(rigAxis))
                            angle2 = abs(foreignAxis2.getAngle(rigAxis))
                            #
                            if angle1<=angle2:
                                axis = foreignAxis1.cross(rigAxis)
                                foreignAxis = foreignAxis1
                            else:
                                axis = foreignAxis2.cross(rigAxis)
                                foreignAxis = foreignAxis2
                            try:
                                axis.normalize()
                                angle = foreignAxis.getAngle(rigAxis)
                                axis.multiply(angle*300)
                                rig.spin = rig.spin.add(axis)
                            except:
                                pass
                
    def moveRigids(self,doc):
        for rig in self.rigids:
            if rig.fixed: continue
            #
            #Linear moving of a rigid
            if rig.moveVectorSum != None:
                mov = rig.moveVectorSum
                mov.multiply(0.5) # stabilize computation, adjust if needed...
                rot = FreeCAD.Rotation()
                center = rig.spinCenter
                pl = FreeCAD.Placement(mov,rot,center)
                rig.applyPlacementStep(pl)
            #    
            #Rotate the rigid...
            if (
                rig.spin != None and
                rig.spin.Length != 0.0
                ):
                mov = Base.Vector(0,0,0) # weitere Verschiebung ist null
                
                # Spinning of more than 360.0 degrees is useless...
                orig = rig.spin.Length
                if orig>359.0: orig=359.0
                if orig>1e-9:
                    try:
                        sq=abs(orig)/300
                        rig.spin.normalize()
                        rig.spin.multiply(sq)
                        rot = FreeCAD.Rotation(rig.spin.multiply(-1.0),rig.spin.Length)
                        cent = rig.spinCenter
                        pl = FreeCAD.Placement(mov,rot,cent)
                        rig.applyPlacementStep(pl)
                    except:
                        pass
                
    def solutionToParts(self,doc):
        for rig in self.rigids:
            if rig.fixed: continue
            ob1 = doc.getObject(rig.objectName)
            ob1.Placement = rig.placement
        
    def doSolverStep(self,doc):
        self.calcMoveData(doc)
        self.moveRigids(doc)
        
#------------------------------------------------------------------------------
class Rigid():
    ''' All data necessary for one rigid body'''
    def __init__(self,
                 name,
                 fixed,
                 placement
                 ):
        self.objectName = name
        self.fixed = fixed
        self.placement = placement
        self.dependencies = []
        self.spinCenter = None
        self.spin = None
        self.moveVectorSum = None
        
    def applyPlacementStep(self,pl):
        self.placement = pl.multiply(self.placement)
        self.spinCenter = pl.multVec(self.spinCenter)
        for dep in self.dependencies:
            if dep.refPoint != None:
                dep.refPoint = pl.multVec(dep.refPoint)
            if dep.refAxisEnd != None:
                dep.refAxisEnd = pl.multVec(dep.refAxisEnd)
        
    def clear(self):
        for d in self.dependencies:
            d.clear()
        self.dependencies = []
        
#------------------------------------------------------------------------------
class Dependency():
    def __init__(self):
        self.Type = None
        self.refType = None
        self.refPoint = None
        self.refAxisEnd = None
        self.direction = None
        self.offset = None
        self.foreignDependency = None
        self.rotationAxis = None
        self.moveVector = None
        
    def clear(self):
        self.Type = None
        self.refType = None
        self.refPoint = None
        self.refAxisEnd = None
        self.direction = None
        self.offset = None
        self.foreignDependency = None
        self.rotationAxis = None
        self.moveVector = None
        
#------------------------------------------------------------------------------
        
        
        



#------------------------------------------------------------------------------
# Code to make system useable within FreeCAD
#------------------------------------------------------------------------------
def solveConstraints2( doc, cache=None ): #cache because of compatibility to hamish...
    ss = SolverSystem()
    ss.loadSystem(doc)
    t1 = time.time()
    for i in range(0,3000): #adapt for your needs
        ss.doSolverStep(doc)
    t2 = time.time()
    FreeCAD.Console.PrintMessage( "Time for iterating: "+  str(t2-t1) )
    ss.solutionToParts(doc)
    


class a3_SolverCommand:
    def Activated(self):
        solveConstraints2( FreeCAD.ActiveDocument ) #the new iterative solver

    def GetResources(self): 
        return {
            #'Pixmap' : path_a3 + '/icons/a3_solver.svg', 
            'MenuText': 'Solve', 
            'ToolTip': 'Solve Assembly 2 constraints'
            } 

FreeCADGui.addCommand('a3_SolverCommand', a3_SolverCommand())
#------------------------------------------------------------------------------