/**
 * SECTION:ufo-averager-task
 * @Short_description: Write TIFF files
 * @Title: averager
 *
 * The averager node writes each incoming image as a TIFF using libtiff to disk.
 * Each file is prefixed with #UfoLaminoBpTask:prefix and written into
 * #UfoLaminoBpTask:path.
 */

#ifdef __APPLE__
#include <OpenCL/cl.h>
#else
#include <CL/cl.h>
#endif
#include <math.h>
#include "ufo-lamino-bp-task.h"
#include "lamino-filter-def.h"

struct _UfoLaminoBpTaskPrivate {
    cl_context      context;
    cl_kernel       bp_kernel;
    cl_kernel       bp_var_kernel;
    cl_kernel       clean_kernel;
    cl_kernel       norm_kernel;
    cl_mem          param_mem;
    gint            proj_idx;
    CLParameters    params;
    gboolean        cleaned;
    gboolean        produced;
};

static void ufo_task_interface_init (UfoTaskIface *iface);

G_DEFINE_TYPE_WITH_CODE (UfoLaminoBpTask, ufo_lamino_bp_task, UFO_TYPE_TASK_NODE,
                         G_IMPLEMENT_INTERFACE (UFO_TYPE_TASK,
                                                ufo_task_interface_init))

#define UFO_LAMINO_BP_TASK_GET_PRIVATE(obj) (G_TYPE_INSTANCE_GET_PRIVATE((obj), UFO_TYPE_LAMINO_BP_TASK, UfoLaminoBpTaskPrivate))

enum {
    PROP_0 = 0,
    PROP_THETA,
    PROP_PSI,
    PROP_ANGLE_STEP,
    PROP_VOL_SX,
    PROP_VOL_SY,
    PROP_VOL_SZ,
    PROP_VOL_OX,
    PROP_VOL_OY,
    PROP_VOL_OZ,
    PROP_PROJ_OX,
    PROP_PROJ_OY,
    PROP_PROJ_OX_VARIATION,
    PROP_Z_SPACING,
    N_PROPERTIES
};

static GParamSpec *properties[N_PROPERTIES] = { NULL, };


UfoNode *
ufo_lamino_bp_task_new (void)
{
    return UFO_NODE (g_object_new (UFO_TYPE_LAMINO_BP_TASK, NULL));
}

static void
ufo_lamino_bp_task_setup (UfoTask *task,
                          UfoResources *resources,
                          GError **error)
{
    UfoLaminoBpTaskPrivate *priv;

    priv = UFO_LAMINO_BP_TASK_GET_PRIVATE (task);
    priv->proj_idx = 0;
    priv->context = ufo_resources_get_context (resources);
    priv->bp_kernel = ufo_resources_get_kernel (resources, "lamino_bp_generic.cl", "lamino_bp_generic", error);
    priv->bp_var_kernel = ufo_resources_get_kernel (resources, "lamino_bp_generic.cl", "lamino_bp_varied", error);
    priv->norm_kernel = ufo_resources_get_kernel (resources, "lamino_bp_generic.cl", "lamino_norm_vol", error);
    priv->clean_kernel = ufo_resources_get_kernel (resources, "lamino_bp_generic.cl", "lamino_clean_vol", error);

    UFO_RESOURCES_CHECK_CLERR (clRetainContext (priv->context));

    if (priv->bp_kernel) {
        UFO_RESOURCES_CHECK_CLERR (clRetainKernel (priv->bp_kernel));
    }

    if (priv->bp_var_kernel) {
        UFO_RESOURCES_CHECK_CLERR (clRetainKernel (priv->bp_var_kernel));
    }

    if (priv->norm_kernel) {
        UFO_RESOURCES_CHECK_CLERR (clRetainKernel (priv->norm_kernel));
    }

    if (priv->clean_kernel) {
        UFO_RESOURCES_CHECK_CLERR (clRetainKernel (priv->clean_kernel));
    }
}

