#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Demonstrateur OpenCL d'interaction NCorps

Emmanuel QUEMENER <emmanuel.quemener@ens-lyon.fr> CeCILLv2
"""
import getopt
import sys
import time
import numpy as np
import pyopencl as cl
import pyopencl.array as cl_array
from numpy.random import randint as nprnd


BlobOpenCL= """
#define znew  ((z=36969*(z&65535)+(z>>16))<<16)
#define wnew  ((w=18000*(w&65535)+(w>>16))&65535)
#define MWC   (znew+wnew)
#define SHR3  (jsr=(jsr=(jsr=jsr^(jsr<<17))^(jsr>>13))^(jsr<<5))
#define CONG  (jcong=69069*jcong+1234567)
#define KISS  ((MWC^CONG)+SHR3)

#define MWCfp MWC * 2.328306435454494e-10f
#define KISSfp KISS * 2.328306435454494e-10f
#define SHR3fp SHR3 * 2.328306435454494e-10f
#define CONGfp CONG * 2.328306435454494e-10f

#define LENGTH 1.

#define PI 3.14159265359

#define SMALL_NUM 0.000000001

float Gravity(float4 m,float4 n)
{
    return((float)(1./pow(distance(m,n),2)));
}

float8 RungeKutta(__global float8* clDataIn,int gid,float h)
{
    float4 c[4],d[4],ct,dt;

    c[0]=h*clDataIn[gid].hi;
    c[1]=h*(clDataIn[gid].hi+c[0]/(float)2.);
    c[2]=h*(clDataIn[gid].hi+c[1]/(float)2.);
    c[3]=h*(clDataIn[gid].hi+c[2]);

    d[0]=(0.,0.,0.,0.);
    for (int i=0;i<get_global_size(0);i++)
    {
        if (gid != i)
        d[0]+=Gravity(clDataIn[gid].lo,clDataIn[i].lo);
    }
    d[0]*=h;
    d[1]=(0.,0.,0.,0.);
    for (int i=0;i<get_global_size(0);i++)
    {
        if (gid != i)
        d[1]+=Gravity(clDataIn[gid].lo+d[0]/(float)2.,clDataIn[i].lo);
    }
    d[1]*=h;
    d[2]=(0.,0.,0.,0.);
    for (int i=0;i<get_global_size(0);i++)
    {
        if (gid != i)
        d[2]+=Gravity(clDataIn[gid].lo+d[1]/(float)2.,clDataIn[i].lo);
    }
    d[2]*=h;
    d[3]=(0.,0.,0.,0.);
    for (int i=0;i<get_global_size(0);i++)
    {
        if (gid != i)
        d[3]+=Gravity(clDataIn[gid].lo+d[2],clDataIn[i].lo);
    }
    d[3]*=h;
    
//    ct=clDataIn[gid].lo+(c[0]+(float)2.*c[1]+(float)2.*c[2]+c[3])/(float)6.;
//    dt=clDataIn[gid].hi+(d[0]+(float)2.*d[1]+(float)2.*d[2]+d[3])/(float)6.;
    if (gid != 0)
        ct=(float4)distance(clDataIn[gid],clDataIn[0]);
    else
        ct=(float4)3.14159;
    dt=(float4)0;
    
    return((float8)(ct.s0,ct.s1,ct.s2,ct.s3,dt.s0,dt.s1,dt.s2,dt.s3));
}

__kernel void SplutterPoints(__global float8* clData, float box, float velocity,
                               uint seed_z,uint seed_w)
{
    int gid = get_global_id(0);
    uint z=seed_z+(uint)gid;
    uint w=seed_w-(uint)gid;
    float theta=MWCfp*PI;
    float phi=MWCfp*PI*2.;
    float sinTheta=sin(theta);
    clData[gid].s01234567 = (float8) (box*MWCfp,box*MWCfp,box*MWCfp,0.,0.,0.,0.,0.);
    clData[gid].s4=velocity*sinTheta*cos(phi);
    clData[gid].s5=velocity*sinTheta*sin(phi);
    clData[gid].s6=velocity*cos(theta);
}

__kernel void Evolution(__global float8* clDataOut,__global float8* clDataIn, float h)
{
    int gid = get_global_id(0);
    
    clDataOut[gid]=RungeKutta(clDataIn,gid,h);
    barrier(CLK_GLOBAL_MEM_FENCE);
}

__kernel void Commit(__global float8* clDataOut,__global float8* clDataIn)
{
    int gid = get_global_id(0);
    
    clDataIn[gid]=clDataOut[gid];
    barrier(CLK_GLOBAL_MEM_FENCE);
}

