import numpy
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy.linalg


class Solide:
    def __init__(self,Rx,Ry,theta,contour=[]):
        self.Rx = Rx
        self.Ry = Ry
        self.theta = theta
        self.contour = contour
        self.line = None
    def coord_point(self,xp,yp):
        # xp,yp : coordonnées d'un point lié au solide dans le repère propre
        cos = numpy.cos(self.theta)
        sin = numpy.sin(self.theta)
        X = self.Rx + xp*cos-yp*sin
        Y = self.Ry + xp*sin+yp*cos
        return [X,Y]
    def dessiner(self,ax):
        liste_x = []
        liste_y = []
        for k in range(len(self.contour)):
            p = self.contour[k]
            XY = self.coord_point(p[0],p[1])
            liste_x.append(XY[0])
            liste_y.append(XY[1])
        self.line, = ax.plot(liste_x,liste_y,"-",linewidth=2)
    def redessiner(self,ax):
        liste_x = []
        liste_y = []
        for k in range(len(self.contour)):
            p = self.contour[k]
            XY = self.coord_point(p[0],p[1])
            liste_x.append(XY[0])
            liste_y.append(XY[1])
        self.line.set_data(liste_x,liste_y)


class Contraintes:
    def __init__(self):
        pass
    def c(q): # définition de la contrainte (c(q)=0)
        pass
        # q : [Rx,Ry,theta,...] listes à plat des coordonnées pour les N solides
    def Jc(q): # définition des lignes du jacobien de C
        pass

class SolideFixe(Contraintes):
    def __init__(self,i,Rx,Ry,theta):
        self.i = i
        self.Rx = Rx
        self.Ry = Ry
        self.theta = theta
    def c(self,q):
        c1 = q[3*self.i]-self.Rx
        c2 = q[3*self.i+1]-self.Ry
        c3 = q[3*self.i+2]-self.theta
        return numpy.array([c1,c2,c3])
    def Jc(self,q):
        Jc1 = numpy.zeros(len(q))
        Jc1[3*self.i] = 1
        Jc2 = numpy.zeros(len(q))
        Jc2[3*self.i+1] = 1
        Jc3 = numpy.zeros(len(q))
        Jc3[3*self.i+2] = 1
        return numpy.array([Jc1,Jc2,Jc3])

class Pivot(Contraintes):
    def __init__(self,i,j,xpi,ypi,xpj,ypj):
        self.i = i
        self.j = j
        self.xpi = xpi
        self.ypi = ypi
        self.xpj = xpj
        self.ypj = ypj
    def c(self,q):
        c1 = q[3*self.i]+self.xpi*numpy.cos(q[3*self.i+2])-self.ypi*numpy.sin(q[3*self.i+2])-q[3*self.j]-self.xpj*numpy.cos(q[3*self.j+2])+self.ypj*numpy.sin(q[3*self.j+2])
        c2 = q[3*self.i+1]+self.xpi*numpy.sin(q[3*self.i+2])+self.ypi*numpy.cos(q[3*self.i+2])-q[3*self.j+1]-self.xpj*numpy.sin(q[3*self.j+2])-self.ypj*numpy.cos(q[3*self.j+2])
        return numpy.array([c1,c2])
    def Jc(self,q):
        Jc1 = numpy.zeros(len(q))
        Jc2 = numpy.zeros(len(q))
        Jc1[3*self.i] = 1
        Jc1[3*self.i+2] = -self.xpi*numpy.sin(q[3*self.i+2])-self.ypi*numpy.cos(q[3*self.i+2])
        Jc1[3*self.j] = -1
        Jc1[3*self.j+2] = self.xpj*numpy.sin(q[3*self.j+2])-self.ypj*numpy.cos(q[3*self.j+2])
        Jc2[3*self.i+1] = 1
        Jc2[3*self.i+2] = self.xpi*numpy.cos(q[3*self.i+2])-self.ypi*numpy.sin(q[3*self.i+2])
        Jc2[3*self.j+1] = -1
        Jc2[3*self.j+2] = -self.xpj*numpy.cos(q[3*self.j+2])+self.ypj*numpy.sin(q[3*self.j+2])
        return numpy.array([Jc1,Jc2])

