"""Colliding Elastic Balls"""

import numpy

# SPH equations
from pysph.sph.equation import Group
from pysph.sph.basic_equations import IsothermalEOS, ContinuityEquation, MonaghanArtificialViscosity,\
     XSPHCorrection, VelocityGradient2D
from pysph.sph.solid_mech.basic import MomentumEquationWithStress2D, HookesDeviatoricStressRate2D,\
    MonaghanArtificialStress

from pysph.base.utils import get_particle_array
from pysph.base.kernels import CubicSpline, WendlandQuintic
from pysph.solver.application import Application
from pysph.solver.solver import Solver
from pysph.sph.integrator import PECIntegrator
from pysph.sph.integrator_step import SolidMechStep

def get_K(G, nu):
    ''' Get the bulk modulus from shear modulus and Poisson ratio '''
    return 2.0*G*(1+nu)/(3*(1-2*nu))

def add_properties(pa, *props):
    for prop in props:
        pa.add_property(name=prop)


# constants
E = 1e7
nu = 0.3975
G = E/(2.0*(1+nu))
K = get_K(G, nu)
rho0 = 1.0
c0 = numpy.sqrt(K/rho0)

dx = 0.0005
hdx = 1.5

# geometry
ri = 0.03
ro = 0.04

def create_particles(**kwargs):
    #x,y = numpy.mgrid[-1.05:1.05+1e-4:dx, -0.105:0.105+1e-4:dx]
    spacing = 0.041 # spacing = 2*5cm

    x,y = numpy.mgrid[-ro:ro:dx, -ro:ro:dx]
    x = x.ravel()
    y = y.ravel()

    d = (x*x+y*y)
    keep = numpy.flatnonzero((ri*ri<=d) * (d<ro*ro))
    x = x[keep]
    y = y[keep]

    x = numpy.concatenate([x-spacing,x+spacing])
    y = numpy.concatenate([y,y])

    print 'Ellastic Collision with %d particles'%(x.size)
    print "Shear modulus G = %g, Young's modulus = %g, Poisson's ratio =%g"%(G,E,nu)

    #print bdry, numpy.flatnonzero(bdry)
    m = numpy.ones_like(x)*dx*dx
    h = numpy.ones_like(x)*hdx*dx
    rho = numpy.ones_like(x)
    z = numpy.zeros_like(x)

    p = 0.5*1.0*100*100*(1 - (x**2 + y**2))

    cs = numpy.ones_like(x) * 10000.0

    # u is set later
    v = z
    u_f = 0.059

    p *= 0
    h *= 1

    # create the particle array
    pa = get_particle_array(
        name="solid", x=x+spacing, y=y, m=m, rho=rho, h=h,
        p=p, cs=cs, u=z, v=v)

    pa.cs[:] = c0
    pa.u = pa.cs*u_f*(2*(x<0)-1)

    # add requisite properties

    # sound speed etc.
    add_properties(pa, 'cs', 'e' )

    # velocity gradient properties
    add_properties(pa, 'v00', 'v01', 'v10', 'v11')

    # artificial stress properties
    add_properties(pa, 'r00', 'r01', 'r02', 'r11', 'r12', 'r22')

    # deviatoric stress components
    add_properties(pa, 's00', 's01', 's02', 's11', 's12', 's22')

    # deviatoric stress accelerations
    add_properties(pa, 'as00', 'as01', 'as02', 'as11', 'as12', 'as22')

    # deviatoric stress initial values
    add_properties(pa, 's000', 's010', 's020', 's110', 's120', 's220')

    # standard acceleration variables
    add_properties(pa, 'arho', 'au', 'av', 'aw', 'ax', 'ay', 'az', 'ae')

    # initial values
    add_properties(pa, 'rho0', 'u0', 'v0', 'w0', 'x0', 'y0', 'z0', 'e0')

    # load balancing properties
    pa.set_lb_props( pa.properties.keys() )

    return [pa,]

# create the Application
app = Application()

# kernel
kernel = CubicSpline(dim=2)
wdeltap = kernel.kernel(rij=dx, h=hdx*dx)

# integrator
integrator = PECIntegrator(solid=SolidMechStep())

# Create a solver
solver = Solver(kernel=kernel, dim=2, integrator=integrator)

# default parameters
dt = 1e-8
tf = 5e-5
solver.set_time_step(dt)
solver.set_final_time(tf)
solver.set_print_freq(500)

# add the equations
equations = [

    # Properties computed set from the current state
    Group(
        equations=[
            # p
            IsothermalEOS(dest='solid', sources=None, rho0=rho0, c0=c0),

            # vi,j : requires properties v00, v01, v10, v11
            VelocityGradient2D(dest='solid', sources=['solid',]),

            # rij : requires properties r00, r01, r02, r11, r12, r22,
            #                           s00, s01, s02, s11, s12, s22
            MonaghanArtificialStress(
             dest='solid', sources=None, eps=0.3),
            ],
        ),

    # Acceleration variables are now computed
    Group(
        equations=[

            # arho
            ContinuityEquation(dest='solid', sources=['solid',]),

            # au, av
            MomentumEquationWithStress2D(
                dest='solid', sources=['solid',], n=4, wdeltap=wdeltap),

            # au, av
            MonaghanArtificialViscosity(
                dest='solid', sources=['solid',], alpha=1.0, beta=1.0),

            # a_s00, a_s01, a_s11
            HookesDeviatoricStressRate2D(
                dest='solid', sources=None, shear_mod=G),

            # ax, ay, az
            XSPHCorrection(
                dest='solid', sources=['solid',], eps=0.5),

            ]

        ) # End Acceleration Group

    ] # End Group list

# Setup the application and solver.  This also generates the particles.
app.setup(solver=solver, equations=equations,
          particle_factory=create_particles,
          name='fluid')

# run
app.run()
