#include "hg_gpdf.h"

#if defined(WIN32)
#include "baseapi.h"
#include "allheaders.h"
#include "renderer.h"
#else
#include <leptonica/allheaders.h>
#include <tesseract/baseapi.h>
#include <tesseract/renderer.h>
#endif

//#define USE_QT
#ifdef USE_QT
#include <QDebug>
#include <QTime>
#endif

Pix* createPix(const unsigned char * imagedata, int width, int height, int bytes_per_pixel, int bytes_per_line, int dpi)
{
    int bpp = bytes_per_pixel * 8;
    if (bpp == 0) bpp = 1;
    Pix* pix = pixCreate(width, height, bpp == 24 ? 32 : bpp);
    pixSetXRes(pix, dpi);
    pixSetYRes(pix, dpi);
    l_uint32* data = pixGetData(pix);
    int wpl = pixGetWpl(pix);
    switch (bpp)
    {
    case 1:
        for (int y = 0; y < height; ++y, data += wpl, imagedata += bytes_per_line)
            for (int x = 0; x < width; ++x)
                if (imagedata[x / 8] & (0x80 >> (x % 8)))
                    CLEAR_DATA_BIT(data, x);
                else
                    SET_DATA_BIT(data, x);
        break;

    case 8:
        // Greyscale just copies the bytes in the right order.
        for (int y = 0; y < height; ++y, data += wpl, imagedata += bytes_per_line)
            for (int x = 0; x < width; ++x)
                SET_DATA_BYTE(data, x, imagedata[x]);
        break;

    case 24:
        // Put the colors in the correct places in the line buffer.
        for (int y = 0; y < height; ++y, imagedata += bytes_per_line)
            for (int x = 0; x < width; ++x, ++data) {
                SET_DATA_BYTE(data, COLOR_RED, imagedata[3 * x]);
                SET_DATA_BYTE(data, COLOR_GREEN, imagedata[3 * x + 1]);
                SET_DATA_BYTE(data, COLOR_BLUE, imagedata[3 * x + 2]);
            }
        break;

    case 32:
        // Maintain byte order consistency across different endianness.
        for (int y = 0; y < height; ++y, imagedata += bytes_per_line, data += wpl)
            for (int x = 0; x < width; ++x)
                data[x] = static_cast<l_uint32>((imagedata[x * 4] << 24) | (imagedata[x * 4 + 1] << 16) |
                (imagedata[x * 4 + 2] << 8) | imagedata[x * 4 + 3]);
        break;

    default:
        break;
    }

    pix->informat = bytes_per_pixel == 1 ? 1 : 2;
    if (bytes_per_pixel == 1)
    {
        PIXCMAP* colormap = pixcmapCreate(8);
        LEPT_FREE(colormap->array);
        colormap->array = reinterpret_cast<l_uint8*>(LEPT_CALLOC(256, sizeof(RGBA_QUAD)));
        colormap->n = 256;
        colormap->nalloc = 256;
        colormap->depth = 8;
        l_uint8* ptr = reinterpret_cast<l_uint8*>(colormap->array);
        for (int i = 0; i < 256; i++)
            ptr[i * 4 + 0] = ptr[i * 4 + 1] = ptr[i * 4 + 2] = ptr[i * 4 + 3] = static_cast<l_uint8>(i);
        pixSetColormap(pix, colormap);
    }
    pixSetXRes(pix, 200);
    pixSetYRes(pix, 200);
    return pix;
}

HG_OCR::HG_OCR()
    : m_ptr(nullptr)
{
}

HG_OCR::~HG_OCR()
{
    if (m_ptr)
    {
        reinterpret_cast<tesseract::TessBaseAPI*>(m_ptr)->End();
        delete reinterpret_cast<tesseract::TessBaseAPI*>(m_ptr);
    }
}

int HG_OCR::init(const char* trainFile, RECOGNITION_MODE mode)
{
    if (m_ptr) delete reinterpret_cast<tesseract::TessBaseAPI*>(m_ptr);
    tesseract::TessBaseAPI* api = new tesseract::TessBaseAPI();
    api->SetPageSegMode(tesseract::PSM_AUTO_OSD);
    m_ptr = reinterpret_cast<void*>(api);
    return  api->Init(trainFile, mode == RECOGNITION_OCR ? "chi_sim" : "osd");
}

bool HG_OCR::ocr(const char** inputFileNames, int numOfFiles, const char* outputFileName, SAVE_MODE flag)
{
    if (m_ptr == nullptr) return false;
    if (inputFileNames == nullptr || numOfFiles == 0) return false;

    tesseract::PointerVector<tesseract::TessResultRenderer> renderers;
    tesseract::TessBaseAPI *api = reinterpret_cast<tesseract::TessBaseAPI*>(m_ptr);
    if (flag & SAVE_PDF)
        renderers.push_back(new tesseract::TessPDFRenderer(outputFileName, api->GetDatapath(), false));
    if (flag & SAVE_TXT)
        renderers.push_back(new tesseract::TessTextRenderer(outputFileName));

#ifdef USE_QT
    QTime timer;
    timer.start();
#endif

    int numOfRenderer = renderers.length();
    for (int i = 0; i < numOfRenderer; i++)
        renderers[i]->BeginDocument("");

    for (int i = 0; i < numOfFiles; i++)
    {
        Pix* img = pixRead(inputFileNames[i]);
        api->SetInputName(outputFileName);
        api->SetImage(img);
        api->Recognize(nullptr);
        for (int i = 0; i < numOfRenderer; i++)
            renderers[i]->AddImage(api);
        pixDestroy(&img);
    }

    for (int i = 0; i < numOfRenderer; i++)
        renderers[i]->EndDocument();

    renderers.clear();

#ifdef USE_QT
    qDebug() << timer.elapsed();
#endif

    return true;
}

int HG_OCR::orientation(const unsigned char *data, int width, int height, int channels, int dpi)
{
    Pix* pix = createPix(data, width, height, channels, (width * channels + 3) / 4 * 4, dpi);
    tesseract::TessBaseAPI *api = reinterpret_cast<tesseract::TessBaseAPI*>(m_ptr);
    api->SetImage(pix);
    int orient_deg = 1;
    api->DetectOrientationScript(&orient_deg, nullptr, nullptr, nullptr);
    pixDestroy(&pix);
    return orient_deg;
}