class Glissiere(Contraintes):
    def __init__(self,i,j,xpi,ypi,xqi,yqi,xpj,ypj,angle):
        self.i = i
        self.j = j
        self.xpi = xpi
        self.ypi = ypi
        self.xqi = xqi
        self.yqi = yqi
        self.xpj = xpj
        self.ypj = ypj
        self.angle = angle
    def c(self,q):
        c1 = q[3*self.i+2]-q[3*self.j+2]-self.angle
        u = q[3*self.i]+self.xpi*numpy.cos(q[3*self.i+2])-self.ypi*numpy.sin(q[3*self.i+2])-q[3*self.j]-self.xpj*numpy.cos(q[3*self.j+2])+self.ypj*numpy.sin(q[3*self.j+2])
        v = q[3*self.i+1]+self.xpi*numpy.sin(q[3*self.i+2])+self.ypi*numpy.cos(q[3*self.i+2])-q[3*self.j+1]-self.xpj*numpy.sin(q[3*self.j+2])-self.ypj*numpy.cos(q[3*self.j+2])
        a = (self.xpi-self.xqi)*numpy.cos(q[3*self.i+2])-(self.ypi-self.yqi)*numpy.sin(q[3*self.i+2])
        b = (self.xpi-self.xqi)*numpy.sin(q[3*self.i+2])+(self.ypi-self.yqi)*numpy.cos(q[3*self.i+2])
        c2 = u*a+v*b
        return numpy.array([c1,c2])
    def Jc(self,q):
        Jc1 = numpy.zeros(len(q))
        Jc2 = numpy.zeros(len(q))
        Jc1[3*self.i+2] = 1
        Jc1[3*self.j+2] = -1
        u = q[3*self.i]+self.xpi*numpy.cos(q[3*self.i+2])-self.ypi*numpy.sin(q[3*self.i+2])-q[3*self.j]-self.xpj*numpy.cos(q[3*self.j+2])+self.ypj*numpy.sin(q[3*self.j+2])
        v = q[3*self.i+1]+self.xpi*numpy.sin(q[3*self.i+2])+self.ypi*numpy.cos(q[3*self.i+2])-q[3*self.j+1]-self.xpj*numpy.sin(q[3*self.j+2])-self.ypj*numpy.cos(q[3*self.j+2])
        a = (self.xpi-self.xqi)*numpy.cos(q[3*self.i+2])-(self.ypi-self.yqi)*numpy.sin(q[3*self.i+2])
        b = (self.xpi-self.xqi)*numpy.sin(q[3*self.i+2])+(self.ypi-self.yqi)*numpy.cos(q[3*self.i+2])
        Jc2[3*self.i] = a
        Jc2[3*self.j] = -a
        Jc2[3*self.i+1] = b
        Jc2[3*self.j+1] = -b
        Jc2[3*self.i+2] = (-self.xpi*numpy.sin(q[3*self.i+2])-self.ypi*numpy.cos(q[3*self.i+2]))*a\
                        +(-(self.xpi-self.xqi)*numpy.sin(q[3*self.i+2])-(self.ypi-self.yqi)*numpy.cos(q[3*self.i+2]))*u\
                        +(self.xpi*numpy.cos(q[3*self.i+2])-self.ypi*numpy.sin(q[3*self.i+2]))*b\
                        +((self.xpi-self.xqi)*numpy.cos(q[3*self.i+2])-(self.ypi-self.yqi)*numpy.sin(q[3*self.i+2]))*v
        Jc2[3*self.j+2] = (self.xpj*numpy.sin(q[3*self.j+2])+self.ypj*numpy.cos(q[3*self.j+2]))*a\
                        + (-self.xpj*numpy.cos(q[3*self.j+2])+self.ypj*numpy.sin(q[3*self.j+2]))*b
        return numpy.array([Jc1,Jc2])
        

class Angle(Contraintes):
    def __init__(self,i,theta):
        self.i = i
        self.theta = theta
    def c(self,q):
        c1 = q[3*self.i+2]-self.theta
        return numpy.array([c1])
    def Jc(self,q):
        Jc1 = numpy.zeros(len(q))
        Jc1[3*self.i+2] = 1
        return numpy.array([Jc1])

class GlissiereHorizontale(Contraintes):
    def __init__(self,i,Y):
        self.i = i
        self.Y = Y
    def c(self,q):
        c1 = q[3*self.i+1]-self.Y
        c2 = q[3*self.i+2]
        return numpy.array([c1,c2])
    def Jc(self,q):
        Jc1 = numpy.zeros(len(q))
        Jc2 = numpy.zeros(len(q))
        Jc1[3*self.i+1] = 1
        Jc2[3*self.i+2] = 1
        return numpy.array([Jc1,Jc2])


class Systeme:
    def __init__(self):
        self.liste_solides = []
        self.liste_contraintes = []
        self.n = 0
    def ajouter_solide(self,solide):
        self.liste_solides.append(solide)
        self.n += 3
    def ajouter_contrainte(self,contrainte):
        self.liste_contraintes.append(contrainte)
    def dessiner(self,ax):
        for solide in self.liste_solides:
            solide.dessiner(ax)
    def redessiner(self,ax):
        for solide in self.liste_solides:
            solide.redessiner(ax)
    def coordonnees(self):
        q = []
        for solide in self.liste_solides:
            q.append(solide.Rx)
            q.append(solide.Ry)
            q.append(solide.theta)
        return q
    def jacobien(self,q):
        J = numpy.zeros((self.n,self.n))
        i = 0
        for contrainte in self.liste_contraintes:
            Jc = contrainte.Jc(q)
            for k in range(len(Jc)):
                J[i] = Jc[k]
                i += 1
        if (i<self.n):
            print("Nombre de contraintes insuffisant")
        return J
    def C(self,q):
        C = numpy.zeros(self.n)
        i = 0
        for contrainte in self.liste_contraintes:
            c = contrainte.c(q)
            for k in range(len(c)):
                C[i] = c[k]
                i +=1
        return C
    def calcul(self,tolerance,max_iter=100):
        iteration = True
        q = self.coordonnees()
        ite = 0
        while iteration:
            J = self.jacobien(q)
            C = self.C(q)
            dq = numpy.linalg.solve(J,-C)
            if numpy.dot(dq,dq) < tolerance:
                iteration = False
            ite += 1
            if ite > max_iter:
                iteration = False
                print("maximum d'itérations atteint")
            q += dq
        i = 0
        for solide in self.liste_solides:
            solide.Rx = q[3*i]
            solide.Ry = q[3*i+1]
            solide.theta = q[3*i+2]
            i += 1