# coding: utf-8

# # Practice: Generating a Simple Kernel
# 
# The purpose of this practice problem is to generate a simple kernel that applies a user-supplied expression to every entry of an array. Implement a class `ExpressionKernel` that can be used as shown in the test at the end of this notebook.

# In[2]:

import numpy as np
import numpy.linalg as la

import pyopencl as cl
import pyopencl.array
import pyopencl.clmath
import pyopencl.clrandom


# In[3]:


class ExpressionKernel:
    def __init__(self, cl_context, expression):
        # ...
        pass
    
    def __call__(self, queue, ary):
        # ...
        pass


# In[4]:

# Solution

class ExpressionKernel:
    def __init__(self, cl_context, expression):
        src = """
            kernel void apply(__global double *out, global double *in)
            {
                int i = get_global_id(0);
                double x = in[i];
                out[i] = RESULT;
            }
            """

        from pymbolic.mapper.c_code import CCodeMapper
        ccm = CCodeMapper()
        src = src.replace("RESULT", ccm(expression))
        self.prg = cl.Program(cl_context, src).build()
        self.knl = self.prg.apply

    def __call__(self, queue, ary):
        result = cl.array.empty_like(ary)
        self.knl(queue, ary.shape, None, result.data, ary.data)
        return result


# To test our implementation, we create a context and an array full of random numbers:

# In[5]:

cl_context = cl.create_some_context()
queue = cl.CommandQueue(cl_context)

ary = cl.clrandom.rand(queue, 500, dtype=np.float64)


# In[6]:


from pymbolic import var

x = var("x")
eknl = ExpressionKernel(cl_context, var("sqrt")(1-x**2))

result = eknl(queue, ary)

diff = result - cl.clmath.sqrt(1-ary**2)
print(la.norm(diff.get()))


# In[ ]: