# -*- coding: utf-8 -*-
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# solve the equations of motion for a number of vortices
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# nvort : number of vortices
# charge [array] : charges of vortices
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
#
from numpy import array
from scipy.integrate import odeint
from pylab import *


# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# return derivatives for the equation of motion
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
def deriv(yin,t):
    
    x1 = array(zeros(nvort))       # positions of vortices
    y1 = array(zeros(nvort))
    yout = array(zeros(2*nvort))   # output of function
    
    for iv in range(nvort): # produce arrays for x and y coordinates
        x1[iv] = yin[2*iv]
        y1[iv] = yin[2*iv+1]
    
    
    for iv in range(nvort):         # rhs of equations
        dx1 = 0
        dy1 = 0
        for iv1 in range(nvort):
            if iv1 != iv:
                dist2 = (x1[iv]-x1[iv1])**2+(y1[iv]-y1[iv1])**2
                dx1 = dx1 + charge[iv1]*(y1[iv]-y1[iv1])/dist2
                dy1 = dy1 - charge[iv1]*(x1[iv]-x1[iv1])/dist2
        
        yout[2*iv] = dx1
        yout[2*iv+1] = dy1
        
    return yout
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -


# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# calculate conserved and other quantities
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
def info():
    x1 = array(zeros(nvort))
    y1 = array(zeros(nvort))
    
    for iv in range(nvort): # produce arrays for x and y coordinates
        x1[iv] = pos0[2*iv]
        y1[iv] = pos0[2*iv+1]
    
    total_charge = 0        # total charge of vortex system
    ix = 0                  # guiding center
    iy = 0
    angmom = 0              # angular momentum
    for iv in range(nvort):
        ix = ix + charge[iv]*x1[iv]
        iy = iy + charge[iv]*y1[iv]
        angmom = angmom + 0.5*charge[iv]*(x1[iv]**2+y1[iv]**2)
    
    if total_charge != 0:    # no guiding center if total charge = 0
        ix = ix/total_charge
        iy = iy/total_charge
    
    energy = 0               # energy
    for iv in range(nvort):
        for iv1 in range(nvort):
            if iv1 != iv:
                dist2 = (x1[iv]-x1[iv1])**2+(y1[iv]-y1[iv1])**2
                energy = energy - 0.5*charge[iv]*charge[iv1]*log(dist2)
    
    return energy,ix,iy,angmom
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -


# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# main program
# - - - - - - - - - - - - - - - - - - - - - - - - - - - -
#
# read data: positions of vortices
fil = open('vortices.dat','r')
lines = fil.readlines()
fil.close()

nvort = len(lines)                # number of vortices
pos0 = array(zeros(2*nvort))      # initialize arrays
charge = array(zeros(nvort))

# set initial values for odeint
iv = 0
for line in lines:
    temp,xx,yy = map(float, line.split())
    charge[iv] = temp
    pos0[2*iv] = xx
    pos0[2*iv+1] = yy
    iv = iv + 1

# print conserved quantities at beginning
conserve = info()
print conserve

time_fin = 40.0
time = arange(0.0,time_fin,0.2)
pos = odeint(deriv,pos0,time)     # numerical integration

# plot trajectories of vortices
figure(figsize=(6,6))             # prepare graph
xlabel("x")
ylabel("y")
title("trajectories of vortices")

for iv in range(nvort):
    plot(pos[:,2*iv],pos[:,2*iv+1])
show()

# plot coordinates of vortices (e.g., vortex 1) with time
figure(figsize=(6,6))             # prepare graph
xlabel("time")
ylabel("x, y")

#for iv in range(nvort):
#iv = 0
#plot(time,pos[:,2*iv],'-r')
#plot(time,pos[:,2*iv+1],'-b')
#show()
iv = 1
plot(time,pos[:,2*iv],'--r')
plot(time,pos[:,2*iv+1],'--b')
show()

# print conserved quantities at final time
fin = size(pos)/(2*nvort)           # final vortex positions
pos0 = pos[fin-1]
    
conserve = info()
print conserve