#include "MandelWidget.h"
#include <cmath>

using namespace mnd;

#include <QOpenGLVertexArrayObject>


Texture::Texture(const Bitmap<RGBColor>& bitmap, GLint param)
{
    glGenTextures(1, &id);
    glBindTexture(GL_TEXTURE_2D, id);

    long lineLength = (bitmap.width * 3 + 3) & ~3;

    unsigned char* pixels = new unsigned char[lineLength * bitmap.height];
    for (int i = 0; i < bitmap.width; i++) {
        for (int j = 0; j < bitmap.height; j++) {
            int index = i * 3 + j * lineLength;
            RGBColor c = bitmap.get(i, j);
            pixels[index] = c.r;
            pixels[index + 1] = c.g;
            pixels[index + 2] = c.b;
        }
    }
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, int(bitmap.width), int(bitmap.height), 0, GL_RGB, GL_UNSIGNED_BYTE, pixels);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, param);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, param);
}


Texture::~Texture(void)
{
    if (id != 0)
        glDeleteTextures(1, &id);
}


Texture::Texture(Texture&& other) :
    id{ other.id }
{
    other.id = 0;
}


Texture& Texture::operator=(Texture&& other)
{
    this->id = other.id;
    other.id = 0;
    return *this;
}


void Texture::bind(void) const
{
    glBindTexture(GL_TEXTURE_2D, id);
}


void Texture::drawRect(float x, float y, float width, float height)
{
    glColor3ub(255, 255, 255);
    glEnable(GL_TEXTURE_2D);
    bind();
    glBegin(GL_TRIANGLE_STRIP);
    glTexCoord2f(0, 0);
    glVertex2f(x, y);
    glTexCoord2f(1, 0);
    glVertex2f(x + width, y);
    glTexCoord2f(0, 1);
    glVertex2f(x, y + height);
    glTexCoord2f(1, 1);
    glVertex2f(x + width, y + height);
    glEnd();
    glDisable(GL_TEXTURE_2D);
}


void TextureClip::drawRect(float x, float y, float width, float height)
{
    glColor3ub(255, 255, 255);
    glEnable(GL_TEXTURE_2D);
    glBindTexture(GL_TEXTURE_2D, texture->getId());
    glBegin(GL_TRIANGLE_STRIP);
    glTexCoord2f(tx, ty);
    glVertex2f(x, y);
    glTexCoord2f(tx + tw, ty);
    glVertex2f(x + width, y);
    glTexCoord2f(tx, ty + th);
    glVertex2f(x, y + height);
    glTexCoord2f(tx + tw, ty + th);
    glVertex2f(x + width, y + height);
    glEnd();
    glDisable(GL_TEXTURE_2D);
}


TextureClip TextureClip::clip(float x, float y, float w, float h)
{
    TextureClip result(this->texture);
    result.tx = this->tx + x * this->tw;
    result.ty = this->ty + y * this->th;
    result.tw = this->tw * w;
    result.th = this->th * h;
    return result;
}


TexGrid::TexGrid(MandelV& owner, int level) :
    owner{ owner },
    level{ level },
    dpp{ owner.getDpp(level) }
{
}


std::pair<GridIndex, GridIndex> TexGrid::getCellIndices(double x, double y)
{
    return { ::floor(x / dpp / MandelV::chunkSize), ::floor(y / dpp / MandelV::chunkSize) };
}


std::pair<double, double> TexGrid::getPositions(GridIndex x, GridIndex y)
{
    return { x * dpp * MandelV::chunkSize, y * dpp * MandelV::chunkSize };
}


GridElement* TexGrid::getCell(GridIndex i, GridIndex j)
{
    auto cIt = cells.find({i, j});
    if (cIt != cells.end()) {
        return cIt->second.get();
    }
    else {
        return nullptr;
    }
}


void TexGrid::setCell(GridIndex i, GridIndex j, std::unique_ptr<GridElement> tex)
{
    cells[{i, j}] = std::move(tex);
}


void TexGrid::clearCells(void)
{
    cells.clear();
}


