// File:      pngfilm.cpp
// Authors:   Philip Rideout (png stuff)
//            Mark Colbert (imaging stuff)
//            Matt Pharr and Greg Humphreys (pbrt itself)
// Brief:     pbrt plugin to write png files.
// ToBeDone:  Param for specifying bit depth of 8 (default) or 16 bits
//            Param for specifying iCCP or sRGB or neither (default)
//            Params for keywords (text annotations)
//            Documentation for the plug-in

#include "pbrt.h"
#include "film.h"
#include "color.h"
#include "paramset.h"
#include "tonemap.h"
#include "sampling.h"
#include <png.h>

extern "C" DLLEXPORT Film* CreateFilm(const ParamSet& params, Filter* filter);

struct Pixel
{
    Pixel() : L(0.f)
    {
        alpha = 0.f;
        weightSum = 0.f;
    }
    Spectrum L;
    float alpha, weightSum;
};


class PngFilm : public Film
{
  public:

    PngFilm::PngFilm(int xres, int yres, Filter *filt, const float crop[4], const string &filename, int wf);
    ~PngFilm();

    void AddSample(const Sample &sample, const Ray &ray, const Spectrum &L, float alpha);
    void GetSampleExtent(int *xstart, int *xend, int *ystart, int *yend) const;
    void WriteImage();

  private:

    static const int filterTableSize = 16;
    Filter *filter;
    int writeFrequency, sampleCount;
    string filename;
    float cropWindow[4];
    int xPixelStart, yPixelStart, xPixelCount, yPixelCount;
    string toneMapper;
    float bloomWidth, bloomRadius, screenGamma, fileGamma, dither;
    BlockedArray<Pixel> *pixels;
    float *filterTable;

    friend extern DLLEXPORT Film* CreateFilm(const ParamSet& params, Filter* filter);
};


PngFilm::~PngFilm()
{
    delete pixels;
    delete filter;
    delete[] filterTable;
}


// PngFilm Method Definitions
PngFilm::PngFilm(int xres, int yres, Filter *filt, const float crop[4], const string &fn, int wf) : Film(xres, yres)
{
    filter = filt;
    memcpy(cropWindow, crop, 4 * sizeof(float));
    filename = fn;
    writeFrequency = sampleCount = wf;

    // Compute film image extent
    xPixelStart = Ceil2Int(xResolution * cropWindow[0]);
    xPixelCount = max(1, Ceil2Int(xResolution * cropWindow[1]) - xPixelStart);
    yPixelStart = Ceil2Int(yResolution * cropWindow[2]);
    yPixelCount = max(1, Ceil2Int(yResolution * cropWindow[3]) - yPixelStart);
    
    // Allocate film image storage
    pixels = new BlockedArray<Pixel>(xPixelCount, yPixelCount);
    
    // Precompute filter weight table
    filterTable = new float[filterTableSize * filterTableSize];
    float *ftp = filterTable;
    for (int y = 0; y < filterTableSize; ++y)
    {
        float fy = ((float)y + .5f) * filter->yWidth / filterTableSize;
        for (int x = 0; x < filterTableSize; ++x)
        {
            float fx = ((float)x + .5f) * filter->xWidth / filterTableSize;
            *ftp++ = filter->Evaluate(fx, fy);
        }
    }
}


