#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <ufo/ufo.h>
#include "reco.h"
#include "io.h"

typedef struct {
    UfoPluginManager *pm;
    UfoTaskGraph *graph;
    UfoTaskNode *pad;
    UfoTaskNode *ramp;
    UfoTaskNode *fft1;
    UfoTaskNode *fft2;
    UfoTaskNode *ifft;
    UfoTaskNode *conv;
    UfoTaskNode *reco;
    UfoTaskNode *slice;
    UfoTaskNode *stats;
    UfoTaskNode *writer;

    /* interactive */
    guint current;
    UfoTaskNode *output;
    UfoTaskNode *input;
    Params *params;

    /* stats file */
    FILE *stats_fp;
} LaminoData;


static void
print_dots (gpointer user)
{
    static int n = 0;

    if ((n++ % 10) == 0)
        g_print (".");
}

static void
check (GError *error)
{
    if (error != NULL) {
        if (error->message != NULL)
            g_printerr ("%s", error->message);

        exit (1);
    }
}

static UfoTaskNode *
make_task (UfoPluginManager *pm, const gchar *name)
{
    GError *error = NULL;
    UfoTaskNode *task;

    task = ufo_plugin_manager_get_task (pm, name, &error);
    check (error);

    return task;
}

static guint
next_power_of_two (guint32 x)
{
    --x;
    x |= x >> 1;
    x |= x >> 2;
    x |= x >> 4;
    x |= x >> 8;
    x |= x >> 16;
    return x + 1;
}

gboolean
params_okay (Params *params)
{
    return params->width != 0 &&
           params->height != 0 &&
           params->num_radios != 0 &&
           params->radios != NULL &&
           params->theta < G_MAXDOUBLE &&
           params->px < G_MAXDOUBLE &&
           params->py < G_MAXDOUBLE;
}

static gboolean
with_flat_field_correction (Params *params)
{
    return params->darks != NULL && params->flats != NULL;
}

static void
stats_finished (UfoTaskNode *stats, UfoBuffer *result, LaminoData *user)
{
    gfloat *data;
    gchar buffer[128];

    data = ufo_buffer_get_host_array (result, NULL);
    g_snprintf (buffer, 128, "%.12f\n", data[0]);
    fwrite (buffer, 1, strlen (buffer), user->stats_fp);
}