void Job::run(void)
{
    auto [absX, absY] = grid->getPositions(i, j);
    double gw = grid->dpp * MandelV::chunkSize;

    Bitmap<float> f(MandelV::chunkSize, MandelV::chunkSize);
    mnd::MandelInfo mi;
    mi.view.x = absX;
    mi.view.y = absY;
    mi.view.width = mi.view.height = gw;
    mi.bWidth = mi.bHeight = MandelV::chunkSize;
    mi.maxIter = maxIter;
    mndContext.getDefaultGenerator().generate(mi, f.pixels.get());
    auto* rgb = new Bitmap<RGBColor>(f.map<RGBColor>([&mi, this](float i) {
        return i >= mi.maxIter ? RGBColor{ 0, 0, 0 } : gradient.get(i);
    }));
    emit done(level, i, j, rgb);
}


void Calcer::setMaxIter(int maxIter)
{
    this->maxIter = maxIter;
    clearAll();
}


void Calcer::clearAll(void)
{
    this->threadPool->clear();
}


void Calcer::calc(TexGrid& grid, int level, GridIndex i, GridIndex j, int priority)
{
    jobsMutex.lock();
    if (jobs.find({ level, i, j }) == jobs.end()) {
        Job* job = new Job(mndContext, gradient, maxIter, &grid, level, i, j);
        connect(job, &Job::done, this, &Calcer::redirect);
        connect(job, &QObject::destroyed, this, [this, level, i, j] () { this->notFinished(level, i, j); });
        jobs.emplace(std::tuple{level, i, j}, job);
        threadPool->start(job, priority);
    }
    jobsMutex.unlock();
}


void Calcer::setCurrentLevel(int level)
{
    if (this->currentLevel != level) {
        this->currentLevel = level;
        std::vector<QRunnable*> toCancel;
        jobsMutex.lock();
        for (auto&[tup, job] : jobs) {
            auto& [level, i, j] = tup;
            if(level != currentLevel) {
                toCancel.push_back(job);
            }
        }
        jobsMutex.unlock();
        for (auto* job : toCancel) {
            threadPool->cancel(job);
        }
    }
}


void Calcer::notFinished(int level, GridIndex i, GridIndex j)
{
    jobsMutex.lock();
    jobs.erase({ level, i, j });
    jobsMutex.unlock();
}


void Calcer::redirect(int level, GridIndex i, GridIndex j, Bitmap<RGBColor>* bmp)
{
    jobsMutex.lock();
    jobs.erase({ level, i, j });
    jobsMutex.unlock();
    emit done(level, i, j, bmp);
}


MandelV::MandelV(mnd::MandelContext& mndContext, Gradient& gradient, int maxIter) :
    mndContext{ mndContext },
    calcer{ mndContext, gradient, maxIter },
    gradient{ gradient },
    maxIter{ maxIter }
{
    Bitmap<RGBColor> emp(8, 8);
    for(auto i = 0; i < emp.width; i++) {
        for(auto j = 0; j < emp.height; j++) {
            if((i + j) & 0x1) { // if i+j is odd
                emp.get(i, j) = RGBColor{ 255, 255, 255 };
            }
            else {
                emp.get(i, j) = RGBColor{ 120, 120, 120 };
            }
        }
    }
    empty = std::make_unique<Texture>(emp, GL_NEAREST);
    connect(&calcer, &Calcer::done, this, &MandelV::cellReady);
}


int MandelV::getLevel(double dpp) {
    return int(::log2(dpp / chunkSize));
}


double MandelV::getDpp(int level)
{
    return ::pow(2, level) * chunkSize;
}


TexGrid& MandelV::getGrid(int level)
{
    auto it = levels.find(level);
    if (it != levels.end()) {
        return it->second;
    }
    else {
        levels.insert(std::pair<int, TexGrid>{ level, TexGrid{ *this, level } });
        return levels.at(level);
    }
}


void MandelV::setMaxIter(int maxIter)
{
    this->maxIter = maxIter;
    calcer.setMaxIter(maxIter);
    clear();
    emit redrawRequested();
}



