/*
 * Copyright (C) 2009  Trustin Heuiseung Lee
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2.1 of
 * the License, or (at your option) any later version.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, write to the Free
 * Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
 * 02110-1301 USA, or see the FSF site: http://www.fsf.org.
 */
package net.gleamynode.img2pdf;

import java.awt.Color;
import java.awt.Graphics;
import java.awt.image.BufferedImage;
import java.awt.image.IndexColorModel;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;

import com.mortennobel.imagescaling.ResampleFilters;
import com.mortennobel.imagescaling.ResampleOp;

public class ImageProducer {

    // Configurable parameters
    private final int maxBlankHeight = 36;
    private final float topMargin = 0.1f;
    private final float bottomMargin = 0.1f;
    private final float leftMargin = 0.075f;
    private final float rightMargin = 0.075f;
    private final float leftCropThreshold = 0.005f;
    private final float rightCropThreshold = 0.005f;
    private final float leftCropAdjustment = 0.002f;
    private final float rightCropAdjustment = 0.0035f;
    private final float lineSeparatorMinWidth = 0.4f;
    private final float lineSeparatorColorTolerance = 128;
    private final float cropImbalanceLimit = 0.3f;
    private final float minSplitHeight = 0.6f;
    private final float splitOverlap = 0.035f;
    private final int whiteThreshold = 236;
    private final int coverPages = 2;
    private final boolean rotate = true;
    private final boolean debug = false;
    // END of configurable parameters

    private static final IndexColorModel COLOR_MODEL;

    private static BufferedImage newImage(int width, int height) {
        return new BufferedImage(width, height, BufferedImage.TYPE_BYTE_GRAY);
    }

    static {
        final byte[] v = new byte[16];
        for (int i = 0; i < v.length; i ++) {
            v[i] = (byte) (i * 17);
        }
        COLOR_MODEL = new IndexColorModel(4, v.length, v, v, v);
    }

    private final BufferedImage buf = newImage(754, 584);
    private final Queue<BufferedImage> output = new LinkedList<BufferedImage>();
    private boolean processingCoverPage = true;
    private boolean finished;
    private int dstY = 0;
    private boolean wasBlank;
    private int offerCount = 0;

    public ImageProducer() {
        Graphics g = buf.getGraphics();
        g.setColor(Color.WHITE);
        g.fillRect(0, 0, buf.getWidth(), buf.getHeight());
        g.dispose();
    }

