#include "loadrunable.h"
#include "imageringlist.h"
#include <QFileInfo>
#include <QPainter>
#include <QElapsedTimer>
#include <QDebug>
#include <algorithm>
#include <fitsio2.h>
#include "rawimage.h"
#include "loadimage.h"
#include <lcms2.h>

LoadRunable::LoadRunable(const QString &file, Image *receiver, AnalyzeLevel level, bool thumbnail) :
    m_file(makeMaxPath(file)),
    m_receiver(receiver),
    m_analyzeLevel(level),
    m_thumbnail(thumbnail)
{
}

void LoadRunable::run()
{
    try
    {
        if(!m_thumbnail && !m_receiver->isCurrent())
        {
            return;
        }
        QElapsedTimer timer;
        ImageInfoData info;
        QFileInfo finfo(m_file);
        info.info.append({QObject::tr("Filename"), finfo.fileName()});

        std::shared_ptr<RawImage> rawImage;
        if(!loadImage(m_file, info, rawImage))
            info.info.append({QObject::tr("Error"), QObject::tr("Failed to load image")});


        if(rawImage && !m_thumbnail)
        {
            rawImage->convertToGLFormat();
            timer.start();
            rawImage->generateLUT();
            qDebug() << "generate LUT" << timer.restart();
            //rawImage->convertTosRGB();
            //qDebug() << "convert" << timer.restart();
            rawImage->calcStats();
            const RawImage::Stats &stats = rawImage->imageStats();
            qDebug() << "image stats" << timer.restart();
            if(rawImage->channels() == 1)
            {
                info.info.append({QObject::tr("Mean"), QString::number(stats.m_mean[0])});
                info.info.append({QObject::tr("Standart deviation"), QString::number(stats.m_stdDev[0])});
                info.info.append({QObject::tr("Median"), QString::number(stats.m_median[0])});
                info.info.append({QObject::tr("Minimum"), QString::number(stats.m_min[0])});
                info.info.append({QObject::tr("Maximum"), QString::number(stats.m_max[0])});
                info.info.append({QObject::tr("MAD"), QString::number(stats.m_mad[0])});
                info.info.append({QObject::tr("Saturated"), QString::number(100.0 * stats.m_saturated[0] / rawImage->size()) + "%"});
            }
            else
            {
                info.info.append({QObject::tr("Mean"), QString("%1 %2 %3").arg(stats.m_mean[0]).arg(stats.m_mean[1]).arg(stats.m_mean[2])});
                info.info.append({QObject::tr("Standart deviation"), QString("%1 %2 %3").arg(stats.m_stdDev[0]).arg(stats.m_stdDev[1]).arg(stats.m_stdDev[2])});
                info.info.append({QObject::tr("Median"), QString("%1 %2 %3").arg(stats.m_median[0]).arg(stats.m_median[1]).arg(stats.m_median[2])});
                info.info.append({QObject::tr("Minimum"), QString("%1 %2 %3").arg(stats.m_min[0]).arg(stats.m_min[1]).arg(stats.m_min[2])});
                info.info.append({QObject::tr("Maximum"), QString("%1 %2 %3").arg(stats.m_max[0]).arg(stats.m_max[1]).arg(stats.m_max[2])});
                info.info.append({QObject::tr("MAD"), QString("%1 %2 %3").arg(stats.m_mad[0]).arg(stats.m_mad[1]).arg(stats.m_mad[2])});
                info.info.append({QObject::tr("Saturated"), QString("%1 %2 %3%").arg(100.0 * stats.m_saturated[0] / rawImage->size())
                                  .arg(100.0 * stats.m_saturated[1] / rawImage->size())
                                  .arg(100.0 * stats.m_saturated[2] / rawImage->size())});
            }
        }

        if(m_thumbnail)
        {
            if(rawImage && rawImage->valid())
            {
                if(QUALITY_RESIZE)
                    rawImage->resize(THUMB_SIZE, THUMB_SIZE);

                rawImage->convertToGLFormat();
                rawImage->convertToThumbnail();
            }
            QMetaObject::invokeMethod(m_receiver, "thumbnailLoadFinish", Qt::QueuedConnection, Q_ARG(std::shared_ptr<RawImage>, rawImage));
        }
        else
        {
            QMetaObject::invokeMethod(m_receiver, "imageLoaded", Qt::QueuedConnection, Q_ARG(std::shared_ptr<RawImage>, rawImage), Q_ARG(ImageInfoData, info));
        }
    }
    catch(std::exception e)
    {
        qDebug() << m_file << e.what();
        std::shared_ptr<RawImage> rawImage;
        if(m_thumbnail)
            QMetaObject::invokeMethod(m_receiver, "thumbnailLoadFinish", Qt::QueuedConnection, Q_ARG(std::shared_ptr<RawImage>, rawImage));
        else
            QMetaObject::invokeMethod(m_receiver, "imageLoaded", Qt::QueuedConnection, Q_ARG(std::shared_ptr<RawImage>, rawImage), Q_ARG(ImageInfoData, ImageInfoData()));
    }
}