void PngFilm::AddSample(const Sample &sample, const Ray &ray, const Spectrum &L, float alpha)
{
    // Compute sample's raster extent
    float dImageX = sample.imageX - 0.5f;
    float dImageY = sample.imageY - 0.5f;
    int x0 = Ceil2Int (dImageX - filter->xWidth);
    int x1 = Floor2Int(dImageX + filter->xWidth);
    int y0 = Ceil2Int (dImageY - filter->yWidth);
    int y1 = Floor2Int(dImageY + filter->yWidth);
    x0 = max(x0, xPixelStart);
    x1 = min(x1, xPixelStart + xPixelCount - 1);
    y0 = max(y0, yPixelStart);
    y1 = min(y1, yPixelStart + yPixelCount - 1);

    // Loop over filter support and add sample to pixel arrays
    // Precompute $x$ and $y$ filter table offsets
    int *ifx = (int*) alloca((x1 - x0 + 1) * sizeof(int));
    for (int x = x0; x <= x1; ++x)
    {
        float fx = fabsf((x - dImageX) * filter->invXWidth * filterTableSize);
        ifx[x-x0] = min(Floor2Int(fx), filterTableSize - 1);
    }

    int *ify = (int*) alloca((y1 - y0 + 1) * sizeof(int));
    for (int y = y0; y <= y1; ++y)
    {
        float fy = fabsf((y - dImageY) * filter->invYWidth * filterTableSize);
        ify[y-y0] = min(Floor2Int(fy), filterTableSize - 1);
    }

    for (int y = y0; y <= y1; ++y)
    {
        for (int x = x0; x <= x1; ++x)
        {
            // Evaluate filter value at $(x,y)$ pixel
            int offset = ify[y - y0] * filterTableSize + ifx[x - x0];
            float filterWt = filterTable[offset];

            // Update pixel values with filtered sample contribution
            Pixel &pixel = (*pixels)(x - xPixelStart, y - yPixelStart);
            pixel.L.AddWeighted(filterWt, L);
            pixel.alpha += alpha * filterWt;
            pixel.weightSum += filterWt;
        }
    }

    // Possibly write out in-progress image
    if (--sampleCount == 0)
    {
        WriteImage();
        sampleCount = writeFrequency;
    }
}


void PngFilm::GetSampleExtent(int *xstart, int *xend, int *ystart, int *yend) const
{
    *xstart = Floor2Int(xPixelStart + .5f - filter->xWidth);
    *xend   = Floor2Int(xPixelStart + .5f + xPixelCount  + filter->xWidth);
    *ystart = Floor2Int(yPixelStart + .5f - filter->yWidth);
    *yend   = Floor2Int(yPixelStart + .5f + yPixelCount + filter->yWidth);
}


void pbrt_png_error(png_structp png_, png_const_charp msg)
{
    Error("libpng error: %s\n", msg);
}