void MandelV::clear(void)
{
    for(auto& [level, grid] : this->levels) {
        grid.clearCells();
    }
}


void MandelV::garbageCollect(int level)
{
    for(auto& [l, grid] : levels) {
        int dist = ::abs(l - level);
        if (dist > 20) {
            grid.clearCells();
        }
        else if (dist > 10) {
            if (grid.countAllocatedCells() > 50)
                grid.clearCells();
        }
        else if (dist > 3) {
            if (grid.countAllocatedCells() > 150)
                grid.clearCells();
        }
        else if (dist > 0) {
            if (grid.countAllocatedCells() > 350)
                grid.clearCells();
        }
        else {
            if (grid.countAllocatedCells() > 2500)
                grid.clearCells();
        }
    }
}


GridElement* MandelV::searchAbove(int level, GridIndex i, GridIndex j, int recursionLevel)
{
    auto& grid = getGrid(level);
    auto& gridAbove = getGrid(level + 1);

    GridIndex ai = (i < 0 ? (i - 1) : i) / 2;
    GridIndex aj = (j < 0 ? (j - 1) : j) / 2;

    GridElement* above = gridAbove.getCell(ai, aj);

    if (above == nullptr && recursionLevel > 0) {
        auto abFound = searchAbove(level + 1, ai, aj, recursionLevel - 1);
        if (abFound)
            above = abFound;
    }

    if (above != nullptr) {
        auto newElement = std::make_unique<GridElement>(
            false, above->img.clip((i & 1) * 0.5f, (j & 1) * 0.5f, 0.5f, 0.5f)
        );
        GridElement* ret = newElement.get();
        grid.setCell(i, j, std::move(newElement));
        return ret;
    }
    else {
        return nullptr;
    }
}


std::unique_ptr<GridElement> MandelV::searchUnder(int level, GridIndex i, GridIndex j, int recursionLevel)
{
    /*if (recursionLevel == 0)
        return nullptr;

    auto& gridUnder = getGrid(level - 1);

    int ai = i * 2;
    int aj = j * 2;

    GridElement* u00 = gridUnder.getCell(ai, aj);
    GridElement* u01 = gridUnder.getCell(ai, aj + 1);
    GridElement* u10 = gridUnder.getCell(ai + 1, aj);
    GridElement* u11 = gridUnder.getCell(ai + 1, aj + 1);

    if (   u00 != nullptr
        && u01 != nullptr
        && u10 != nullptr
        && u11 != nullptr) {
        GLuint FramebufferName = 0;
        auto newElement = std::make_unique<GridElement>(
            false, above->img.clip((i & 1) * 0.5f, (j & 1) * 0.5f, 0.5f, 0.5f)
        );
        return newElement;
    }
    else {
        return searchAbove(level - 1, ai, aj, recursionLevel - 1);
    }*/
    return nullptr;
}


void MandelV::paint(const mnd::MandelViewport& mvp)
{
    double dpp = mvp.width / width;
    int level = getLevel(dpp) - 1;
    garbageCollect(level);
    emit calcer.setCurrentLevel(level);

    auto& grid = getGrid(level);
    double gw = getDpp(level) * chunkSize;

    double w = width * gw / mvp.width;
    //double h = height * gw / mvp.height;
    //printf("level: %d, dpp: %f, width: %f\n", level, dpp, w);

    auto [left, top] = grid.getCellIndices(mvp.x, mvp.y);
    auto [right, bottom] = grid.getCellIndices(mvp.right(), mvp.bottom());
    auto [realXLeft, realYTop] = grid.getPositions(left, top);
    realXLeft = (realXLeft - mvp.x) * width / mvp.width;
    realYTop = (realYTop - mvp.y) * height / mvp.height;
    for(GridIndex i = left; i <= right; i++) {
        for(GridIndex j = top; j <= bottom; j++) {
            double x = realXLeft + (i - left) * w;
            double y = realYTop + (j - top) * w;

            GridElement* t = grid.getCell(i, j);

            if (t == nullptr) {
                auto above = searchAbove(level, i, j, 2);
                if (above) {
                    t = above;
                }
            }

            if (t != nullptr) {
                t->img.drawRect(x, y, w, w);
                /*glBegin(GL_LINE_LOOP);
                glVertex2f(x, y);
                glVertex2f(x + w, y);
                glVertex2f(x + w, y + w);
                glVertex2f(x, y + w);
                glEnd();*/

                if (!t->enoughResolution) {
                    calcer.calc(grid, level, i, j, 100);
                }
            }
            else {
                calcer.calc(grid, level, i, j, 1000);
                this->empty->drawRect(x, y, w, w);
            }
        }
    }
}