    public void offer(BufferedImage img) {
        if (finished) {
            throw new IllegalStateException();
        }

        if (processingCoverPage) {
            processingCoverPage = offerCount ++ < coverPages;
        }

        img = grayscale(img);
        img = crop(img);
        img = scale(img, buf.getWidth());

        List<Chunk> chunks = decompose(img);
        chunks = cropLargeBlanks(chunks);
        //chunks = splitLargeChunks(chunks);

        Graphics g = buf.getGraphics();
        try {
            loop:
            for (int i = 0; i < chunks.size(); i ++) {
                Chunk c = chunks.get(i);
                if (c.isBlank()) {
                    // Skip the blank chunks in the page top.
                    if (dstY == 0 && c.isBlank()) {
                        wasBlank = false;
                        continue;
                    }

                    // Skip continuous blanks.
                    if (wasBlank) {
                        continue;
                    }
                }

                wasBlank = c.isBlank();

                BufferedImage ci = c.getImage();
                int dstH = buf.getHeight() - dstY;
                final int x = (buf.getWidth() - ci.getWidth()) / 2;
                boolean wasSplitting = false;
                while (ci.getHeight() > dstH) {
                    // Split if:
                    // 1) we already decided to split the current chunk,or
                    // 2) the height of the chunk is greater than the threshold.
                    if (wasSplitting || ci.getHeight() > buf.getHeight() * minSplitHeight || processingCoverPage) {
                        wasSplitting = true;

                        // Draw the first part
                        g.drawImage(ci, x, dstY, null);
                        if (debug) {
                            g.setColor(Color.BLACK);
                            g.drawLine(0, dstY, buf.getWidth() - 1, dstY);
                            g.drawLine(0, dstY + ci.getHeight(), buf.getWidth() - 1, dstY + ci.getHeight());
                            g.drawString(ci.getWidth() + "x" + ci.getHeight() + " (SPLIT)", 5, dstY + 12);
                        }

                        // The second part will be drawn at the end of this loop.
                        int overlap = Math.min(dstH, (int) (buf.getHeight() * splitOverlap));
                        ci = ci.getSubimage(0, dstH - overlap, ci.getWidth(), ci.getHeight() - dstH + overlap);
                    } else {
                        // The chunk will not be split and will be drawn at the
                        // end of this loop.
                    }

                    // Finish the current page.
                    BufferedImage dupe = newImage(buf.getWidth(), buf.getHeight());
                    Graphics dupeG = dupe.getGraphics();
                    dupeG.drawImage(buf, 0, 0, null);
                    dupeG.dispose();
                    output.offer(dupe);
                    dstY = 0;
                    dstH = buf.getHeight();

                    // Prepare a new empty page.
                    g.setColor(Color.WHITE);
                    g.fillRect(0, 0, buf.getWidth(), buf.getHeight());

                    if (c.isBlank()) {
                        continue loop;
                    }
                }

                g.drawImage(ci, x, dstY, null);
                if (debug) {
                    g.setColor(Color.BLACK);
                    g.drawLine(0, dstY, buf.getWidth() - 1, dstY);
                    g.drawLine(0, dstY + ci.getHeight(), buf.getWidth() - 1, dstY + ci.getHeight());
                    g.drawString(ci.getWidth() + "x" + ci.getHeight() + (c.isBlank()? " (BLANK)" : ""), 5, dstY + 12);
                }

                dstY += ci.getHeight();
            }
        } finally {
            g.dispose();
        }
    }

    public void finish() {
        finished = true;
        if (dstY != 0) {
            output.offer(buf);
        } else {
            // Do not offer the current buffer if it's empty.
        }
    }

    public BufferedImage poll() {
        BufferedImage img = output.poll();
        if (img == null) {
            return null;
        }

        img = darken(img);

        // Switch to more compact color model
        BufferedImage colorReduced = new BufferedImage(
                img.getWidth(), img.getHeight(),
                //BufferedImage.TYPE_BYTE_GRAY);
                BufferedImage.TYPE_BYTE_BINARY, COLOR_MODEL);
        Graphics g = colorReduced.getGraphics();
        g.drawImage(img, 0, 0, null);
        g.dispose();
        img = colorReduced;

        // Rotate
        img = rotate(img);
        return img;
    }

    private static BufferedImage grayscale(BufferedImage img) {
        final BufferedImage result = newImage(img.getWidth(), img.getHeight());
        final int[] rgb = new int[img.getWidth()];

        for (int y = 0; y < img.getHeight(); y ++) {
            img.getRGB(0, y, img.getWidth(), 1, rgb, 0, rgb.length);
            for (int x = 0; x < img.getWidth(); x ++) {
                int pixel = rgb[x];
                int a = (pixel & 0xff000000) >>> 24;
                int r, g, b;
                if (a == 255) {
                    r = (pixel & 0x00ff0000) >>> 16;
                    g = (pixel & 0x0000ff00) >>> 8;
                    b = pixel & 0x000000ff;
                } else {
                    r = ((pixel & 0x00ff0000) >>> 16) * a / 255 + 255 - a;
                    g = ((pixel & 0x0000ff00) >>> 8) * a / 255 + 255 - a;
                    b = (pixel & 0x000000ff) * a / 255 + 255 - a;
                }

                //int v = Math.min(255, (int) (0.299 * r + 0.587 * g + 0.114 * b + 0.5));
                int v = Math.min(255, (int) (Math.sqrt(
                        0.241 * r * r + 0.691 * g * g + 0.068 * b * b) + 0.5));
                rgb[x] = (v << 8 | v) << 8 | v | 0xff000000;
            }
            result.setRGB(0, y, img.getWidth(), 1, rgb, 0, rgb.length);
        }

        return result;
    }