static void
ufo_lamino_bp_task_get_requisition (UfoTask *task,
                                    UfoBuffer **inputs,
                                    UfoRequisition *requisition)
{
    UfoLaminoBpTaskPrivate *priv;
    UfoRequisition in_req;

    priv = UFO_LAMINO_BP_TASK_GET_PRIVATE (task);
    ufo_buffer_get_requisition (inputs[0], &in_req);

    priv->params.proj_sx = in_req.dims[0];
    priv->params.proj_sy = in_req.dims[1];

    if (priv->param_mem == NULL) {
        priv->param_mem = clCreateBuffer (priv->context,
                                          CL_MEM_READ_ONLY, sizeof (CLParameters),
                                          NULL, NULL);
    }

    requisition->n_dims = 3;
    requisition->dims[0] = priv->params.vol_sx;
    requisition->dims[1] = priv->params.vol_sy;
    requisition->dims[2] = priv->params.vol_sz;
}

static guint
ufo_lamino_bp_task_get_num_inputs (UfoTask *task)
{
    return 1;
}

static guint
ufo_lamino_bp_task_get_num_dimensions (UfoTask *task,
                                           guint input)
{
    g_return_val_if_fail (input == 0, 0);
    return 2;
}

static UfoTaskMode
ufo_lamino_bp_task_get_mode (UfoTask *task)
{
    return UFO_TASK_MODE_REDUCTOR | UFO_TASK_MODE_GPU;
}

static void
backproject_regularly (UfoLaminoBpTaskPrivate *priv,
                       UfoRequisition *requisition,
                       cl_command_queue cmd_queue,
                       cl_mem projection_mem,
                       cl_mem volume_mem)
{
    cl_event event;
    gfloat cf, ct, cg;
    gfloat sf, st, sg;

    priv->params.alpha = - 3 * G_PI/2 + priv->params.theta;
    priv->params.phi   = priv->params.angle_step * ((float) priv->proj_idx);

    cf = cos(priv->params.phi);
    ct = cos(priv->params.alpha);
    cg = cos(priv->params.psi);
    sf = sin(priv->params.phi);
    st = sin(priv->params.alpha);
    sg = sin(priv->params.psi);

    priv->params.mat_0 =  cg * cf - sg * st * sf;
    priv->params.mat_1 = -cg * sf - sg * st * cf;
    priv->params.mat_2 = -sg * ct;
    priv->params.mat_3 =  sg * cf + cg * st * sf;
    priv->params.mat_4 = -sg * sf + cg * st * cf;
    priv->params.mat_5 =  cg * ct;

    UFO_RESOURCES_CHECK_CLERR (clEnqueueWriteBuffer (cmd_queue,
                                                     priv->param_mem, CL_FALSE,
                                                     0, sizeof(CLParameters), &priv->params,
                                                     0, NULL, &event));

    UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->bp_kernel, 0, sizeof(cl_mem), (void *) &projection_mem));
    UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->bp_kernel, 1, sizeof(cl_mem), (void *) &volume_mem));
    UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->bp_kernel, 2, sizeof(cl_mem), (void *) &priv->param_mem));

    UFO_RESOURCES_CHECK_CLERR (clEnqueueNDRangeKernel (cmd_queue, priv->bp_kernel,
                                                       3, NULL, requisition->dims, NULL,
                                                       1, &event, NULL));

    UFO_RESOURCES_CHECK_CLERR (clReleaseEvent (event));
}

static void
backproject_varied (UfoLaminoBpTaskPrivate *priv,
                    UfoRequisition *requisition,
                    cl_command_queue cmd_queue,
                    cl_mem projection_mem,
                    cl_mem volume_mem)
{
    cl_event event;

    priv->params.phi   = priv->params.angle_step * ((float) priv->proj_idx);

    UFO_RESOURCES_CHECK_CLERR (clEnqueueWriteBuffer (cmd_queue,
                                                     priv->param_mem, CL_FALSE,
                                                     0, sizeof(CLParameters), &priv->params,
                                                     0, NULL, &event));

    UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->bp_var_kernel, 0, sizeof(cl_mem), (void *) &projection_mem));
    UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->bp_var_kernel, 1, sizeof(cl_mem), (void *) &volume_mem));
    UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->bp_var_kernel, 2, sizeof(cl_mem), (void *) &priv->param_mem));

    UFO_RESOURCES_CHECK_CLERR (clEnqueueNDRangeKernel (cmd_queue, priv->bp_var_kernel,
                                                       3, NULL, requisition->dims, NULL,
                                                       1, &event, NULL));

    UFO_RESOURCES_CHECK_CLERR (clReleaseEvent (event));
}