void MandelV::cellReady(int level, GridIndex i, GridIndex j, Bitmap<RGBColor>* bmp)
{

    this->getGrid(level).setCell(i, j,
        std::make_unique<GridElement>(true, TextureClip{ std::make_shared<Texture>(*bmp) }));
    delete bmp;
    emit redrawRequested();
}


MandelWidget::MandelWidget(mnd::MandelContext& ctxt, QWidget* parent) :
    QOpenGLWidget{ parent },
    mndContext{ ctxt }
{
    this->setContentsMargins(0, 0, 0, 0);
    this->setSizePolicy(QSizePolicy::Expanding,
        QSizePolicy::Expanding);
    qRegisterMetaType<GridIndex>("GridIndex");
}


MandelWidget::~MandelWidget()
{
}


void MandelWidget::initializeGL(void)
{
    this->context()->functions()->glClearColor(0, 0, 0, 0);
    this->context()->makeCurrent(nullptr);

    glDisable(GL_DEPTH_TEST);

    // looks not even better
    //glDisable(GL_FRAMEBUFFER_SRGB);

    //glShadeModel(GL_SMOOTH);

    v = nullptr;
    requestRecalc();
}


void MandelWidget::paintGL(void)
{
    if (v == nullptr) {
        v = std::make_unique<MandelV>(mndContext, gradient, maxIterations);
        QObject::connect(v.get(), &MandelV::redrawRequested, this, static_cast<void(QOpenGLWidget::*)(void)>(&QOpenGLWidget::update));
    }

    int width = this->width();
    int height = this->height();
    v->width = width;
    v->height = height;

    glViewport(0, 0, width, height);

    glMatrixMode(GL_PROJECTION);
    glLoadIdentity();
#ifdef QT_OPENGL_ES_1
    glOrthof(0, width, height, 0, -1.0, 1.0);
#else
    glOrtho(0, width, height, 0, -1.0, 1.0);
#endif
    glMatrixMode(GL_MODELVIEW);

    glClear(GL_COLOR_BUFFER_BIT);
    glLoadIdentity();

    v->paint(this->viewport);

    if (rubberbanding)
        drawRubberband();
}


void MandelWidget::drawRubberband(void)
{
    glColor3ub(10, 200, 10);
    glBegin(GL_LINE_LOOP);
    glVertex2d(rubberband.x(), rubberband.y());
    glVertex2d(rubberband.right(), rubberband.y());
    glVertex2d(rubberband.right(), rubberband.bottom());
    glVertex2d(rubberband.x(), rubberband.bottom());
    glEnd();

    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
    glColor4f(0.1f, 0.9f, 0.1f, 0.2f);
    glBegin(GL_TRIANGLE_FAN);
    glVertex2d(rubberband.x(), rubberband.y());
    glVertex2d(rubberband.right(), rubberband.y());
    glVertex2d(rubberband.right(), rubberband.bottom());
    glVertex2d(rubberband.x(), rubberband.bottom());
    glEnd();
    glDisable(GL_BLEND);
}


void MandelWidget::zoom(float scale, float x, float y)
{
    viewport.zoom(scale, x, y);
    //viewport.zoomCenter(scale);
    requestRecalc();
}


void MandelWidget::setMaxIterations(int maxIter)
{
    this->maxIterations = maxIter;
    if (v)
        v->setMaxIter(maxIter);
    requestRecalc();
}