ConvertRunable::ConvertRunable(const QString &in, const QString &out, const QString &format, const ConvertParams &params, QSemaphore *semaphore) :
    m_infile(makeMaxPath(in)),
    m_outfile(makeMaxPath(out)),
    m_format(format),
    m_params(params),
    m_semaphore(semaphore)
{
}

void writeFITSImage(fitsfile *fw, std::shared_ptr<RawImage> rawimage, ImageInfoData &imageinfo)
{
    static QStringList skipKeys = {"SIMPLE", "BITPIX", "NAXIS", "NAXIS1", "NAXIS2", "NAXIS3", "BZERO", "BSCALE", "EXTEND"};

    int status = 0;
    long firstpix[3] = {1,1,1};

    int channels = rawimage->channels();
    int naxis = channels == 1 ? 2 : 3;
    long naxes[3] = {(int)rawimage->width(), (int)rawimage->height(), rawimage->channels()};

    std::vector<RawImage> planes;
    if(channels == 1)
        planes.push_back(*rawimage);
    else
        planes = rawimage->split();

    switch(rawimage->type())
    {
    case RawImage::UINT8:
        fits_create_img(fw, BYTE_IMG, naxis, naxes, &status);
        for(int i=0; i<channels; i++)
        {
            firstpix[2] = i+1;
            fits_write_pix(fw, TBYTE, firstpix, rawimage->size(), planes[i].data(), &status);
        }
        break;
    case RawImage::UINT16:
        fits_create_img(fw, USHORT_IMG, naxis, naxes, &status);
        for(int i=0; i<channels; i++)
        {
            firstpix[2] = i+1;
            fits_write_pix(fw, TUSHORT, firstpix, rawimage->size(), planes[i].data(), &status);
        }
        break;
    case RawImage::FLOAT32:
        fits_create_img(fw, FLOAT_IMG, naxis, naxes, &status);
        for(int i=0; i<channels; i++)
        {
            firstpix[2] = i+1;
            fits_write_pix(fw, TFLOAT, firstpix, rawimage->size(), planes[i].data(), &status);
        }
        break;
    default:
        return;
    }
    for(const FITSRecord &record : imageinfo.fitsHeader)
    {
        if(skipKeys.contains(record.key) || record.xisf)continue;

        bool isdouble;
        bool isint;
        bool isbool = record.value.toString() == "T" || record.value.toString() == "F";
        double vald = record.value.toDouble(&isdouble);
        int valb = record.value.toString() == "T";
        long long vall = record.value.toLongLong(&isint);
        QByteArray str = record.value.toString().toLatin1();
        if(isdouble)
            fits_write_key(fw, TDOUBLE, record.key.data(), &vald, record.comment.isEmpty() ? nullptr : record.comment.data(), &status);
        else if(isint)
            fits_write_key(fw, TLONGLONG, record.key.data(), &vall, record.comment.isEmpty() ? nullptr : record.comment.data(), &status);
        else if(isbool)
            fits_write_key(fw, TLOGICAL, record.key.data(), &valb, record.comment.isEmpty() ? nullptr : record.comment.data(), &status);
        else if(record.key == "COMMENT")
            fits_write_comment(fw, record.comment.isEmpty() ? nullptr : record.comment.data(), &status);
        else if(record.key == "HISTORY")
            fits_write_history(fw, record.comment.isEmpty() ? nullptr : record.comment.data(), &status);
        else
            fits_write_key(fw, TSTRING, record.key.data(), str.isEmpty() ? nullptr : str.data(), record.comment.isEmpty() ? nullptr : record.comment.data(), &status);
    }
}