static gboolean
ufo_lamino_bp_task_process (UfoTask *task,
                            UfoBuffer **inputs,
                            UfoBuffer *output,
                            UfoRequisition *requisition)
{
    UfoLaminoBpTaskPrivate *priv;
    UfoGpuNode *node;
    cl_command_queue cmd_queue;
    cl_mem in_mem;
    cl_mem out_mem;

    priv = UFO_LAMINO_BP_TASK_GET_PRIVATE (task);
    node = UFO_GPU_NODE (ufo_task_node_get_proc_node (UFO_TASK_NODE (task)));
    cmd_queue = ufo_gpu_node_get_cmd_queue (node);

    in_mem = ufo_buffer_get_device_array (inputs[0], cmd_queue);
    out_mem = ufo_buffer_get_device_array (output, cmd_queue);

    if (!priv->cleaned) {
        cl_event clean_event;
        UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->clean_kernel, 0, sizeof(cl_mem), (void *) &out_mem));
        UFO_RESOURCES_CHECK_CLERR (clEnqueueNDRangeKernel (cmd_queue, priv->clean_kernel,
                                                           3, NULL, requisition->dims, NULL,
                                                           0, NULL, &clean_event));
        UFO_RESOURCES_CHECK_CLERR (clWaitForEvents (1, &clean_event));
        UFO_RESOURCES_CHECK_CLERR (clReleaseEvent (clean_event));
        priv->cleaned = TRUE;
    }

    if (priv->params.proj_ox_variation == 0.0f ||
        priv->params.vol_sx < 2 ||
        priv->params.vol_sy < 2 ||
        priv->params.vol_sz < 2 ||
        priv->params.z_spacing > 1.0f) {
        if (priv->proj_idx == 0)
            g_print ("INFO: Using regular backprojection");

        backproject_regularly (priv, requisition, cmd_queue, in_mem, out_mem);
    }
    else {
        if (priv->proj_idx == 0)
            g_print ("INFO: Using backprojection with parameter variation");

        backproject_varied (priv, requisition, cmd_queue, in_mem, out_mem);
    }

    priv->proj_idx++;
    return TRUE;
}

static gboolean
ufo_lamino_bp_task_generate (UfoTask *task,
                             UfoBuffer *output,
                             UfoRequisition *requisition)
{
    UfoLaminoBpTaskPrivate *priv;
    UfoGpuNode *node;
    cl_command_queue cmd_queue;
    cl_mem out_mem;

    priv = UFO_LAMINO_BP_TASK_GET_PRIVATE (task);

    if (priv->produced)
        return FALSE;

    node = UFO_GPU_NODE (ufo_task_node_get_proc_node (UFO_TASK_NODE (task)));
    cmd_queue = ufo_gpu_node_get_cmd_queue (node);
    out_mem = ufo_buffer_get_device_array (output, cmd_queue);

    UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->norm_kernel, 0, sizeof(cl_mem), (void *) &out_mem));
    UFO_RESOURCES_CHECK_CLERR (clSetKernelArg (priv->norm_kernel, 1, sizeof(float), &priv->params.angle_step));

    UFO_RESOURCES_CHECK_CLERR (clEnqueueNDRangeKernel (cmd_queue, priv->norm_kernel,
                                                       3, NULL, requisition->dims, NULL,
                                                       0, NULL, NULL));
    priv->produced = TRUE;
    return TRUE;
}

