from math import *
import numpy as np

#Paramètres des ressorts
kBD, LBD = 1, 1
kCD, LCD = 1, 2
kAC, LAC = 1, sqrt(2)
kAE, LAE = 1, 1
kBE, LBE = 1, 2
kCE, LCE = 1, 1
kDE, LDE = 1, sqrt(5)

def norm(x, y):
    """
    Renvoie la norme 2 du vecteur (x,y)^T
    """
    return np.linalg.norm(np.array([x,y]))





def GradEpot(q):
    """
    Renvoie le gradient de l'énergie potentielle des ressorts "nabla V(q) = -F"
    le système ayant la configuration de position q
    q est un tableau de taille 6
    """
    xA, yA = 0, 0
    xB, yB = 3, 0
    xC, yC = q[0], q[1]
    xD, yD = q[2], q[3]
    xE, yE = q[4], q[5]
    AB = norm(xB - xA, yB - yA)
    AC = norm(xC - xA, yC - yA)
    CD = norm(xD - xC, yD - yC)
    BD = norm(xD - xB, yD - yB)
    AE = norm(xE - xA, yE - yA)
    BE = norm(xE - xB, yE - yB)
    CE = norm(xE - xC, yE - yC)
    DE = norm(xE - xD, yE - yD)
    F = [   kAC * (1 - LAC/AC) * (xC-xA) + kCD * (1 - LCD/CD)
            * (xC-xD) + kCE * (1 - LCE/CE) * (xC-xE), 
            kAC * (1 - LAC/AC) * (yC-yA) + kCD * (1 - LCD/CD) * (yC-yD) + kCE * (1 - LCE/CE) * (yC-yE),
            kCD * (1 - LCD/CD) * (xD-xC) + kBD * (1 - LBD/BD) * (xD-xB) + kDE * (1 - LDE/DE) * (xD-xE),
            kCD * (1 - LCD/CD) * (yD-yC) + kBD * (1 - LBD/BD) * (yD-yB) + kDE * (1 - LDE/DE) * (yD-yE),
            kAE * (1 - LAE/AE) * (xE-xA) + kBE * (1 - LBE/BE) * (xE-xB) + kCE * (1 - LCE/CE) * (xE-xC) + kDE * (1 - LDE/DE) * (xE-xD), 
            kAE * (1 - LAE/AE) * (yE-yA) + kBE * (1 - LBE/BE) * (yE-yB) + kCE * (1 - LCE/CE) * (yE-yC) + kDE * (1 - LDE/DE) * (yE-yD)]
    return np.array(F)

def Epot(q):
    """
    Renvoie l'énergie potentielle des ressorts, le système ayant la configuration de position q
    """
    xA, yA = 0, 0
    xB, yB = 3, 0
    xC, yC = q[0], q[1]
    xD, yD = q[2], q[3]
    xE, yE = q[4], q[5]
    AB = norm(xB - xA, yB - yA)
    AC = norm(xC - xA, yC - yA)
    CD = norm(xD - xC, yD - yC)
    BD = norm(xD - xB, yD - yB)
    AE = norm(xE - xA, yE - yA)
    BE = norm(xE - xB, yE - yB)
    CE = norm(xE - xC, yE - yC)
    DE = norm(xE - xD, yE - yD)
    e = 0.5 * (kAC * (LAC - AC)**2 + kCD * (LCD - CD)**2 + kBD * (LBD - BD)**2 + kAE * (LAE - AE)**2 + kBE * (LBE - BE)**2 + kCE * (LCE - CE)**2 + kDE * (LDE - DE)**2);
    return e

def Ekin(p):
    """
    Renvoie l'énergie cinétique des ressorts, le système ayant la configuration de moments p
    p est un tableau de taille 6
    """
    m = 1
    dxC = p[0]
    dyC = p[1]
    dxD = p[2]
    dyD = p[3]
    dxE = p[4]
    dyE = p[5]
    vC = norm(dxC, dyC)
    vD = norm(dxD, dyD)
    vE = norm(dxE, dyE)
    e = 0.5 * m * (vC**2 + vD**2 + vE**2)
    return e

def H(p,q):
    """
    Renvoie le Hamiltonien du système, le système ayant la configuration de position q et la configuration de moments p
    """
    return Epot(q) + Ekin(p)

def Geometry(q):
    """
    Prépare le tracé de la configuration dans le plan R^2
    """
    xA, yA = 0, 0
    xB, yB = 3, 0
    xC, yC = q[0], q[1]
    xD, yD = q[2], q[3]
    xE, yE = q[4], q[5]
    X = [xB, xD, xC, xA, xE, xB, xD, xE, xC]
    Y = [yB, yD, yC, yA, yE, yB, yD, yE, yC]
    return X, Y