void PngFilm::WriteImage()
{
    // Convert image to RGB and compute final pixel values
    int nPix = xPixelCount * yPixelCount;
    float *rgb = new float[3 * nPix], *alpha = new float[nPix];
    int offset = 0;
    for (int y = 0; y < yPixelCount; ++y)
    {
        for (int x = 0; x < xPixelCount; ++x)
        {
            // Convert pixel spectral radiance to RGB
            float xyz[3];
            (*pixels)(x, y).L.XYZ(xyz);
            const float rWeight[3] = { 3.240479f, -1.537150f, -0.498535f };
            const float gWeight[3] = {-0.969256f,  1.875991f,  0.041556f };
            const float bWeight[3] = { 0.055648f, -0.204043f,  1.057311f };
            rgb[3*offset  ] = rWeight[0]*xyz[0] +
                              rWeight[1]*xyz[1] +
                              rWeight[2]*xyz[2];
            rgb[3*offset+1] = gWeight[0]*xyz[0] +
                              gWeight[1]*xyz[1] +
                              gWeight[2]*xyz[2];
            rgb[3*offset+2] = bWeight[0]*xyz[0] +
                              bWeight[1]*xyz[1] +
                              bWeight[2]*xyz[2];
            alpha[offset] = (*pixels)(x, y).alpha;

            // Normalize pixel with weight sum
            float weightSum = (*pixels)(x, y).weightSum;
            if (weightSum != 0.f)
            {
                float invWt = 1.f / weightSum;
                rgb[3*offset  ] =
                    Clamp(rgb[3*offset  ] * invWt, 0.f, INFINITY);
                rgb[3*offset+1] =
                    Clamp(rgb[3*offset+1] * invWt, 0.f, INFINITY);
                rgb[3*offset+2] =
                    Clamp(rgb[3*offset+2] * invWt, 0.f, INFINITY);
                alpha[offset] = Clamp(alpha[offset] * invWt, 0.f, 1.f);
            }

            ++offset;
        }
    }

    // Apply the tone mapper.
    // Note that no gamma processing is done here.
    // PNG can encode gamma information, so applying gamma would infer a needless loss of information.
    ParamSet toneParams;
    ApplyImagingPipeline(rgb,xPixelCount,yPixelCount,0,bloomRadius,bloomWidth,toneMapper.c_str(),&toneParams,1,dither,255);

    //
    // PNG writing starts here!
    //
    //   xPixelCount, yPixelCount....size of cropped region
    //   xResolution, yResolution....size of entire image (usually the same as the cropped image)
    //   xPixelStart, yPixelStart....offset of cropped region (usually 0,0)
    //

    FILE *fp = fopen(filename.c_str(), "wb");
    png_structp png = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, (png_error_ptr) pbrt_png_error, NULL);
    png_infop info = png_create_info_struct(png);
    png_init_io(png, fp);

    // Registered PNG keywords are:
    // Title           7 Mar 95       PNG-1.2
    // Author          7 Mar 95       PNG-1.2
    // Description     7 Mar 95       PNG-1.2
    // Copyright       7 Mar 95       PNG-1.2
    // Creation Time   7 Mar 95       PNG-1.2
    // Software        7 Mar 95       PNG-1.2
    // Disclaimer      7 Mar 95       PNG-1.2
    // Warning         7 Mar 95       PNG-1.2
    // Source          7 Mar 95       PNG-1.2
    // Comment         7 Mar 95       PNG-1.2

    png_text text;
    text.compression = PNG_TEXT_COMPRESSION_NONE;
    text.key = (png_charp) "Software";
    text.text = (png_charp) "pbrt";
    text.text_length = 4;
    png_set_text(png, info, &text, 1);

    png_color_16 black = {0};
    png_set_background(png, &black, PNG_BACKGROUND_GAMMA_SCREEN, 0, 255.0);

    // gAMA:
    //
    // From the PNG spec: (http://www.w3.org/TR/PNG/)
    //
    // Computer graphics renderers often do not perform gamma encoding, instead
    // making sample values directly proportional to scene light intensity. If
    // the PNG encoder receives sample values that have already been 
    // quantized into integer values, there is no point in doing gamma encoding 
    // on them; that would just result in further loss of information. The 
    // encoder should just write the sample values to the PNG datastream. This
    // does not imply that the gAMA chunk should contain a gamma value of 1.0 
    // because the desired end-to-end transfer function from scene intensity to 
    // display output intensity is not necessarily linear. However, the desired 
    // gamma value is probably not far from 1.0. It may depend on whether the 
    // scene being rendered is a daylight scene or an indoor scene, etc.
    //
    // Given the above text, I decided that the user's requested gamma value
    // should be directly encoded into the PNG gAMA chunk; it is NOT a request
    // for pbrt to do gamma processing, since doing so would infer a needless 
    // loss of information.  The PNG authors say that the gamma value should
    // not necessarily be 1.0, but given the physically rigorous nature of 
    // pbrt, I think it would be rare to use something other than 1.0.
    //
    // Note there is no option for premultiply alpha.
    // It was purposely omitted because the PNG spec states that PNG never uses this.
    //
    // cHRM: is: CIE x,y chromacities of R, G, B and white.
    //
    // x = X / (X + Y + Z)
    // y = Y / (X + Y + Z)

    png_set_gamma(png, screenGamma, fileGamma);

    float rgbWhite[3] = {1, 1, 1};
    float rgbRed[3] = {1, 0, 0};
    float rgbGreen[3] = {0, 1, 0};
    float rgbBlue[3] = {0, 0, 1};
    float xyzWhite[3];
    float xyzRed[3];
    float xyzGreen[3];
    float xyzBlue[3];

    Spectrum(rgbWhite).XYZ(xyzWhite);
    Spectrum(rgbRed).XYZ(xyzRed);
    Spectrum(rgbGreen).XYZ(xyzGreen);
    Spectrum(rgbBlue).XYZ(xyzBlue);

    float whiteX = xyzWhite[0] / (xyzWhite[0] + xyzWhite[1] + xyzWhite[2]);
    float whiteY = xyzWhite[1] / (xyzWhite[0] + xyzWhite[1] + xyzWhite[2]);
    float redX = xyzRed[0] / (xyzRed[0] + xyzRed[1] + xyzRed[2]);
    float redY = xyzRed[1] / (xyzRed[0] + xyzRed[1] + xyzRed[2]);
    float greenX = xyzGreen[0] / (xyzGreen[0] + xyzGreen[1] + xyzGreen[2]);
    float greenY = xyzGreen[1] / (xyzGreen[0] + xyzGreen[1] + xyzGreen[2]);
    float blueX = xyzBlue[0] / (xyzBlue[0] + xyzBlue[1] + xyzBlue[2]);
    float blueY = xyzBlue[1] / (xyzBlue[0] + xyzBlue[1] + xyzBlue[2]);

    png_set_cHRM(png, info, whiteX, whiteY, redX, redY, greenX, greenY, blueX, blueY);

    png_set_IHDR(
        png, info,
        xPixelCount, yPixelCount, 8,
        PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE,
        PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);

    unsigned char** rows = (png_bytep*) malloc(yPixelCount * sizeof(png_bytep));
    rows[0] = (png_bytep) malloc(xPixelCount * yPixelCount * 4);
    for (int i = 1; i < yPixelCount; i++)
        rows[i] = rows[0] + i * xPixelCount * 4;

    for (int x = xPixelStart; x < xPixelStart + xPixelCount; ++x)
    {
        for (int y = yPixelStart; y < yPixelStart + yPixelCount; ++y)
        {
            char r = (char) rgb[(x + y * xResolution) * 3 + 0];
            char g = (char) rgb[(x + y * xResolution) * 3 + 1];
            char b = (char) rgb[(x + y * xResolution) * 3 + 2];
            char a = (char) (255.0 * alpha[x + y * xResolution]);
            rows[y - yPixelStart][x * 4 + 0] = r;
            rows[y - yPixelStart][x * 4 + 1] = g;
            rows[y - yPixelStart][x * 4 + 2] = b;
            rows[y - yPixelStart][x * 4 + 3] = a;
        }
    }

    png_set_rows(png, info, rows);
    png_write_png(png, info, PNG_TRANSFORM_IDENTITY, NULL);

    fclose(fp);

    // Release temporary image memory
    delete[] alpha;
    delete[] rgb;
}