void MandelWidget::requestRecalc()
{
    //emit needsUpdate(MandelInfo{ viewport, this->width(), this->height(), maxIterations });
    this->update();
}


void MandelWidget::resizeGL(int width, int height)
{
    //glViewport(0, 0, (GLint) width, (GLint) height);
    this->update();
}


/*void MandelWidget::redraw(void)
{
    /*CpuGenerator<double> cpg;
    MandelInfo mi;
    mi.bWidth = this->geometry().width();//ql.geometry().width();
    mi.bHeight = this->geometry().height(); //ql.geometry().height();
    mi.maxIter = 250;
    mi.view = viewport;*/
    //update();
    //emit needsUpdate(viewport);

    //auto bitmap = cpg.generate(mi).map<uint32_t>([](RGBColor rgb) { return 255 << 24 | rgb.b << 16 | rgb.g << 8 | rgb.r; });
//}


void MandelWidget::resizeEvent(QResizeEvent* re)
{
    QOpenGLWidget::resizeEvent(re);
    double aspect = double(geometry().width()) / geometry().height();

    //if (viewport.width > viewport.height * aspect)
        viewport.height = (viewport.width / aspect);
    //else
    //    viewport.width = (viewport.height * aspect);

    if (v.get() != nullptr) {
        v->width = this->width();
        v->height = this->height();
    }
    //printf("resized\n");
    requestRecalc();
    //redraw();
}


void MandelWidget::mousePressEvent(QMouseEvent* me)
{
    QOpenGLWidget::mousePressEvent(me);
    if (me->button() == Qt::RightButton) {
        rubberbanding = true;
        rubberband.setCoords(me->x(), me->y(), 0, 0);
        //emit repaint();
        me->accept();
    }
    else if (me->button() == Qt::LeftButton) {
        dragging = true;
        dragX = me->x();
        dragY = me->y();
        me->accept();
    }
}


void MandelWidget::mouseMoveEvent(QMouseEvent* me)
{
    QOpenGLWidget::mouseMoveEvent(me);
    if (rubberbanding) {
        QRectF& rect = rubberband;
        double aspect = double(geometry().width()) / geometry().height();
        rect.setBottomRight(QPoint(me->x(), me->y()));
        if (rect.width() > rect.height() * aspect)
            rect.setHeight(rect.width() / aspect);
        else
            rect.setWidth(rect.height() * aspect);

        emit repaint();
    }
    else if (dragging) {
        double deltaX = me->x() - dragX;
        double deltaY = me->y() - dragY;

        this->viewport.x -= deltaX * viewport.width / this->width();
        this->viewport.y -= deltaY * viewport.height / this->height();
        dragX = me->x(); dragY = me->y();

        emit repaint();
    }
    me->accept();
}


void MandelWidget::mouseReleaseEvent(QMouseEvent* me)
{
    QOpenGLWidget::mouseReleaseEvent(me);
    if (rubberbanding) {
        QRect rect = rubberband.toRect();
        QRect full = this->geometry();

        viewport.x += double(rect.left()) * viewport.width / full.width();
        viewport.y += double(rect.top()) * viewport.height / full.height();
        viewport.width *= double(rect.width()) / full.width();
        viewport.height *= double(rect.height()) / full.height();
        viewport.normalize();
        requestRecalc();
        rubberbanding = false;
    }
    dragging = false;

    //requestRecalc();
}

void MandelWidget::wheelEvent(QWheelEvent* we)
{
    QOpenGLWidget::wheelEvent(we);
    float x = float(we->x()) / this->width();
    float y = float(we->y()) / this->height();
    float scale = ::pow(0.9975, we->angleDelta().y());
    //printf("scale: %f\n", double(scale));
    zoom(scale, x, y);
    we->accept();
}


void MandelWidget::viewUpdated(Bitmap<RGBColor>* bitmap)
{
    if (bitmap != nullptr) {
        tex = std::make_unique<Texture>(*bitmap);
        delete bitmap;
        emit repaint();
    }
}