static UfoNode *
ufo_lamino_bp_task_copy (UfoNode *node,
                         GError **error)
{
    UfoLaminoBpTask *orig;
    UfoLaminoBpTask *copy;

    orig = UFO_LAMINO_BP_TASK (node);
    copy = UFO_LAMINO_BP_TASK (ufo_lamino_bp_task_new ());

    copy->priv->params.theta = orig->priv->params.theta;
    copy->priv->params.psi = orig->priv->params.psi;
    copy->priv->params.angle_step = orig->priv->params.angle_step;
    copy->priv->params.vol_sx = orig->priv->params.vol_sx;
    copy->priv->params.vol_sy = orig->priv->params.vol_sy;
    copy->priv->params.vol_sz = orig->priv->params.vol_sz;
    copy->priv->params.vol_ox = orig->priv->params.vol_ox;
    copy->priv->params.vol_oy = orig->priv->params.vol_oy;
    copy->priv->params.vol_oz = orig->priv->params.vol_oz;
    copy->priv->params.proj_ox = orig->priv->params.proj_ox;
    copy->priv->params.proj_oy = orig->priv->params.proj_oy;
    copy->priv->params.z_spacing = orig->priv->params.z_spacing;

    return UFO_NODE (copy);
}

static void
ufo_lamino_bp_task_set_property (GObject *object,
                                 guint property_id,
                                 const GValue *value,
                                 GParamSpec *pspec)
{
    UfoLaminoBpTaskPrivate *priv = UFO_LAMINO_BP_TASK_GET_PRIVATE (object);

    switch (property_id) {
        case PROP_THETA:
            priv->params.theta = (float) g_value_get_double(value);
            break;
        case PROP_PSI:
            priv->params.psi = (float) g_value_get_double(value);
            break;
        case PROP_ANGLE_STEP:
            priv->params.angle_step = (float) g_value_get_double(value);
            break;
        case PROP_VOL_SX:
            priv->params.vol_sx = g_value_get_uint(value);
            break;
        case PROP_VOL_SY:
            priv->params.vol_sy = g_value_get_uint(value);
            break;
        case PROP_VOL_SZ:
            priv->params.vol_sz = g_value_get_uint(value);
            break;
        case PROP_VOL_OX:
            priv->params.vol_ox = (float)g_value_get_double(value);
            break;
        case PROP_VOL_OY:
            priv->params.vol_oy = (float)g_value_get_double(value);
            break;
        case PROP_VOL_OZ:
            priv->params.vol_oz = (float)g_value_get_double(value);
            break;
        case PROP_PROJ_OX:
            priv->params.proj_ox = (float)g_value_get_double(value);
            break;
        case PROP_PROJ_OY:
            priv->params.proj_oy = (float)g_value_get_double(value);
            break;
        case PROP_PROJ_OX_VARIATION:
            priv->params.proj_ox_variation = (float) g_value_get_double (value);
            break;
        case PROP_Z_SPACING:
            priv->params.z_spacing = (float) g_value_get_double (value);
            break;
        default:
            G_OBJECT_WARN_INVALID_PROPERTY_ID (object, property_id, pspec);
            break;
    }
}

static void
ufo_lamino_bp_task_get_property (GObject *object,
                                 guint property_id,
                                 GValue *value,
                                 GParamSpec *pspec)
{
    UfoLaminoBpTaskPrivate *priv = UFO_LAMINO_BP_TASK_GET_PRIVATE (object);

    switch (property_id) {
        case PROP_THETA:
            g_value_set_double(value, (double) priv->params.theta);
            break;
        case PROP_PSI:
            g_value_set_double(value, (double) priv->params.psi);
            break;
        case PROP_ANGLE_STEP:
            g_value_set_double(value, (double) priv->params.angle_step);
            break;
        case PROP_VOL_SX:
            g_value_set_uint(value, priv->params.vol_sx);
            break;
        case PROP_VOL_SY:
            g_value_set_uint(value, priv->params.vol_sy);
            break;
        case PROP_VOL_SZ:
            g_value_set_uint(value, priv->params.vol_sz);
            break;
        case PROP_VOL_OX:
            g_value_set_double(value, (double)priv->params.vol_ox);
            break;
        case PROP_VOL_OY:
            g_value_set_double(value, (double)priv->params.vol_oy);
            break;
        case PROP_VOL_OZ:
            g_value_set_double(value, (double)priv->params.vol_oz);
            break;
        case PROP_PROJ_OX:
            g_value_set_double(value, (double)priv->params.proj_ox);
            break;
        case PROP_PROJ_OY:
            g_value_set_double(value, (double)priv->params.proj_oy);
            break;
        case PROP_PROJ_OX_VARIATION:
            g_value_set_double (value, (double) priv->params.proj_ox_variation);
            break;
        case PROP_Z_SPACING:
            g_value_set_double (value, (double) priv->params.z_spacing);
            break;
        default:
            G_OBJECT_WARN_INVALID_PROPERTY_ID(object, property_id, pspec);
            break;
    }
}