"""

if __name__=='__main__':
    
    # Set defaults values
  
    # Id of Device : 1 is for first find !
    Device=2
    # Iterations is integer
    Number=2
    # Size of box
    SizeOfBox=1000.
    # Initial velocity of particules
    Velocity=10.
    # Redo the last process
    Redo=1
    # Step
    Step=1.

    HowToUse='%s -d <DeviceId> -n <NumberOfSegments> -z <SizeOfBox> -v <Velocity> -s <Step>'

    try:
        opts, args = getopt.getopt(sys.argv[1:],"hd:n:z:v:r:",["device=","number=","size=","velocity=","redo=","step=1"])
    except getopt.GetoptError:
        print HowToUse % sys.argv[0]
        sys.exit(2)

    for opt, arg in opts:
        if opt == '-h':
            print HowToUse % sys.argv[0]

            print "\nInformations about devices detected under OpenCL:"
            try:
                Id=1
                for platform in cl.get_platforms():
                    for device in platform.get_devices():
                        deviceType=cl.device_type.to_string(device.type)
                        print "Device #%i from %s of type %s : %s" % (Id,platform.vendor.lstrip(),deviceType,device.name.lstrip())
                        Id=Id+1
                sys.exit()
            except ImportError:
                print "Your platform does not seem to support OpenCL"
                sys.exit()

        elif opt in ("-d", "--device"):
            Device=int(arg)
        elif opt in ("-n", "--number"):
            Number=int(arg)
        elif opt in ("-z", "--size"):
            SizeOfBox=np.float32(arg)
        elif opt in ("-v", "--velocity"):
            Velocity=np.float32(arg)
        elif opt in ("-s", "--step"):
            Step=np.float32(arg)
        elif opt in ("-r", "--redo"):
            Redo=int(arg)
            
    print "Device choosed : %s" % Device
    print "Number of segments : %s" % Number
    print "Size of Box : %s" % SizeOfBox
    print "Initial velocity % s" % Velocity
    print "Redo the last process % s" % Redo
    print "Step of iteration % s" % Step
    
    MyData = np.zeros(Number, dtype=cl_array.vec.float8)

    Id=1
    HasXPU=False
    for platform in cl.get_platforms():
        for device in platform.get_devices():
            if Id==Device:
                PlatForm=platform
                XPU=device
                print "CPU/GPU selected: ",device.name.lstrip()
                HasXPU=True
            Id+=1

    if HasXPU==False:
        print "No XPU #%i found in all of %i devices, sorry..." % (Device,Id-1)
        sys.exit()      

    # Je cree le contexte et la queue pour son execution
    try:
        ctx = cl.Context([XPU])
        queue = cl.CommandQueue(ctx,properties=cl.command_queue_properties.PROFILING_ENABLE)
    except:
        print "Crash during context creation"
   

    MyRoutines = cl.Program(ctx, BlobOpenCL).build()
    MyData[1][0]=1.

    mf = cl.mem_flags
    # clDataIn = cl.Buffer(ctx, mf.READ_WRITE, MyData.nbytes)
    clDataIn = cl.Buffer(ctx, mf.WRITE_ONLY|mf.COPY_HOST_PTR,hostbuf=MyData)
    clDataOut = cl.Buffer(ctx, mf.READ_WRITE, MyData.nbytes)

    print 'Tous au meme endroit',MyData

    # MyRoutines.SplutterPoints(queue,(Number,1),None,clDataIn,np.float32(SizeOfBox),np.float32(Velocity),np.uint32(nprnd(2**32)),np.uint32(nprnd(2**32)))

    print 'Tous distribues',MyData

    CLLaunch=MyRoutines.Evolution(queue,(Number,1),None,clDataOut,clDataIn,np.float32(Step))
    CLLaunch.wait()
    cl.enqueue_copy(queue, MyData, clDataOut)
    print 'Tous calcules',MyData
    CLLaunch=MyRoutines.Commit(queue,(Number,1),None,clDataOut,clDataIn)
    CLLaunch.wait()

    time_start=time.time()
    for i in xrange(Redo):
        #CLLaunch=MyRoutines.ShortestDistance(queue, (Number,Number), None, clData, clDistance)
        sys.stdout.write('.')
        #CLLaunch.wait()
    print "\nDuration on %s for each %s" % (Device,(time.time()-time_start)/Redo)

    clDataIn.release()
    clDataOut.release()