static LaminoData *
reco_graph_new (Params *params)
{
    LaminoData *data;
    guint xl, xr, yt, yb;
    guint padded_width;
    guint padded_height;
    gdouble theta_rad;
    gdouble angle_step;
    guint fwidth;
    gchar *stats_name;

    data = g_malloc0 (sizeof (LaminoData));

    data->params = params;
    data->pm = ufo_plugin_manager_new ();
    data->graph = UFO_TASK_GRAPH (ufo_task_graph_new ());

    data->pad = make_task (data->pm, "padding-2d");
    data->ramp = make_task (data->pm, "lamino-ramp");
    data->fft1 = make_task (data->pm, "fft");
    data->fft2 = make_task (data->pm, "fft");
    data->ifft = make_task (data->pm, "ifft");
    data->conv = make_task (data->pm, "lamino-conv");
    data->reco = make_task (data->pm, "lamino-bp");
    data->slice = make_task (data->pm, "slice");
    data->stats = make_task (data->pm, "measure");
    data->writer = make_task (data->pm, "write");

    fwidth = ((guint) params->px) * 2;
    padded_width = next_power_of_two ((guint32) params->width + 1);
    padded_height = next_power_of_two ((guint32) params->height + 1);

    xl = xr = (padded_width - params->width) / 2;
    yt = yb = (padded_height - params->height) / 2;

    angle_step = (G_PI * 2.0) / params->num_radios * params->radio_step;
    theta_rad = params->theta / 360. * G_PI * 2;

    info ("Axis     x=%.1f  y=%.1f  variation=%.1f\n",
          params->px, params->py, params->px_variation);

    info ("Lamino   theta=%.3f  tau=%.3f  step=%.5f\n",
          theta_rad, params->tau, angle_step);

    info ("Padding  size=[%d %d]  (xl=%d xr=%d yt=%d yb=%d)\n",
          padded_width, padded_height, xl, xr, yt, yb);

    info ("Volume   origin=[%.1f %.1f %.1f]  size=[%d %d %d]\n",
          params->v_origin[0], params->v_origin[1], params->v_origin[2],
          params->v_size[0], params->v_size[1], params->v_size[2]);

    g_object_set (data->fft1, "dimensions", 2, NULL);
    g_object_set (data->fft2, "dimensions", 2, NULL);
    g_object_set (data->ifft, "dimensions", 2, NULL);

    g_object_set (data->pad,
                  "xl", xl, "xr", xr,
                  "yt", yt, "yb", yb,
                  "mode", "brep",
                  NULL);

    g_object_set (data->ramp,
                  "width", padded_width,
                  "height", padded_height,
                  "theta", theta_rad,
                  "tau", params->tau,
                  "fwidth", fwidth,
                  NULL);

    g_object_set (data->reco,
                  "angle-step", angle_step,
                  "theta", theta_rad,
                  "psi", params->psi,
                  "proj-ox", params->px + xl,
                  "proj-oy", params->py + yt,
                  "proj-ox-variation", params->px_variation,
                  "vol-ox", params->v_size[0] / 2 + params->v_origin[0],
                  "vol-oy", params->v_size[1] / 2 + params->v_origin[1],
                  "vol-oz", params->v_size[2] / 2 + params->v_origin[2],
                  "vol-sx", params->v_size[0],
                  "vol-sy", params->v_size[1],
                  "vol-sz", params->v_size[2],
                  "vol-z-spacing", params->z_spacing,
                  NULL);

    g_object_set (data->stats,
                  "pass-through", TRUE,
                  NULL);

    g_object_set (data->writer,
                  "filename", params->output,
                  "single-file", TRUE,
                  NULL);

    g_signal_connect (data->stats, "result", G_CALLBACK (stats_finished), data);

    ufo_task_graph_connect_nodes (data->graph, data->pad, data->fft1);
    ufo_task_graph_connect_nodes (data->graph, data->ramp, data->fft2);
    ufo_task_graph_connect_nodes_full (data->graph, data->fft1, data->conv, 0);
    ufo_task_graph_connect_nodes_full (data->graph, data->fft2, data->conv, 1);
    ufo_task_graph_connect_nodes (data->graph, data->conv, data->ifft);
    ufo_task_graph_connect_nodes (data->graph, data->ifft, data->reco);
    ufo_task_graph_connect_nodes (data->graph, data->reco, data->slice);
    ufo_task_graph_connect_nodes (data->graph, data->slice, data->stats);
    ufo_task_graph_connect_nodes (data->graph, data->stats, data->writer);

    stats_name = g_strdup_printf ("%s.stats", params->output);
    data->stats_fp = fopen (stats_name, "wb");
    g_free (stats_name);

    return data;
}

static void
reco_graph_free (LaminoData *data)
{
    g_object_unref (data->pm);
    g_object_unref (data->graph);

    g_object_unref (data->pad);
    g_object_unref (data->fft1);
    g_object_unref (data->fft2);
    /* g_object_unref (conv); */
    g_object_unref (data->reco);
    g_object_unref (data->slice);
    g_object_unref (data->stats);
    g_object_unref (data->writer);

    fclose (data->stats_fp);
}

static void
update_reader (UfoTaskNode *reader, Params *params)
{
    g_object_set (reader,
                  "path", params->radios,
                  "number", params->num_radios - 1,
                  "step", params->radio_step,
                  NULL);
}

static void
show_time (UfoBaseScheduler *sched)
{
    gdouble time = 0.0;

    g_object_get (sched, "time", &time, NULL);
    g_print ("\n");
    info ("Finished in %3.3fs\n", time);
}