    private BufferedImage crop(BufferedImage img) {
        assert img.getType() == BufferedImage.TYPE_BYTE_GRAY;

        final int startY = (int) (img.getHeight() * topMargin);
        final int endY = (int) (img.getHeight() * (1.0f - bottomMargin));
        final int startX = (int) (img.getWidth() * leftMargin);
        final int endX = (int) (img.getWidth() * (1.0f - rightMargin));
        final int leftCropThreshold = (int) (img.getHeight() * this.leftCropThreshold);
        final int rightCropThreshold = (int) (img.getHeight() * this.rightCropThreshold);
        final int lineSeparatorLength = (int) (img.getWidth() * lineSeparatorMinWidth);

        int cnt;
        int firstX, lastX;

        loop:
        for (firstX = startX; firstX < endX; firstX ++) {
            cnt = 0;
            for (int y = startY; y < endY; y ++) {
                if (!isWhite(img.getRGB(firstX, y))) {
                    // Determine if the current line is a line separator.
                    int horizLength = 0;
                    if (!processingCoverPage) {
                        final int v1 = img.getRGB(firstX, y) & 0xff;
                        for (int x = firstX; x < img.getWidth(); x ++) {
                            final int v2 = img.getRGB(x, y) & 0xff;
                            if (Math.abs(v1 - v2) > lineSeparatorColorTolerance) {
                                break;
                            } else {
                                horizLength ++;
                                if (horizLength >= lineSeparatorLength) {
                                    break;
                                }
                            }
                        }
                    } else {
                        // A cover page does not have a line separator in general.
                    }

                    if (horizLength < lineSeparatorLength) {
                        cnt ++;
                        if (cnt == leftCropThreshold) {
                            break loop;
                        }
                    } else {
                        // Do not count a line separator.
                    }
                }
            }
        }

        loop:
        for (lastX = endX - 1; lastX >= startX; lastX --) {
            cnt = 0;
            for (int y = startY; y < endY; y ++) {
                if (!isWhite(img.getRGB(lastX, y))) {
                    // Determine if the current line is a line separator.
                    int horizLength = 0;
                    if (!processingCoverPage) {
                        final int v1 = img.getRGB(lastX, y) & 0xff;
                        for (int x = lastX; x >= 0; x --) {
                            final int v2 = img.getRGB(x, y) & 0xff;
                            if (Math.abs(v1 - v2) > lineSeparatorColorTolerance) {
                                break;
                            } else {
                                horizLength ++;
                                if (horizLength >= lineSeparatorLength) {
                                    break;
                                }
                            }
                        }
                    } else {
                        // A cover page does not have a line separator in general.
                    }

                    if (horizLength < lineSeparatorLength) {
                        cnt ++;
                        if (cnt == rightCropThreshold) {
                            break loop;
                        }
                    } else {
                        // Do not count a line separator.
                    }
                }
            }
        }

        firstX = Math.max(firstX - (int) (img.getWidth() * leftCropAdjustment), 0);
        lastX  = Math.min(lastX + (int) (img.getWidth() * rightCropAdjustment),  img.getWidth() - 1);

        if (firstX >= lastX) {
            return img;
        }

        // Correct imbalance.
        int leftCropAmount = firstX;
        int rightCropAmount = img.getWidth() - lastX - 1;
        if (Math.abs(leftCropAmount - rightCropAmount) > img.getWidth() * cropImbalanceLimit) {
            if (leftCropAmount > rightCropAmount) {
                firstX = rightCropAmount + (int) (img.getWidth() * cropImbalanceLimit);
            } else {
                lastX = img.getWidth() - (leftCropAmount + (int) (img.getWidth() * cropImbalanceLimit)) - 1;
            }
            firstX = Math.max(firstX, 0);
            lastX  = Math.min(lastX,  img.getWidth() - 1);
        }

        return img.getSubimage(firstX, 0, lastX - firstX + 1, img.getHeight());
    }

