import pycuda.autoinit
import pycuda.driver as cuda
from pycuda.compiler import SourceModule
import numpy
import math
import time

class FpuCuda:
    def __init__(self,nvar):
        self.nvar = nvar
        f1 = """__device__ void dPotential(int k, int nvar, float b, float n, float *q, float *dp) {
            float x1,x2;
            if (k==0) {x1=q[0]; x2=q[1]-q[0];}
            else if (k==nvar-1) {x1 = q[k]-q[k-1]; x2=-q[k];}
            else {x1=q[k]-q[k-1]; x2=-q[k]+q[k+1];}
            dp[k] = x1-x2+b*(powf(x1,n)-powf(x2,n));
        }"""
            
        f2 = """__device__ void verletStep(int k, float h, int nvar, float b, float n, float *q, float *p, float *dp) {
            dPotential(k,nvar,b,n,q,dp);
            p[k] -= 0.5*h*dp[k];
            q[k] += h*p[k];
            __syncthreads();
            dPotential(k,nvar,b,n,q,dp);
            p[k] -= 0.5*h*dp[k];
            __syncthreads();
        }"""    
            
        f3 = """__device__ void energy(int nvar,float b, float n, float *q, float *p, float *e) {
            float x1,x2;
            int k = threadIdx.x;
            if (k==0) {x1=q[k]; x2=0.0;}
            else if (k==nvar-1) {x1=q[k]-q[k-1]; x2=-q[k];}
            else {x1=q[k]-q[k-1]; x2=0.0;}
            e[k] = 0.5*p[k]*p[k]+0.5*(x1*x1+x2*x2)+1.0/(n+1)*b*(powf(x1,n+1)+powf(x2,n+1));
            __syncthreads();
            int i = blockDim.x/2;
            while (i!=0) {
                if (k < i) e[k] += e[k+i];
                __syncthreads();
                i /= 2;
            }
        }"""
            
        f4="""__global__ void verletMultiSteps(int nsteps, float h, int nvar, float b, float n, float *q0, float *p0, float *en) {
            __shared__ float q[%d];
            __shared__ float p[%d];
            __shared__ float dp[%d];
            __shared__ float e[%d];
            int i;
            int k = threadIdx.x;
            q[k] = q0[k];
            p[k] = p0[k];
            __syncthreads();
            for (i=0; i<nsteps; i++) verletStep(k,h,nvar,b,n,q,p,dp);
            q0[k] = q[k];
            p0[k] = p[k];
            energy(nvar,b,n,q,p,e);
            if (k==0) en[0] = e[0];
        }"""%(nvar,nvar,nvar,nvar)
              
        f5="""__global__ void modeEnergy(float *q, float *p, float *A, float *omeg2, float *emode) {
            __shared__ float xi[%d];
            __shared__ float dxi[%d];
            int i = blockIdx.x; // mode
            int j = threadIdx.x; // variable
            int index = j+blockDim.x*i;
            xi[j] = A[index]*q[j];
            dxi[j] = A[index]*p[j];
            __syncthreads();
            int j1 = blockDim.x/2;
            while (j1!=0) {
                if (j < j1) {
                    xi[j] += xi[j+j1];
                    dxi[j] += dxi[j+j1];
                }
                __syncthreads();
                j1 /= 2;
            }
            if (j==0) emode[i] = 0.5*(dxi[0]*dxi[0]+omeg2[i]*xi[0]*xi[0]);
        }"""%(nvar,nvar)     
              
        mod=SourceModule(f1+"\n"+f2+"\n"+f3+"\n"+f4+"\n"+f5)
        self.verletMultiSteps = mod.get_function("verletMultiSteps")
        self.modeEnergy = mod.get_function("modeEnergy")
              
        self.A = numpy.empty(shape=(nvar,nvar),dtype=numpy.float32)
        s=math.sqrt(2.0/(nvar+1))
        r=math.pi/(nvar+1)
        for i in range(nvar):
            for j in range(nvar):
                self.A[i,j]=s*math.sin(r*(i+1)*(j+1))
        self.omeg2=numpy.empty(shape=nvar,dtype=numpy.float32)
        r=math.pi/(2*(nvar+1))
        for i in range(nvar):
            self.omeg2[i] = 4*math.pow(math.sin((i+1)*r),2)
              
        self.q0 = numpy.empty(shape=nvar,dtype=numpy.float32)
        self.p0 = numpy.empty(shape=nvar,dtype=numpy.float32)
              
    def init(self,m):
        for k in range(self.nvar):
            self.q0[k] = math.sin((k+1)*m*math.pi/(self.nvar+1))
            self.p0[k] = 0.0
              
    def run(self,h,n1,n2,b,n,nm):
        if nm > self.nvar:
            nm = self.nvar
        self.t = numpy.arange(n2)*h*n1
        self.qmat=numpy.empty(shape=(n2,self.nvar),dtype=numpy.float32)
        self.pmat=numpy.empty(shape=(n2,self.nvar),dtype=numpy.float32)
        en=numpy.empty(shape=(1),dtype=numpy.float32)
        self.emat=numpy.empty(shape=(n2,1),dtype=numpy.float32)
        self.emodemat=numpy.empty(shape=(n2,self.nvar),dtype=numpy.float32)
        #memoire GPU
        dev_q0 = cuda.mem_alloc(self.q0.nbytes)
        dev_p0 = cuda.mem_alloc(self.p0.nbytes)
        dev_en = cuda.mem_alloc(en.nbytes)
        dev_A = cuda.mem_alloc(self.A.nbytes)
        dev_omeg2 = cuda.mem_alloc(self.omeg2.nbytes)
        dev_emode = cuda.mem_alloc(self.q0.nbytes)
        cuda.memcpy_htod(dev_q0,self.q0)
        cuda.memcpy_htod(dev_p0,self.p0)
        cuda.memcpy_htod(dev_A,self.A)
        cuda.memcpy_htod(dev_omeg2,self.omeg2)
        ta = time.time()
        # boucle de calcul
        for i2 in range(n2): 
            self.verletMultiSteps(numpy.uint32(n1),numpy.float32(h),numpy.uint32(self.nvar),numpy.float32(b),numpy.float32(n),dev_q0,dev_p0,dev_en,grid=(1,1),block=(self.nvar,1,1))
            self.modeEnergy(dev_q0,dev_p0,dev_A,dev_omeg2,dev_emode,grid=(nm,1),block=(self.nvar,1,1))
            cuda.memcpy_dtoh(self.qmat[i2],dev_q0)
            cuda.memcpy_dtoh(self.pmat[i2],dev_p0)
            cuda.memcpy_dtoh(self.emat[i2],dev_en)
            cuda.memcpy_dtoh(self.emodemat[i2],dev_emode)
            print "%d sur %d"%(i2,n2)
        cuda.memcpy_dtoh(self.q0,dev_q0)
        cuda.memcpy_dtoh(self.p0,dev_p0)
        print "Duree du calcul (s) : %f"%(time.time()-ta)
        self.deltaE = numpy.std(self.emat)/numpy.average(self.emat)
        print "Ecart type relatif de E : %e"%self.deltaE