void
run_simple_reconstruction (Params *params, gchar **argv)
{
    LaminoData *data;
    UfoBaseScheduler *sched = NULL;
    UfoTaskNode *radio_reader = NULL;
    UfoTaskNode *dark_reader = NULL;
    UfoTaskNode *flat_before_reader = NULL;
    UfoTaskNode *flat_before_stack = NULL;
    UfoTaskNode *flat_before_median = NULL;
    UfoTaskNode *flat_after_reader = NULL;
    UfoTaskNode *flat_after_stack = NULL;
    UfoTaskNode *flat_after_median = NULL;
    UfoTaskNode *flat_interpolate = NULL;
    UfoTaskNode *ffc = NULL;
    UfoTaskNode *copy = NULL;
    UfoTaskNode *sum = NULL;
    UfoTaskNode *sum_writer = NULL;
    GError *error = NULL;

    if (!params_okay (params)) {
        err ("Parameters missing.\n");
        return;
    }

    data = reco_graph_new (params);
    g_signal_connect (data->reco, "processed", G_CALLBACK (print_dots), NULL);

    /* Set up input */
    radio_reader = make_task (data->pm, "read");
    update_reader (radio_reader, params);

    if (with_flat_field_correction (params)) {
        flat_before_reader = make_task (data->pm, "read");
        flat_before_stack = make_task (data->pm, "stack");
        flat_before_median = make_task (data->pm, "flatten");
        dark_reader = make_task (data->pm, "read");
        ffc = make_task (data->pm, "flat-field-correction");
        copy = UFO_TASK_NODE (ufo_copy_task_new ());
        sum = make_task (data->pm, "flatten-inplace");
        sum_writer = make_task (data->pm, "writer");

        g_object_set (flat_before_reader, "path", params->flats, NULL);
        g_object_set (flat_before_stack, "number", params->num_flats, NULL);
        g_object_set (dark_reader, "path", params->darks, NULL);

        g_object_set (ffc, "dark-scale", params->dark_scale,
                           "absorption-correction", TRUE,
                           NULL);


        ufo_task_graph_connect_nodes (data->graph, flat_before_reader, flat_before_stack);
        ufo_task_graph_connect_nodes (data->graph, flat_before_stack, flat_before_median);
        ufo_task_graph_connect_nodes_full (data->graph, radio_reader, ffc, 0);
        ufo_task_graph_connect_nodes_full (data->graph, dark_reader, ffc, 1);

        if (params->write_summed) {
            gchar *sum_name = g_strdup_printf ("%s.sum.tif", params->output);
            info ("Write sum image to `%s'\n", sum_name);

            g_object_set (sum_writer, "filename", sum_name, NULL);
            ufo_task_graph_connect_nodes (data->graph, ffc, copy);
            ufo_task_graph_connect_nodes_full (data->graph, copy, sum, 0);
            ufo_task_graph_connect_nodes_full (data->graph, copy, data->pad, 0);
            ufo_task_graph_connect_nodes (data->graph, sum, sum_writer);
            g_free (sum_name);
        }
        else {
            ufo_task_graph_connect_nodes (data->graph, ffc, data->pad);
        }

        if (params->flats_after != NULL) {
            flat_after_reader = make_task (data->pm, "read");
            flat_after_stack = make_task (data->pm, "stack");
            flat_after_median = make_task (data->pm, "flatten");

            flat_interpolate = make_task (data->pm, "interpolate");

            g_object_set (flat_after_reader, "path", params->flats_after, NULL);
            g_object_set (flat_after_stack, "number", params->num_flats, NULL);

            ufo_task_graph_connect_nodes (data->graph, flat_after_reader, flat_after_stack);
            ufo_task_graph_connect_nodes (data->graph, flat_after_stack, flat_after_median);
            ufo_task_graph_connect_nodes_full (data->graph, flat_before_median, flat_interpolate, 0);
            ufo_task_graph_connect_nodes_full (data->graph, flat_after_median, flat_interpolate, 1);
            ufo_task_graph_connect_nodes_full (data->graph, flat_interpolate, ffc, 2);
        }
        else {
            ufo_task_graph_connect_nodes_full (data->graph, flat_before_median, ffc, 2);
        }
    }
    else {
        warn ("> No flat/dark field correction\n");
        ufo_task_graph_connect_nodes (data->graph, radio_reader, data->pad);
    }

    /* Run reconstruction */
    info ("Processing ");
    sched = UFO_BASE_SCHEDULER (ufo_scheduler_new ());
    ufo_base_scheduler_run (sched, data->graph, &error);
    check (error);
    show_time (sched);

    /* Clean up */
    reco_graph_free (data);

    g_object_unref (sched);
    g_object_unref (radio_reader);

    if (with_flat_field_correction (params)) {
        g_object_unref (ffc);
        g_object_unref (flat_before_reader);
        g_object_unref (flat_before_stack);
        g_object_unref (flat_before_median);
        g_object_unref (dark_reader);
        g_object_unref (copy);
        g_object_unref (sum);
        g_object_unref (sum_writer);

        if (params->flats_after) {
            g_object_unref (flat_after_reader);
            g_object_unref (flat_after_stack);
            g_object_unref (flat_after_median);
        }
    }
}

