#! /usr/bin/env python2

import numpy as np
from scipy.integrate import odeint

def Data_init(**kwargs):
  """
  if there is a kwarg['file'], then read from file, otherwise find
  kwargs['dimensions'] and kwargs['particles'] and randomize data.  If
  this fails, print an error.

  kwargs['time'] is a tuple of final time and number of steps.
  """
  if 'file' in kwargs:
    print "Reading the file"
  else:
    print "Randomizing the initial data"
    X = np.random.rand(kwargs['particles'],kwargs['dimensions']) * 2 - 1
    V = np.random.rand(kwargs['particles'],kwargs['dimensions']) * 2 - 1
    M = np.random.rand(kwargs['particles'])

  t_f,num = kwargs['time']
  t = np.linspace(0,t_f,num)

  return X,V,M,t

def Gravitational_force(X,M):
  # Expand the collections
  x1, x2 = X
  m1, m2 = M

  # Calculate a displacement vector
  r = x2-x1
  # Calculate it's magnitude
  r_mod = np.sqrt(r.dot(r))
  # Calculate the force itself
  f = m1*m2*r / np.power(r_mod,3)

  return f

def Net_force_1(X,M):
  F = np.zeros(X.shape)
  for i in xrange(X.shape[0]):
    for j in xrange(X.shape[0]):
      if i == j:
        continue
      F += Gravitational_force(
          [ X[i,:], X[j,:] ],
          [ M[i], M[j] ]
          )

  return F

def ODE(W,t,M):
  """
  This is the function which is needed for odeint function
  """
  X, V = W
  F = Net_force_1(X,M)
  function = []
  for i in xrange(X.shape[0]):
    for j in xrange(X.shape[1]):
      function += [
          V[i,j],F[i,j]
          ]
  return function

# Initiate the data:
X0,V0,M,t = Data_init(particles=4,dimensions=2,time=[10,255])

# Solve the ODE system
solution = odeint(ODE,[X0,V0],t,args=(M,))