static void
ufo_lamino_bp_task_finalize (GObject *object)
{
    UfoLaminoBpTaskPrivate *priv = UFO_LAMINO_BP_TASK_GET_PRIVATE (object);

    if (priv->param_mem) {
        UFO_RESOURCES_CHECK_CLERR (clReleaseMemObject (priv->param_mem));
        priv->param_mem = NULL;
    }

    if (priv->bp_kernel) {
        UFO_RESOURCES_CHECK_CLERR (clReleaseKernel (priv->bp_kernel));
        priv->bp_kernel = NULL;
    }

    if (priv->bp_var_kernel) {
        UFO_RESOURCES_CHECK_CLERR (clReleaseKernel (priv->bp_var_kernel));
        priv->bp_var_kernel = NULL;
    }

    if (priv->clean_kernel) {
        UFO_RESOURCES_CHECK_CLERR (clReleaseKernel (priv->clean_kernel));
        priv->clean_kernel = NULL;
    }

    if (priv->norm_kernel) {
        UFO_RESOURCES_CHECK_CLERR (clReleaseKernel (priv->norm_kernel));
        priv->norm_kernel = NULL;
    }

    if (priv->context) {
        UFO_RESOURCES_CHECK_CLERR (clReleaseContext (priv->context));
    }

    G_OBJECT_CLASS (ufo_lamino_bp_task_parent_class)->finalize (object);
}

static void
ufo_task_interface_init (UfoTaskIface *iface)
{
    iface->setup = ufo_lamino_bp_task_setup;
    iface->get_num_inputs = ufo_lamino_bp_task_get_num_inputs;
    iface->get_num_dimensions = ufo_lamino_bp_task_get_num_dimensions;
    iface->get_mode = ufo_lamino_bp_task_get_mode;
    iface->get_requisition = ufo_lamino_bp_task_get_requisition;
    iface->process = ufo_lamino_bp_task_process;
    iface->generate = ufo_lamino_bp_task_generate;
}