static void
move_data (UfoTaskNode *reader, LaminoData *data)
{
    UfoBuffer *buffer;
    gfloat *mem;
    gsize size;
    static int n = 0;

    /* FIXME: we miss one projection */
    /* We will block here, even though reader sent the "generated" signal */
    if (data->current == 0) {
        data->current++;
        return;
    }

    buffer = ufo_output_task_get_output_buffer (UFO_OUTPUT_TASK (data->output));
    mem = ufo_buffer_get_host_array (buffer, NULL);
    size = ufo_buffer_get_size (buffer);
    memcpy (data->params->cache + (data->current - 1) * size, mem, size);
    ufo_output_task_release_output_buffer (UFO_OUTPUT_TASK (data->output), buffer);
    data->current++;

    if ((n++ % 10) == 0)
        g_print (".");
}

static gpointer
input_thread (LaminoData *data)
{
    UfoBuffer *buffer;
    UfoRequisition req;
    guint current;
    guint num_radios;

    current = 0;
    req.n_dims = 2;
    req.dims[0] = data->params->width;
    req.dims[1] = data->params->height;
    buffer = ufo_buffer_new (&req, NULL);
    num_radios = data->params->num_radios / data->params->radio_step;

    while (current < num_radios) {
        gfloat *mem;
        gsize size;

        mem = ufo_buffer_get_host_array (buffer, NULL);
        size = ufo_buffer_get_size (buffer);
        memcpy (mem, data->params->cache + current * size, size);
        ufo_input_task_release_input_buffer (UFO_INPUT_TASK (data->input), buffer);
        current++;

        buffer = ufo_input_task_get_input_buffer (UFO_INPUT_TASK (data->input));

        if ((current % 10) == 0)
            g_print (".");
    }

    ufo_input_task_stop (UFO_INPUT_TASK (data->input));
    return NULL;
}