    private static BufferedImage scale(BufferedImage img, int newWidth) {
        final int width = img.getWidth();
        if (width <= newWidth) {
            return img;
        }

        final int newHeight = img.getHeight() * newWidth / width;

        if (newHeight == 0){
            return img;
        }

        ResampleOp resampleOp = new ResampleOp(newWidth, newHeight);
        resampleOp.setFilter(ResampleFilters.getLanczos3Filter());
        return resampleOp.filter(img, null);
    }

    private BufferedImage rotate(BufferedImage img) {
        if (!rotate) {
            return img;
        }

        BufferedImage rotated = newImage(img.getHeight(), img.getWidth());

        for (int x = 0; x < img.getWidth(); x ++) {
            for (int y = 0; y < img.getHeight(); y ++) {
                rotated.setRGB(y, img.getWidth() - x - 1, img.getRGB(x, y));
            }
        }

        return rotated;
    }

    private static BufferedImage darken(BufferedImage img) {
        assert img.getType() == BufferedImage.TYPE_BYTE_GRAY;

        for (int y = 0; y < img.getHeight(); y ++) {
            for (int x = 0; x < img.getWidth(); x ++) {
                int v = img.getRGB(x, y) & 0xff;

                //v = 255 - (int) Math.min((255 - v) * 2.5, 255);
                v = v * v / 255;
                img.setRGB(x, y, (v << 8 | v) << 8 | v | 0xff000000);
            }
        }

        return img;
    }

    private List<Chunk> decompose(BufferedImage img) {
        assert img.getType() == BufferedImage.TYPE_BYTE_GRAY;

        List<Chunk> chunks = new ArrayList<Chunk>();
        boolean firstLineIsBlank = true;
        for (int x = 0; x < img.getWidth(); x++) {
            if (!isWhite(img.getRGB(x, 0))) {
                firstLineIsBlank = false;
                break;
            }
        }

        boolean searchingForBlank = firstLineIsBlank;
        int lastY = 0;

        for (int y = 1; y < img.getHeight(); y ++) {
            boolean blankLine = true;
            for (int x = 0; x < img.getWidth(); x ++) {
                if (!isWhite(img.getRGB(x, y))) {
                    blankLine = false;
                    break;
                }
            }

            if (searchingForBlank) {
                if (!blankLine) {
                    chunks.add(new Chunk(
                            img.getSubimage(0, lastY, img.getWidth(), y - lastY),
                            true));
                    lastY = y;
                    searchingForBlank = false;
                }
            } else {
                if (blankLine) {
                    chunks.add(new Chunk(
                            img.getSubimage(0, lastY, img.getWidth(), y - lastY),
                            false));
                    lastY = y;
                    searchingForBlank = true;
                }
            }
        }

        chunks.add(new Chunk(
                img.getSubimage(0, lastY, img.getWidth(), img.getHeight() - lastY),
                searchingForBlank));

        return chunks;
    }

    private List<Chunk> cropLargeBlanks(Iterable<Chunk> chunks) {
        List<Chunk> newChunks = new ArrayList<Chunk>();
        for (Chunk c: chunks) {
            if (c.isBlank() && c.getImage().getHeight() > maxBlankHeight) {
                c = new Chunk(c.getImage().getSubimage(
                        0, 0, c.getImage().getWidth(), maxBlankHeight), true);
            }
            newChunks.add(c);
        }
        return newChunks;
    }

    private boolean isWhite(int pixel) {
        // We assume grayscale here, hence no need to consider all channels.
        return (pixel & 0xff) > whiteThreshold;
    }

    private static class Chunk {
        private final BufferedImage image;
        private final boolean blank;

        Chunk(BufferedImage image, boolean blank) {
            this.image = image;
            this.blank = blank;
        }

        BufferedImage getImage() {
            return image;
        }

        boolean isBlank() {
            return blank;
        }
    }
}