static void
ufo_lamino_bp_task_class_init (UfoLaminoBpTaskClass *klass)
{
    GObjectClass *oclass;
    UfoNodeClass *node_class;

    oclass = G_OBJECT_CLASS (klass);
    node_class = UFO_NODE_CLASS (klass);

    oclass->set_property = ufo_lamino_bp_task_set_property;
    oclass->get_property = ufo_lamino_bp_task_get_property;
    oclass->finalize = ufo_lamino_bp_task_finalize;

    node_class->copy = ufo_lamino_bp_task_copy;

    properties[PROP_THETA] =
        g_param_spec_double("theta",
                "Laminographic angle in radians",
                "Laminographic angle in radians",
                -4.0 * G_PI, +4.0 * G_PI, 0.0,
                G_PARAM_READWRITE);

    properties[PROP_PSI] =
        g_param_spec_double("psi",
                "Axis misalignment angle in radians",
                "Axis misalignment angle in radians",
                -4.0 * G_PI, +4.0 * G_PI, 0.0,
                G_PARAM_READWRITE);

    properties[PROP_ANGLE_STEP] =
        g_param_spec_double("angle-step",
                "Increment of rotation angle phi in radians",
                "Increment of rotation angle phi in radians",
                -4.0 * G_PI, +4.0 * G_PI, 0.0,
                G_PARAM_READWRITE);

    properties[PROP_VOL_SX] =
        g_param_spec_uint("vol-sx",
                "Size of reconstructed volume along the 0X-axis in voxels",
                "Size of reconstructed volume along the 0X-axis in voxels",
                0, 1024*8, 512,
                G_PARAM_READWRITE);

    properties[PROP_VOL_SY] =
        g_param_spec_uint("vol-sy",
                "Size of reconstructed volume along the 0Y-axis in voxels",
                "Size of reconstructed volume along the 0Y-axis in voxels",
                0, 1024*8, 512,
                G_PARAM_READWRITE);

    properties[PROP_VOL_SZ] =
        g_param_spec_uint("vol-sz",
                "Size of reconstructed volume along the 0Z-axis in voxels",
                "Size of reconstructed volume along the 0Z-axis in voxels",
                0, 1024*8, 512,
                G_PARAM_READWRITE);

    properties[PROP_VOL_OX] =
        g_param_spec_double("vol-ox",
                "Volume origin offset from the center of a reco-box along the OX-axis in voxels",
                "Volume origin offset from the center of a reco-box along the OX-axis in voxels",
                -1024*8, 1024*8, 0,
                G_PARAM_READWRITE);

    properties[PROP_VOL_OY] =
        g_param_spec_double("vol-oy",
                "Volume origin offset from the center of a reco-box along the OY-axis in voxels",
                "Volume origin offset from the center of a reco-box along the OY-axis in voxels",
                -1024*8, 1024*8, 0,
                G_PARAM_READWRITE);

    properties[PROP_VOL_OZ] =
        g_param_spec_double("vol-oz",
                "Volume origin offset from the center of a reco-box along the OZ-axis in voxels",
                "Volume origin offset from the center of a reco-box along the OZ-axis in voxels",
                -1024*8, 1024*8, 0,
                G_PARAM_READWRITE);

    properties[PROP_Z_SPACING] =
        g_param_spec_double("vol-z-spacing",
                "Spacing factor in z-direction",
                "Spacing factor in z-direction",
                -G_MAXDOUBLE, G_MAXDOUBLE, 0.0,
                G_PARAM_READWRITE);

    properties[PROP_PROJ_OX] =
        g_param_spec_double("proj-ox",
                "Projection of the rotation center on the radiograph origin on the OX-axis",
                "Projection of the rotation center on the radiograph origin on the OX-axis",
                -1024*8, 1024*8, 0,
                G_PARAM_READWRITE);

    properties[PROP_PROJ_OY] =
        g_param_spec_double("proj-oy",
                "Projection of the rotation center on the radiograph origin on the OY-axis",
                "Projection of the rotation center on the radiograph origin on the OY-axis",
                -1024*8, 1024*8, 0,
                G_PARAM_READWRITE);

    properties[PROP_PROJ_OX_VARIATION] =
        g_param_spec_double("proj-ox-variation",
                "Vary x-axis between [proj-ox - x-axis-variation, proj-ox + x-axis-variation]",
                "Vary rotation center x-axis between [proj-ox - x-axis-variation, proj-ox + x-axis-variation]",
                0.0, G_MAXDOUBLE, 0.0,
                G_PARAM_READWRITE);

    for (guint i = PROP_0 + 1; i < N_PROPERTIES; i++)
        g_object_class_install_property (oclass, i, properties[i]);

    g_type_class_add_private (G_OBJECT_CLASS (klass), sizeof(UfoLaminoBpTaskPrivate));
}

static void
ufo_lamino_bp_task_init(UfoLaminoBpTask *self)
{
    UfoLaminoBpTaskPrivate *priv;

    self->priv = priv = UFO_LAMINO_BP_TASK_GET_PRIVATE(self);
    priv->param_mem = NULL;
    priv->bp_kernel = NULL;
    priv->bp_var_kernel = NULL;
    priv->clean_kernel = NULL;
    priv->norm_kernel = NULL;
    priv->cleaned = FALSE;
    priv->produced = FALSE;

    priv->params.phi = 0.0f;
    priv->params.psi = 0.0f;
    priv->params.theta = G_PI / 2;
    priv->params.alpha = G_PI / 2;
    priv->params.angle_step = 0.0f;
    priv->params.vol_sx = 512;
    priv->params.vol_sy = 512;
    priv->params.vol_sz = 512;
    priv->params.vol_ox = 0;
    priv->params.vol_oy = 0;
    priv->params.vol_oz = 0;
    priv->params.proj_ox = 0;
    priv->params.proj_oy = 0;
    priv->params.proj_ox_variation = 0.0f;
    priv->params.z_spacing = 1.0f;
}