void
run_cached_reconstruction (Params *params, gchar **argv)
{
    LaminoData *data;
    GThread *thread;
    UfoBaseScheduler *sched = NULL;
    GError *error = NULL;

    data = reco_graph_new (params);

    if (params->cache == NULL) {
        UfoTaskGraph *read_graph;
        UfoTaskNode *radio_reader;
        UfoBaseScheduler *sched;
        GError *error = NULL;
        UfoTaskNode *dark_reader = NULL;
        UfoTaskNode *flat_before_reader = NULL;
        UfoTaskNode *flat_before_stack = NULL;
        UfoTaskNode *flat_before_median = NULL;
        UfoTaskNode *flat_after_reader = NULL;
        UfoTaskNode *flat_after_stack = NULL;
        UfoTaskNode *flat_after_median = NULL;
        UfoTaskNode *flat_interpolate = NULL;
        UfoTaskNode *ffc = NULL;

        params->cache = g_malloc (((gsize) params->num_radios) * ((gsize) params->width) * ((gsize) params->height) * sizeof (gfloat));

        /* fill the cache and reconstruct simultaneously */
        read_graph = UFO_TASK_GRAPH (ufo_task_graph_new ());

        radio_reader = make_task (data->pm, "read");
        update_reader (radio_reader, params);
        data->output = UFO_TASK_NODE (ufo_output_task_new (2));

        if (with_flat_field_correction (params)) {
            flat_before_reader = make_task (data->pm, "read");
            flat_before_stack = make_task (data->pm, "stack");
            flat_before_median = make_task (data->pm, "flatten");
            dark_reader = make_task (data->pm, "read");
            ffc = make_task (data->pm, "flat-field-correction");

            g_object_set (dark_reader, "path", params->darks, NULL);
            g_object_set (flat_before_reader, "path", params->flats, NULL);
            g_object_set (flat_before_stack, "number", params->num_flats, NULL);

            g_object_set (ffc,
                          "dark-scale", params->dark_scale,
                          "absorption-correction", TRUE,
                          NULL);

            ufo_task_graph_connect_nodes (read_graph, flat_before_reader, flat_before_stack);
            ufo_task_graph_connect_nodes (read_graph, flat_before_stack, flat_before_median);
            ufo_task_graph_connect_nodes_full (read_graph, radio_reader, ffc, 0);
            ufo_task_graph_connect_nodes_full (read_graph, dark_reader, ffc, 1);
            ufo_task_graph_connect_nodes (read_graph, ffc, data->output);

            if (params->flats_after) {
                flat_after_reader = make_task (data->pm, "read");
                flat_after_stack = make_task (data->pm, "stack");
                flat_after_median = make_task (data->pm, "flatten");

                flat_interpolate = make_task (data->pm, "interpolate");

                g_object_set (flat_after_stack, "number", params->num_flats, NULL);
                g_object_set (flat_after_reader, "path", params->flats_after, NULL);

                ufo_task_graph_connect_nodes (read_graph, flat_after_reader, flat_after_stack);
                ufo_task_graph_connect_nodes (read_graph, flat_after_stack, flat_after_median);
                ufo_task_graph_connect_nodes_full (read_graph, flat_before_median, flat_interpolate, 0);
                ufo_task_graph_connect_nodes_full (read_graph, flat_after_median, flat_interpolate, 1);
                ufo_task_graph_connect_nodes_full (read_graph, flat_interpolate, ffc, 2);
            }
            else {
                ufo_task_graph_connect_nodes_full (read_graph, flat_before_median, ffc, 2);
            }
        }
        else {
            warn ("> No flat/dark field correction\n");
            ufo_task_graph_connect_nodes (read_graph, radio_reader, data->output);
        }

        g_signal_connect (radio_reader, "generated", G_CALLBACK (move_data), data);
        data->current = 0;

        sched = UFO_BASE_SCHEDULER (ufo_scheduler_new ());
        info ("Reading ");
        ufo_base_scheduler_run (sched, read_graph, &error);
        check (error);
        g_print ("\n");

        g_object_unref (radio_reader);
        g_object_unref (read_graph);
        g_object_unref (data->output);

        if (with_flat_field_correction (params)) {
            g_object_unref (ffc);
            g_object_unref (flat_before_reader);
            g_object_unref (flat_before_stack);
            g_object_unref (flat_before_median);
            g_object_unref (dark_reader);

            if (params->flats_after) {
                g_object_unref (flat_after_reader);
                g_object_unref (flat_after_stack);
                g_object_unref (flat_after_median);
            }
        }
    }

    data->current = 0;
    data->input = UFO_TASK_NODE (ufo_input_task_new ());
    ufo_task_graph_connect_nodes (data->graph, data->input, data->pad);

    /* Run reconstruction */
    info ("Processing ");
    thread = g_thread_new (NULL, (GThreadFunc) input_thread, data);

    sched = UFO_BASE_SCHEDULER (ufo_scheduler_new ());
    ufo_base_scheduler_run (sched, data->graph, &error);
    check (error);
    show_time (sched);

    g_thread_unref (thread);
    g_object_unref (data->input);
    reco_graph_free (data);
}