extern "C" DLLEXPORT Film* CreateFilm(const ParamSet& params, Filter* filter)
{
    string filename = params.FindOneString("filename", "pbrt.png");

    int xres = params.FindOneInt("xresolution", 640);
    int yres = params.FindOneInt("yresolution", 480);
    float crop[4] = { 0, 1, 0, 1 };
    int cwi;
    const float *cr = params.FindFloat("cropwindow", &cwi);
    if (cr && cwi == 4)
    {
        crop[0] = Clamp(min(cr[0], cr[1]), 0., 1.);
        crop[1] = Clamp(max(cr[0], cr[1]), 0., 1.);
        crop[2] = Clamp(min(cr[2], cr[3]), 0., 1.);
        crop[3] = Clamp(max(cr[2], cr[3]), 0., 1.);
    }
    int writeFrequency = params.FindOneInt("writefrequency", -1);

    PngFilm* film = new PngFilm(xres, yres, filter, crop, filename, writeFrequency);

    film->toneMapper = params.FindOneString("tonemapper", "maxwhite");
    film->bloomWidth = params.FindOneFloat("bloomwidth", 0.0f);
    film->bloomRadius = params.FindOneFloat("bloomradius", 0.0f);
    film->fileGamma = params.FindOneFloat("filegamma", 1.0f);
    film->screenGamma = params.FindOneFloat("screengamma", 2.2f);
    film->dither = params.FindOneFloat("dither", 0.0f);

    return film;
}