from sympy import diff, symbols, sqrt, Symbol
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.cm as cm
import matplotlib.colors as colors


#***************YOUR INPUT HERE*************#

sig, eps = 1, 2
x, y, z = symbols('x, y, z')
r = sqrt(x**2 + y**2 + z**2)
V = 4*eps * ((sig/r)**12 - (sig/r)**6)
vdiff = [diff(V, x), diff(V, y), diff(V, z)]

#**************END YOUR INPUT HERE***********#

#Here we input the coordinates of the atoms. You can have as many atoms as you want, 
#and change the coordinates to whatever you like

atoms = np.array([[0.,0.,0.], [0.,0.,1.], [0.,1.,0.], [1.,0.,0.], [1.,1.,0.], [1.,0.,1.], [0.,1.,1.], [1.,1.,1.], [0.5, 0.5, 0.5]])

#We write the output to a textfile

file = open("output.txt", "w")

file.write('    Atom 1    ' + '    Atom2   ' + '    Force(x-dir)   Force(y-dir)   Force(z-dir) ' + '  Total Force(abs value) ' + '\n')

F_list = []
for i in range(len(atoms)):
    for j in range(i):
        res = atoms[i] - atoms[j]
        Fxyz = [tv.subs([(x,res[0]), (y,res[1]),  (z,res[2])]) for tv in vdiff]
        print("Between", atoms[i], "and", atoms[j], "we have the forces", Fxyz)
        F = sqrt(Fxyz[0]**2 + Fxyz[1]**2 + Fxyz[2]**2)
        F_list.append(F)
        print("and the total force is", F)
        file.write(str(atoms[i]) + " " + str(atoms[j]))
        file.write( "%14.2f %14.2f %14.2f %14.2f" % (Fxyz[0], Fxyz[1], Fxyz[2], F) + '\n')

cmap = get_cmap('cool')
norm = colors.LogNorm(min(F_list), max(F_list))
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
f = 0
for i in range(len(atoms)):
    for j in range(i):
        print(F_list[f]/max(F_list))
        z = np.array([atoms[i,2], atoms[j,2]])
        x = np.array([atoms[i,0], atoms[j,0]])
        y = np.array([atoms[i,1], atoms[j,1]])
        ax.plot(x, y, z, color=cmap(norm(float(F_list[f]))))
        f += 1

plt.show()