void ConvertRunable::run()
{
    QSemaphoreReleaser release;
    if(m_semaphore)release = QSemaphoreReleaser(m_semaphore);

    ImageInfoData imageinfo;
    std::shared_ptr<RawImage> rawimage;
    loadImage(m_infile, imageinfo, rawimage);
    QFileInfo info(m_outfile);
    info.dir().mkpath(".");

    if(m_params.autostretch)
    {
        rawimage->calcStats();
        MTFParam mtfParam = rawimage->calcMTFParams();
        rawimage->applySTF(mtfParam);
    }
    if(m_params.binning > 1)
    {
        rawimage->resizeInt(m_params.binning, m_params.average);
    }
    if(m_params.resize.isValid() && !m_params.resize.isEmpty())
    {
        QSize imgSize(rawimage->width(), rawimage->height());
        imgSize = imgSize.scaled(m_params.resize, m_params.aspect);
        rawimage->resize(imgSize.width(), imgSize.height());
    }

    if(rawimage)
    {
        if(m_format == "xisf")
        {
            try
            {
                LibXISF::XISFWriter xisf;
                int channelCount = rawimage->channels();
                LibXISF::Image::SampleFormat sampleFormat;
                switch(rawimage->type())
                {
                case RawImage::UINT8: sampleFormat = LibXISF::Image::UInt8; break;
                case RawImage::UINT16: sampleFormat = LibXISF::Image::UInt16; break;
                case RawImage::FLOAT32: sampleFormat = LibXISF::Image::Float32; break;
                default: return;
                }

                LibXISF::Image image(rawimage->width(), rawimage->height(), channelCount, sampleFormat, channelCount == 1 ? LibXISF::Image::Gray : LibXISF::Image::RGB, LibXISF::Image::Planar);
                if(channelCount == 1)
                {
                    std::memcpy(image.imageData(), rawimage->data(), image.imageDataSize());
                }
                else
                {
                    size_t off = 0;
                    std::vector<RawImage> planes = rawimage->split();
                    for(const auto &plane : planes)
                    {
                        std::memcpy(image.imageData<uint8_t>() + off, plane.data(), plane.size() * RawImage::typeSize(plane.type()));
                        off += plane.size() * RawImage::typeSize(plane.type());
                    }
                }
                for(auto &record : imageinfo.fitsHeader)
                {
                    if(record.xisf)continue;

                    if(record.value.typeId() == QMetaType::Bool)
                        image.addFITSKeyword({record.key.toStdString(), record.value.toBool() ? "T" : "F", record.comment.toStdString()});
                    else
                        image.addFITSKeyword({record.key.toStdString(), record.value.toString().toStdString(), record.comment.toStdString()});
                }

                if(m_params.compressionType.startsWith("zstd") && LibXISF::DataBlock::CompressionCodecSupported(LibXISF::DataBlock::ZSTD))
                    image.setCompression(LibXISF::DataBlock::ZSTD, m_params.compressionLevel);
                else if(m_params.compressionType.startsWith("lz4hc"))
                    image.setCompression(LibXISF::DataBlock::LZ4HC, m_params.compressionLevel);
                else if(m_params.compressionType.startsWith("lz4"))
                    image.setCompression(LibXISF::DataBlock::LZ4, m_params.compressionLevel);
                else if(m_params.compressionType.startsWith("zlib"))
                    image.setCompression(LibXISF::DataBlock::Zlib, m_params.compressionLevel);

                if(m_params.compressionType.endsWith("+sh"))
                    image.setByteshuffling(true);

                xisf.writeImage(image);
                xisf.save(m_outfile.toLocal8Bit().data());
            }
            catch(LibXISF::Error &err)
            {
                qDebug() << "Failed to save XISF image" <<  err.what();
            }
            return;
        }

        if(m_format == "fits")
        {
            int status = 0;
            fitsfile *fw;
            if(QFileInfo(m_outfile).exists())QFile::remove(m_outfile);
            fits_create_diskfile(&fw, m_outfile.toLocal8Bit().data(), &status);
            if(!m_params.compressionType.isEmpty())
            {
                if(m_params.compressionType == "gzip")
                    fits_set_compression_type(fw, GZIP_1, &status);
                else if(m_params.compressionType == "rice")
                    fits_set_compression_type(fw, RICE_1, &status);
            }
            writeFITSImage(fw, rawimage, imageinfo);
            fits_close_file(fw, &status);
            return;
        }

        // if nothing else try QImage
        {
            QImage::Format format = QImage::Format_Invalid;
            int width = rawimage->widthBytes();

            switch(rawimage->type())
            {
            case RawImage::UINT8:
                if(rawimage->channels() == 1)format = QImage::Format_Grayscale8;
                else if(rawimage->channels() == 3)format = QImage::Format_RGBX8888;
                else if(rawimage->channels() == 4)format = QImage::Format_RGBA8888;
                break;
            case RawImage::UINT16:
                if(rawimage->channels() == 1)format = QImage::Format_Grayscale16;
                else if(rawimage->channels() == 3)format = QImage::Format_RGBX64;
                else if(rawimage->channels() == 4)format = QImage::Format_RGBA64;
                width *= 2;
                break;
            case RawImage::FLOAT16:
            case RawImage::FLOAT32:
            case RawImage::FLOAT64:
            case RawImage::UINT32:
                rawimage->convertToType(RawImage::UINT16);
                if(rawimage->channels() == 1)format = QImage::Format_Grayscale16;
                else if(rawimage->channels() == 3)format = QImage::Format_RGBX64;
                else if(rawimage->channels() == 4)format = QImage::Format_RGBA64;
                width *= 2;
                break;
            }

            if(format == QImage::Format_Invalid)return;

            QImage qimage(rawimage->width(), rawimage->height(), format);
            for(uint32_t i=0; i < rawimage->height(); i++)
                std::memcpy(qimage.scanLine(i), rawimage->data(i), width);
            qimage.save(m_outfile);
        }
    }
}

ConvertRunable::ConvertParams::ConvertParams(const QVariantMap &map)
{
    bool ok = false;
    if(map.contains("compressionLevel"))
        compressionLevel = std::clamp(map["compressionLevel"].toInt(&ok), -1, 100);

    if(!ok)compressionLevel = -1;

    if(map.contains("compressionType"))
        compressionType = map["compressionType"].toString();

    if(map.contains("binning"))
        binning = map["binning"].toInt();

    if(map.contains("average"))
        average = map["average"].toBool();

    if(map.contains("resize"))
    {
        QVariantMap size = map["resize"].toMap();
        if(size.contains("width") && size.contains("height"))
        {
            int w = size["width"].toInt();
            int h = size["height"].toInt();
            resize = QSize(w, h);
        }
        if(size.contains("aspect"))
        {
            QString aspectStr = map["aspect"].toString();
            if(aspectStr == "keep")
                aspect = Qt::KeepAspectRatio;
            else if(aspectStr == "expand")
                aspect = Qt::KeepAspectRatioByExpanding;
            else if(aspectStr == "ignore")
                aspect = Qt::IgnoreAspectRatio;
        }
    }

    if(map.contains("autostretch"))
        autostretch = map["autostretch"].toBool();
}