/*! \file output.c
 *
 * \author Nick Seltzer
 */

/*
 * COPYRIGHT (c) 2011 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 */

#include <stdio.h>
#include "Diderot/diderot.h"
#include <teem/nrrd.h>
#include <teem/air.h>

// FIXME: to make this compatible with OpenCL (and CUDA), we should be copying the
// output state to the output array on the GPU and then just uploading the output
// array to the CPU.  This will require changing the interface of these functions.

/*! \brief Print the strand output state to a text file.
 */
void Diderot_Print (Diderot_World_t *wrld)
{
    WorldPrefix_t *hdr = (WorldPrefix_t *)wrld;
    airArray *mop = airMopNew();

  // here we have the final state of all of the strands in the "in" buffer
    int outFileNameLen = strlen(hdr->name) + 5;
    char *outFileName = (char *)CheckedAlloc(outFileNameLen);
    airMopAdd(mop, outFileName, airFree, airMopAlways);
    snprintf (outFileName, outFileNameLen, "%s.txt", hdr->name);
    FILE *outS = fopen(outFileName, "w");
    if (outS == NULL) {
        fprintf (stderr, "Cannot open output file %s\n", outFileName);
        airMopError(mop);
        exit (8);
    }
    airMopAdd(mop, outS, (airMopper)airFclose, airMopAlways);

  // print the strand output states
    for (int i = 0;  i < hdr->numStrands;  i++) {
        if (hdr->status[i] == DIDEROT_STABLE) {
            Diderot_Strands[0]->print (outS, Diderot_OutState(wrld, i));
        }
    }

    airMopOkay(mop);
}

/*! \brief Output the strand output state to a nrrd file.
 */
void Diderot_Output (Diderot_World_t *wrld, size_t outStateSzb)
{
    WorldPrefix_t *hdr = (WorldPrefix_t *)wrld;
    airArray *mop = airMopNew();
    int status;
    char *err;

  // get correct output file name
    int outFileNameLen = strlen(hdr->name) + 6;
    char *outFileName = (char *)CheckedAlloc(outFileNameLen);
    airMopAdd(mop, outFileName, airFree, airMopAlways);
    snprintf (outFileName, outFileNameLen, "%s.nrrd", hdr->name);

  // calculate the amount of space needed for the output array
    int numStabilized;
    if (hdr->isArray) {
        numStabilized = hdr->numStrands;
    }
    else {
        numStabilized = 0;
        for (int i = 0;  i < hdr->numStrands;  i++) {
            if (hdr->status[i] == DIDEROT_STABLE)
                numStabilized++;
        }
    }

  // allocate the output object
    void *outData = CheckedAlloc(numStabilized * outStateSzb);
    airMopAdd(mop, outData, airFree, airMopAlways);

  // copy data into outData
    void *outPos = outData;
    for (int i = 0;  i < hdr->numStrands;  i++) {
        if (hdr->status[i] == DIDEROT_STABLE) {
            Diderot_Strands[0]->output (outPos, Diderot_OutState(wrld, i));
            outPos += outStateSzb;
        }
    }

  // make a new nrrd
    Nrrd *outNrrd;
    outNrrd = nrrdNew();
    if (outNrrd == NULL) {
        fprintf(stderr, "Trouble allocating nrrd struct:\n%s", err);
        airMopError(mop);
        exit(1);
    }
    airMopAdd(mop, outNrrd, (airMopper)nrrdNix, airMopAlways);

  // wrap the data (outData) with the nrrd (outNrrd)
    if (hdr->isArray) {
        size_t *sizes = NEWVEC(size_t, hdr->nDims + 1);
        airMopAdd(mop, sizes, airFree, airMopAlways);
        switch(Diderot_Strands[0]->nrrdType) {
            case nrrdTypeInt:
            case nrrdTypeFloat:
                sizes[0] = outStateSzb >> 2;
                break;
            case nrrdTypeLLong:
            case nrrdTypeDouble:
                sizes[0] = outStateSzb >> 3;
                break;
            default:
                fprintf(stderr, "Bad nrrdType\n");
                airMopError(mop);
                exit(1);
                break;
        }
        for (int i = 1; i <= hdr->nDims; i++) {
            sizes[i] = hdr->size[hdr->nDims - i];
        }
        status = nrrdWrap_nva(outNrrd, outData, Diderot_Strands[0]->nrrdType, hdr->nDims + 1, sizes);
        if (status) {
            err = biffGetDone(NRRD);
            fprintf(stderr, "Trouble wrapping nrrd struct:\n%s", err);
            airMopError(mop);
            free(err);
            exit(1);
        }
    }
    else {
        size_t *sizes = NEWVEC(size_t, 2);
        airMopAdd(mop, sizes, airFree, airMopAlways);
        switch(Diderot_Strands[0]->nrrdType) {
            case nrrdTypeInt:
            case nrrdTypeFloat:
                sizes[0] = outStateSzb >> 2;
                break;
            case nrrdTypeLLong:
            case nrrdTypeDouble:
                sizes[0] = outStateSzb >> 3;
                break;
            default:
                fprintf(stderr, "Bad nrrdType\n");
                airMopError(mop);
                exit(1);
                break;
        }
        sizes[1] = numStabilized;
        status = nrrdWrap_nva(outNrrd, outData, Diderot_Strands[0]->nrrdType, hdr->nDims, sizes);
        if (status != 0) {
            err = biffGetDone(NRRD);
            fprintf(stderr, "Trouble wrapping nrrd struct:\n%s", err);
            airMopError(mop);
            free(err);
            exit(1);
        }
    }

  // save the data
    status = nrrdSave(outFileName, outNrrd, NULL);
    if(status) {
        err = biffGetDone(NRRD);
        fprintf(stderr, "Trouble saving nrrd struct:\n%s", err);
        airMopError(mop);
        free(err);
        exit(1);
    }

    airMopOkay(mop);

}