/*
 * Copyright 2016 Google Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "guetzli/output_image.h"

#include <algorithm>
#include <assert.h>
#include <stdio.h>
#include <string.h>
#include <cmath>
#include <cstdlib>

#include "guetzli/idct.h"
#include "guetzli/color_transform.h"
#include "guetzli/dct_double.h"
#include "guetzli/gamma_correct.h"
#include "guetzli/preprocess_downsample.h"
#include "guetzli/quantize.h"

namespace guetzli {

OutputImageComponent::OutputImageComponent(int w, int h)
    : width_(w), height_(h) {
  Reset(1, 1);
}

void OutputImageComponent::Reset(int factor_x, int factor_y) {
  factor_x_ = factor_x;
  factor_y_ = factor_y;
  width_in_blocks_ = (width_ + 8 * factor_x_ - 1) / (8 * factor_x_);
  height_in_blocks_ = (height_ + 8 * factor_y_ - 1) / (8 * factor_y_);
  num_blocks_ = width_in_blocks_ * height_in_blocks_;
  coeffs_ = std::vector<coeff_t>(num_blocks_ * kDCTBlockSize);
  pixels_ = std::vector<uint16_t>(width_ * height_, 128 << 4);
  for (int i = 0; i < kDCTBlockSize; ++i) quant_[i] = 1;
}

bool OutputImageComponent::IsAllZero() const {
  int numcoeffs = num_blocks_ * kDCTBlockSize;
  for (int i = 0; i < numcoeffs; ++i) {
    if (coeffs_[i] != 0) return false;
  }
  return true;
}

void OutputImageComponent::GetCoeffBlock(int block_x, int block_y,
                                         coeff_t block[kDCTBlockSize]) const {
  assert(block_x < width_in_blocks_);
  assert(block_y < height_in_blocks_);
  int offset = (block_y * width_in_blocks_ + block_x) * kDCTBlockSize;
  memcpy(block, &coeffs_[offset], kDCTBlockSize * sizeof(coeffs_[0]));
}

void OutputImageComponent::ToPixels(int xmin, int ymin, int xsize, int ysize,
                                    uint8_t* out, int stride) const {
  assert(xmin >= 0);
  assert(ymin >= 0);
  assert(xmin < width_);
  assert(ymin < height_);
  const int yend1 = ymin + ysize;
  const int yend0 = std::min(yend1, height_);
  int y = ymin;
  for (; y < yend0; ++y) {
    const int xend1 = xmin + xsize;
    const int xend0 = std::min(xend1, width_);
    int x = xmin;
    int px = y * width_ + xmin;
    for (; x < xend0; ++x, ++px, out += stride) {
      *out = static_cast<uint8_t>((pixels_[px] + 8 - (x & 1)) >> 4);
    }
    const int offset = -stride;
    for (; x < xend1; ++x) {
      *out = out[offset];
      out += stride;
    }
  }
  for (; y < yend1; ++y) {
    const int offset = -stride * xsize;
    for (int x = 0; x < xsize; ++x) {
      *out = out[offset];
      out += stride;
    }
  }
}

void OutputImageComponent::ToFloatPixels(float* out, int stride) const {
  assert(factor_x_ == 1);
  assert(factor_y_ == 1);
  for (int block_y = 0; block_y < height_in_blocks_; ++block_y) {
    for (int block_x = 0; block_x < width_in_blocks_; ++block_x) {
      coeff_t block[kDCTBlockSize];
      GetCoeffBlock(block_x, block_y, block);
      double blockd[kDCTBlockSize];
      for (int k = 0; k < kDCTBlockSize; ++k) {
        blockd[k] = block[k];
      }
      ComputeBlockIDCTDouble(blockd);
      for (int iy = 0; iy < 8; ++iy) {
        for (int ix = 0; ix < 8; ++ix) {
          int y = block_y * 8 + iy;
          int x = block_x * 8 + ix;
          if (y >= height_ || x >= width_) continue;
          out[(y * width_ + x) * stride] = static_cast<float>(blockd[8 * iy + ix] + 128.0);
        }
      }
    }
  }
}

void OutputImageComponent::SetCoeffBlock(int block_x, int block_y,
                                         const coeff_t block[kDCTBlockSize]) {
  assert(block_x < width_in_blocks_);
  assert(block_y < height_in_blocks_);
  int offset = (block_y * width_in_blocks_ + block_x) * kDCTBlockSize;
  memcpy(&coeffs_[offset], block, kDCTBlockSize * sizeof(coeffs_[0]));
  uint8_t idct[kDCTBlockSize];
  ComputeBlockIDCT(&coeffs_[offset], idct);
  UpdatePixelsForBlock(block_x, block_y, idct);
}

void OutputImageComponent::UpdatePixelsForBlock(
    int block_x, int block_y, const uint8_t idct[kDCTBlockSize]) {
  if (factor_x_ == 1 && factor_y_ == 1) {
    for (int iy = 0; iy < 8; ++iy) {
      for (int ix = 0; ix < 8; ++ix) {
        int x = 8 * block_x + ix;
        int y = 8 * block_y + iy;
        if (x >= width_ || y >= height_) continue;
        int p = y * width_ + x;
        pixels_[p] = idct[8 * iy + ix] << 4;
      }
    }
  } else if (factor_x_ == 2 && factor_y_ == 2) {
    // Fill in the 10x10 pixel area in the subsampled image that will be the
    // basis of the upsampling. This area is enough to hold the 3x3 kernel of
    // the fancy upsampler around each pixel.
    static const int kSubsampledEdgeSize = 10;
    uint16_t subsampled[kSubsampledEdgeSize * kSubsampledEdgeSize];
    for (int j = 0; j < kSubsampledEdgeSize; ++j) {
      // The order we fill in the rows is:
      //   8 rows intersecting the block, row below, row above
      const int y0 = block_y * 16 + (j < 9 ? j * 2 : -2);
      for (int i = 0; i < kSubsampledEdgeSize; ++i) {
        // The order we fill in each row is:
        //   8 pixels within the block, left edge, right edge
        const int ix = ((j < 9 ? (j + 1) * kSubsampledEdgeSize : 0) +
                        (i < 9 ? i + 1 : 0));
        const int x0 = block_x * 16 + (i < 9 ? i * 2 : -2);
        if (x0 < 0) {
          subsampled[ix] = subsampled[ix + 1];
        } else if (y0 < 0) {
          subsampled[ix] = subsampled[ix + kSubsampledEdgeSize];
        } else if (x0 >= width_) {
          subsampled[ix] = subsampled[ix - 1];
        } else if (y0 >= height_) {
          subsampled[ix] = subsampled[ix - kSubsampledEdgeSize];
        } else if (i < 8 && j < 8) {
          subsampled[ix] = idct[j * 8 + i] << 4;
        } else {
          // Reconstruct the subsampled pixels around the edge of the current
          // block by computing the inverse of the fancy upsampler.
          const int y1 = std::max(y0 - 1, 0);
          const int x1 = std::max(x0 - 1, 0);
          subsampled[ix] = (pixels_[y0 * width_ + x0] * 9 +
                            pixels_[y1 * width_ + x1] +
                            pixels_[y0 * width_ + x1] * -3 +
                            pixels_[y1 * width_ + x0] * -3) >> 2;
        }
      }
    }

    // Determine area to update.
    int xmin = std::max(block_x * 16 - 1, 0);
    int xmax = std::min(block_x * 16 + 16, width_ - 1);
    int ymin = std::max(block_y * 16 - 1, 0);
    int ymax = std::min(block_y * 16 + 16, height_ - 1);

    // Apply the fancy upsampler on the subsampled block.
    for (int y = ymin; y <= ymax; ++y) {
      const int y0 = ((y & ~1) / 2 - block_y * 8 + 1) * kSubsampledEdgeSize;
      const int dy = ((y & 1) * 2 - 1) * kSubsampledEdgeSize;
      uint16_t* rowptr = &pixels_[y * width_];
      for (int x = xmin; x <= xmax; ++x) {
        const int x0 = (x & ~1) / 2 - block_x * 8 + 1;
        const int dx = (x & 1) * 2 - 1;
        const int ix = x0 + y0;
        rowptr[x] = (subsampled[ix] * 9 + subsampled[ix + dy] * 3 +
                     subsampled[ix + dx] * 3 + subsampled[ix + dx + dy]) >> 4;
      }
    }
  } else {
    printf("Sampling ratio not supported: factor_x = %d factor_y = %d\n",
           factor_x_, factor_y_);
    exit(1);
  }
}

void OutputImageComponent::CopyFromJpegComponent(const JPEGComponent& comp,
                                                 int factor_x, int factor_y,
                                                 const int* quant) {
  Reset(factor_x, factor_y);
  assert(width_in_blocks_ <= comp.width_in_blocks);
  assert(height_in_blocks_ <= comp.height_in_blocks);
  const size_t src_row_size = comp.width_in_blocks * kDCTBlockSize;
  for (int block_y = 0; block_y < height_in_blocks_; ++block_y) {
    const coeff_t* src_coeffs = &comp.coeffs[block_y * src_row_size];
    for (int block_x = 0; block_x < width_in_blocks_; ++block_x) {
      coeff_t block[kDCTBlockSize];
      for (int i = 0; i < kDCTBlockSize; ++i) {
        block[i] = src_coeffs[i] * quant[i];
      }
      SetCoeffBlock(block_x, block_y, block);
      src_coeffs += kDCTBlockSize;
    }
  }
  memcpy(quant_, quant, sizeof(quant_));
}

void OutputImageComponent::ApplyGlobalQuantization(const int q[kDCTBlockSize]) {
  for (int block_y = 0; block_y < height_in_blocks_; ++block_y) {
    for (int block_x = 0; block_x < width_in_blocks_; ++block_x) {
      coeff_t block[kDCTBlockSize];
      GetCoeffBlock(block_x, block_y, block);
      if (QuantizeBlock(block, q)) {
        SetCoeffBlock(block_x, block_y, block);
      }
    }
  }
  memcpy(quant_, q, sizeof(quant_));
}

OutputImage::OutputImage(int w, int h)
    : width_(w),
      height_(h),
      components_(3, OutputImageComponent(w, h)) {}

void OutputImage::CopyFromJpegData(const JPEGData& jpg) {
  for (size_t i = 0; i < jpg.components.size(); ++i) {
    const JPEGComponent& comp = jpg.components[i];
    assert(jpg.max_h_samp_factor % comp.h_samp_factor == 0);
    assert(jpg.max_v_samp_factor % comp.v_samp_factor == 0);
    int factor_x = jpg.max_h_samp_factor / comp.h_samp_factor;
    int factor_y = jpg.max_v_samp_factor / comp.v_samp_factor;
    assert(comp.quant_idx < jpg.quant.size());
    components_[i].CopyFromJpegComponent(comp, factor_x, factor_y,
                                         &jpg.quant[comp.quant_idx].values[0]);
  }
}

namespace {

void SetDownsampledCoefficients(const std::vector<float>& pixels,
                                int factor_x, int factor_y,
                                OutputImageComponent* comp) {
  assert(pixels.size() == comp->width() * comp->height());
  comp->Reset(factor_x, factor_y);
  for (int block_y = 0; block_y < comp->height_in_blocks(); ++block_y) {
    for (int block_x = 0; block_x < comp->width_in_blocks(); ++block_x) {
      double blockd[kDCTBlockSize];
      int x0 = 8 * block_x * factor_x;
      int y0 = 8 * block_y * factor_y;
      assert(x0 < comp->width());
      assert(y0 < comp->height());
      for (int iy = 0; iy < 8; ++iy) {
        for (int ix = 0; ix < 8; ++ix) {
          float avg = 0.0;
          for (int j = 0; j < factor_y; ++j) {
            for (int i = 0; i < factor_x; ++i) {
              int x = std::min(x0 + ix * factor_x + i, comp->width() - 1);
              int y = std::min(y0 + iy * factor_y + j, comp->height() - 1);
              avg += pixels[y * comp->width() + x];
            }
          }
          avg /= factor_x * factor_y;
          blockd[iy * 8 + ix] = avg;
        }
      }
      ComputeBlockDCTDouble(blockd);
      blockd[0] -= 1024.0;
      coeff_t block[kDCTBlockSize];
      for (int k = 0; k < kDCTBlockSize; ++k) {
        block[k] = static_cast<coeff_t>(std::round(blockd[k]));
      }
      comp->SetCoeffBlock(block_x, block_y, block);
    }
  }
}

}  // namespace

void OutputImage::Downsample(const DownsampleConfig& cfg) {
  if (components_[1].IsAllZero() && components_[2].IsAllZero()) {
    // If the image is already grayscale, nothing to do.
    return;
  }
  if (cfg.use_silver_screen &&
      cfg.u_factor_x == 2 && cfg.u_factor_y == 2 &&
      cfg.v_factor_x == 2 && cfg.v_factor_y == 2) {
    std::vector<uint8_t> rgb = ToSRGB();
    std::vector<std::vector<float> > yuv = RGBToYUV420(rgb, width_, height_);
    SetDownsampledCoefficients(yuv[0], 1, 1, &components_[0]);
    SetDownsampledCoefficients(yuv[1], 2, 2, &components_[1]);
    SetDownsampledCoefficients(yuv[2], 2, 2, &components_[2]);
    return;
  }
  // Get the floating-point precision YUV array represented by the set of
  // DCT coefficients.
  std::vector<std::vector<float> > yuv(3, std::vector<float>(width_ * height_));
  for (int c = 0; c < 3; ++c) {
    components_[c].ToFloatPixels(&yuv[c][0], 1);
  }

  yuv = PreProcessChannel(width_, height_, 2, 1.3f, 0.5f,
                          cfg.u_sharpen, cfg.u_blur, yuv);
  yuv = PreProcessChannel(width_, height_, 1, 1.3f, 0.5f,
                          cfg.v_sharpen, cfg.v_blur, yuv);

  // Do the actual downsampling (averaging) and forward-DCT.
  if (cfg.u_factor_x != 1 || cfg.u_factor_y != 1) {
    SetDownsampledCoefficients(yuv[1], cfg.u_factor_x, cfg.u_factor_y,
                               &components_[1]);
  }
  if (cfg.v_factor_x != 1 || cfg.v_factor_y != 1) {
    SetDownsampledCoefficients(yuv[2], cfg.v_factor_x, cfg.v_factor_y,
                               &components_[2]);
  }
}

void OutputImage::ApplyGlobalQuantization(const int q[3][kDCTBlockSize]) {
  for (int c = 0; c < 3; ++c) {
    components_[c].ApplyGlobalQuantization(&q[c][0]);
  }
}

void OutputImage::SaveToJpegData(JPEGData* jpg) const {
  assert(components_[0].factor_x() == 1);
  assert(components_[0].factor_y() == 1);
  jpg->width = width_;
  jpg->height = height_;
  jpg->max_h_samp_factor = 1;
  jpg->max_v_samp_factor = 1;
  jpg->MCU_cols = components_[0].width_in_blocks();
  jpg->MCU_rows = components_[0].height_in_blocks();
  int ncomp = components_[1].IsAllZero() && components_[2].IsAllZero() ? 1 : 3;
  for (int i = 1; i < ncomp; ++i) {
    jpg->max_h_samp_factor = std::max(jpg->max_h_samp_factor,
                                      components_[i].factor_x());
    jpg->max_v_samp_factor = std::max(jpg->max_h_samp_factor,
                                      components_[i].factor_y());
    jpg->MCU_cols = std::min(jpg->MCU_cols, components_[i].width_in_blocks());
    jpg->MCU_rows = std::min(jpg->MCU_rows, components_[i].height_in_blocks());
  }
  jpg->components.resize(ncomp);
  int q[3][kDCTBlockSize];
  for (int c = 0; c < 3; ++c) {
    memcpy(&q[c][0], components_[c].quant(), kDCTBlockSize * sizeof(q[0][0]));
  }
  for (int c = 0; c < ncomp; ++c) {
    JPEGComponent* comp = &jpg->components[c];
    assert(jpg->max_h_samp_factor % components_[c].factor_x() == 0);
    assert(jpg->max_v_samp_factor % components_[c].factor_y() == 0);
    comp->id = c;
    comp->h_samp_factor = jpg->max_h_samp_factor / components_[c].factor_x();
    comp->v_samp_factor = jpg->max_v_samp_factor / components_[c].factor_y();
    comp->width_in_blocks = jpg->MCU_cols * comp->h_samp_factor;
    comp->height_in_blocks = jpg->MCU_rows * comp->v_samp_factor;
    comp->num_blocks = comp->width_in_blocks * comp->height_in_blocks;
    comp->coeffs.resize(kDCTBlockSize * comp->num_blocks);

    int last_dc = 0;
    const coeff_t* src_coeffs = components_[c].coeffs();
    coeff_t* dest_coeffs = &comp->coeffs[0];
    for (int block_y = 0; block_y < comp->height_in_blocks; ++block_y) {
      for (int block_x = 0; block_x < comp->width_in_blocks; ++block_x) {
        if (block_y >= components_[c].height_in_blocks() ||
            block_x >= components_[c].width_in_blocks()) {
          dest_coeffs[0] = last_dc;
          for (int k = 1; k < kDCTBlockSize; ++k) {
            dest_coeffs[k] = 0;
          }
        } else {
          for (int k = 0; k < kDCTBlockSize; ++k) {
            const int quant = q[c][k];
            int coeff = src_coeffs[k];
            assert(coeff % quant == 0);
            dest_coeffs[k] = coeff / quant;
          }
          src_coeffs += kDCTBlockSize;
        }
        last_dc = dest_coeffs[0];
        dest_coeffs += kDCTBlockSize;
      }
    }
  }
  SaveQuantTables(q, jpg);
}

std::vector<uint8_t> OutputImage::ToSRGB(int xmin, int ymin,
                                         int xsize, int ysize) const {
  std::vector<uint8_t> rgb(xsize * ysize * 3);
  for (int c = 0; c < 3; ++c) {
    components_[c].ToPixels(xmin, ymin, xsize, ysize, &rgb[c], 3);
  }
  for (size_t p = 0; p < rgb.size(); p += 3) {
    ColorTransformYCbCrToRGB(&rgb[p]);
  }
  return rgb;
}

std::vector<uint8_t> OutputImage::ToSRGB() const {
  return ToSRGB(0, 0, width_, height_);
}

void OutputImage::ToLinearRGB(int xmin, int ymin, int xsize, int ysize,
                              std::vector<std::vector<float> >* rgb) const {
  const double* lut = Srgb8ToLinearTable();
  std::vector<uint8_t> rgb_pixels = ToSRGB(xmin, ymin, xsize, ysize);
  for (int p = 0; p < xsize * ysize; ++p) {
    for (int i = 0; i < 3; ++i) {
      (*rgb)[i][p] = static_cast<float>(lut[rgb_pixels[3 * p + i]]);
    }
  }
}

void OutputImage::ToLinearRGB(std::vector<std::vector<float> >* rgb) const {
  ToLinearRGB(0, 0, width_, height_, rgb);
}

std::string OutputImage::FrameTypeStr() const {
  char buf[128];
  int len = snprintf(buf, sizeof(buf), "f%d%d%d%d%d%d",
                     component(0).factor_x(), component(0).factor_y(),
                     component(1).factor_x(), component(1).factor_y(),
                     component(2).factor_x(), component(2).factor_y());
  return std::string(buf, len);
}

}  // namespace guetzli
