#include "loadrunable.h"
#include "imageringlist.h"
#include <libraw/libraw.h>
#include "imageinfo.h"
#include <QFileInfo>
#include <QPainter>
#include <QElapsedTimer>
#include <QDebug>
#include <iostream>
#include <libexif/exif-data.h>
#include <fitsio2.h>
#include <pcl/XISF.h>
#include "rawimage.h"
#include "starfit.h"
#include "wcslib/wcshdr.h"

#ifdef COLOR_MANAGMENT
#include <QColorSpace>
static pcl::ICCProfile sRgbIccProfile((void*)QColorSpace(QColorSpace::SRgb).iccProfile().data());
#endif

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

void loadExifEntry(ImageInfoData &info, ExifContent *content, ExifTag tag)
{
    char val[1024];
    ExifEntry *entry = exif_content_get_entry(content, tag);
    if(entry)
    {
        exif_entry_get_value(entry, val, sizeof(val));
        info.info.append({exif_tag_get_title(tag), val});
    }
}

void drawPeaks(QImage &img, const std::vector<Peak> &peaks)
{
    QPixmap pix = QPixmap::fromImage(img);
    QPainter painter(&pix);
    painter.setPen(Qt::red);
    for(auto peak : peaks)
    {
        painter.drawEllipse(QPoint(peak.x(), peak.y()), 5, 5);
    }
    img = pix.toImage();
}

void drawStars(QImage &img, const std::vector<Star> &stars)
{
    QPixmap pix = QPixmap::fromImage(img);
    QPainter painter(&pix);
    painter.setPen(Qt::red);
    for(auto star : stars)
    {
        painter.drawEllipse(QPointF(star.m_x, star.m_y), star.hw20X(), star.hw20Y());
    }
    img = pix.toImage();
}

void printStarModel(int radius, const std::vector<double> &data, const Star &star)
{
    QString d = "d=[";
    QString m = "m=[";
    for(int y=0; y<radius; y++)
    {
        for(int x=0; x<radius; x++)
        {
            d += QString::number(data[y*radius+x]) + ",";
            m += QString::number(gauss_model(star.m_am, star.m_x, star.m_y, star.m_sx, star.m_sy, x, y)) + ",";
        }
        d += ";";
        m += ";";
    }
    d += "];";
    m += "];";
    //std::cout << star.m_am << " " << star.m_sx << star.m_sy << std::endl;
    std::cout << d.toStdString() << std::endl;
    std::cout << m.toStdString() << std::endl << std::endl;
}

bool loadRAW(const QString path, ImageInfoData &info, RawImage **image)
{
    if(!image)
        return false;

    std::unique_ptr<LibRaw> raw = std::make_unique<LibRaw>();
    raw->open_file(path.toLocal8Bit().data());
    raw->imgdata.params.half_size = true;
    raw->imgdata.params.use_camera_wb = true;
    raw->imgdata.params.user_flip = 0;
    if(raw->unpack())
        return false;

    if(image)
    {
        libraw_rawdata_t rawdata = raw->imgdata.rawdata;
        size_t size = rawdata.sizes.width*rawdata.sizes.height;

        std::vector<uint16_t> out;
        out.resize(size);
        size_t d = 0;
        uint h=rawdata.sizes.top_margin+rawdata.sizes.height;
        uint w=rawdata.sizes.left_margin+rawdata.sizes.width;
        size_t pitch = rawdata.sizes.raw_pitch/sizeof(uint16_t);

        for(size_t i=rawdata.sizes.top_margin;i<h;i++)
        {
            for(size_t o=rawdata.sizes.left_margin;o<w;o++)
            {
                uint16_t p = rawdata.raw_image[i*pitch+o];
                out[d++] = p;
            }
        }
        *image = new RawImage(rawdata.sizes.width, rawdata.sizes.height, RawImage::UINT16);
        memcpy((*image)->data(), &out[0], sizeof(uint16_t)*d);
    }

    QString shutterSpeed = QString::number(raw->imgdata.other.shutter);
    if(raw->imgdata.other.shutter < 1)
    {
        shutterSpeed = QString("1/%1s").arg(1.0f/raw->imgdata.other.shutter);
    }
    //info.append(StringPair(QObject::tr("Width"), QString::number(rawImg->width)));
    //info.append(StringPair(QObject::tr("Height"), QString::number(rawImg->height)));
    info.info.append({QObject::tr("ISO"), QString::number(raw->imgdata.other.iso_speed)});
    info.info.append({QObject::tr("Shutter speed"), shutterSpeed});
#if LIBRAW_MINOR_VERSION>=19
//    info.append(StringPair(QObject::tr("Camera temperature"), QString::number(raw.imgdata.other.CameraTemperature)));
#endif
    return true;
}

int loadFITSHeader(fitsfile *file, ImageInfoData &info)
{
    int imgtype;
    int naxis;
    long naxes[3] = {0};
    int nexist;
    int status = 0;
    char key[FLEN_KEYWORD];
    char val[FLEN_VALUE];
    char comm[FLEN_COMMENT];
    char strval[FLEN_VALUE];
    QVariant var;
    fits_get_img_param(file, 3, &imgtype, &naxis, naxes, &status);
    fits_get_hdrspace(file, &nexist, nullptr, &status);
    for(int i=1; i<=nexist; i++)
    {
        fits_read_keyn(file, i, key, val, comm, &status);
        fits_read_key(file, TSTRING, key, strval, nullptr, &status);
        if(status == 0 || status == VALUE_UNDEFINED)
        {
            QString string(strval);
            bool isint;
            bool isdouble;
            double vald = string.toDouble(&isdouble);
            long long vall = string.toLongLong(&isint);
            if(isint)
                var = vall;
            else if(isdouble)
                var = vald;
            else if(status == VALUE_UNDEFINED)
                var = QVariant();
            else if(string == "T" || string == "F")
                var = string == "T";
            else
                var = string;
            status = 0;
            info.fitsHeader.append(FITSRecord(key, var, comm));
        }
        else
        {
            return status;
        }
    }

    char *header = nullptr;
    int nrec = 0;
    const char *exclist[] = {"PV1_1", "PV1_2"};
    fits_hdr2str(file, TRUE, (char**)exclist, 2, &header, &nrec, &status);
    if(status == 0)
    {
        info.wcs = std::make_shared<WCSData>(naxes[0], naxes[1], header, nrec);
        if(!info.wcs->valid())info.wcs.reset();
    }
    fits_free_memory(header, &status);
    return status;
}

bool loadFITS(const QString path, ImageInfoData &info, RawImage **image)
{
    if(!image)
        return false;

    fitsfile *file;
    int status = 0;
    int type;
    fits_open_image(&file, path.toLocal8Bit().data(), READONLY, &status);
    fits_get_hdu_type(file, &type, &status);

    if(type == IMAGE_HDU)
    {
        int imgtype;
        int naxis;
        long naxes[3] = {0};
        fits_get_img_param(file, 3, &imgtype, &naxis, naxes, &status);
        fits_get_img_equivtype(file, &imgtype, &status);

        if(naxis >= 2 && naxis <= 3 && status == 0)
        {
            int cvtype;
            int fitstype;
            std::vector<cv::Mat> cvimg;
            long fpixel[3] = {1,1,1};
            switch(imgtype)
            {
            case BYTE_IMG:
                cvtype = CV_8U;
                fitstype = TBYTE;
                break;
            case SHORT_IMG:
                cvtype = CV_16S;
                fitstype = TSHORT;
                break;
            case USHORT_IMG:
                cvtype = CV_16U;
                fitstype = TUSHORT;
                break;
            case FLOAT_IMG:
                cvtype = CV_32F;
                fitstype = TFLOAT;
                break;
            default:
                info.info.append({QObject::tr("Error"), QObject::tr("Unsupported sample format")});
                goto noload;
                break;
            }

            size_t size = naxes[0]*naxes[1];
            size_t w = naxes[0];
            size_t h = naxes[1];

            info.info.append({QObject::tr("Width"), QString::number(naxes[0])});
            info.info.append({QObject::tr("Height"), QString::number(naxes[1])});

            for (int i=1; i==1 || i<=naxes[2]; i++)
            {
                cv::Mat tmp(h, w, cvtype);
                fpixel[2] = i;
                fits_read_pix(file, fitstype, fpixel, size, NULL, tmp.ptr(), NULL, &status);
                if(cvtype == CV_16S)
                    tmp.convertTo(tmp, CV_16U, 1, 32767);
                cvimg.push_back(tmp);
            }

            if(cvimg.size() == 1)
            {
                *image = new RawImage(cvimg[0]);
            }
            if(cvimg.size() == 3)
            {
                cv::Mat rgb;
                cv::merge(cvimg, rgb);
                *image = new RawImage(rgb);
            }
        }
    }
    noload:
    if(file)
        loadFITSHeader(file, info);

    fits_close_file(file, &status);
    if(status)
    {
        char err[100];
        fits_get_errstatus(status, err);
        info.info.append({QObject::tr("Error"), QString(err)});
        qDebug() << err;
    }

    return true;
}

#include "pcl/ICCProfileTransformation.h"

template<typename T, typename PCLtype, int CVtype>
bool loadPCLImage(pcl::XISFReader &xisf, RawImage **image)
{
    PCLtype pclImage;
    xisf.ReadImage(pclImage);
    pclImage.Status().DisableInitialization();

#ifdef COLOR_MANAGMENT
    pcl::ICCProfile iccProfile = xisf.ReadICCProfile();
    if(iccProfile.IsProfile())
    {
        pcl::ICCProfileTransformation iccTran;
        iccTran.DisableParallelProcessing();
        iccTran.Add(iccProfile);
        iccTran.Add(sRgbIccProfile);
        iccTran >> pclImage;
    }
#endif

    int numChannels = pclImage.NumberOfChannels();
    cv::Mat cvImg[numChannels];
    for(int i=0; i<numChannels; i++)
    {
        T *p = pclImage.PixelData(i);
        cvImg[i].create(pclImage.Height(), pclImage.Width(), CVtype);
        memcpy(cvImg[i].ptr(0), p, pclImage.Width()*pclImage.Height()*sizeof(T));
    }

    if(numChannels==3)
    {
        cv::Mat merged;
        cv::merge(cvImg, 3, merged);
        *image = new RawImage(merged);
        return true;
    }
    if(numChannels==1)
    {
        *image = new RawImage(cvImg[0]);
        return true;
    }
    return false;
}

bool loadXISF(const QString &path, ImageInfoData &info, RawImage **image)
{
    try
    {
        pcl::XISF::EnsurePTLUTInitialized();
        pcl::String pclPath = path.utf16();
        pcl::XISFReader xisf;
        xisf.Open(pclPath);
        uint8_t bps = xisf.ImageOptions().bitsPerSample;
        bool floatType = xisf.ImageOptions().ieeefpSampleFormat;
        bool signedInt = xisf.ImageOptions().signedIntegers;
        bool complex = xisf.ImageOptions().complexSample;

        auto fitskeywords = xisf.ReadFITSKeywords();
        for(auto fits : fitskeywords)
        {
            info.fitsHeader.append(fits);
        }
        info.wcs = std::make_shared<WCSData>(xisf.ImageInfo().width, xisf.ImageInfo().height, info.fitsHeader);
        if(!info.wcs->valid())info.wcs.reset();

        if(floatType && bps == 32)
            return loadPCLImage<float, pcl::FImage, CV_32F>(xisf, image);
        if(!complex && !signedInt)
        {
            switch(bps)
            {
            case 8:
                return loadPCLImage<uint8_t, pcl::UInt8Image, CV_8U>(xisf, image);
            case 16:
                return loadPCLImage<uint16_t, pcl::UInt16Image, CV_16U>(xisf, image);
            }
        }
    }
    catch (pcl::Error err)
    {
        info.info.append(QPair<QString, QString>("Error", err.FormatInfo().ToUTF8().c_str()));
        qDebug() << err.FormatInfo().ToUTF8().c_str();
    }
    info.info.append({QObject::tr("Error"), QObject::tr("Unsupported sample format")});
    return false;
}

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

    RawImage *rawImage = nullptr;
    bool raw = false;
    if(m_file.endsWith(".CR2", Qt::CaseInsensitive) || m_file.endsWith(".NEF", Qt::CaseInsensitive) || m_file.endsWith(".DNG", Qt::CaseInsensitive))
    {
        timer.start();
        loadRAW(m_file, info, &rawImage);
        raw = true;
        qDebug() << "LoadRaw" << timer.elapsed();
    }
    else if(m_file.endsWith(".FIT", Qt::CaseInsensitive) || m_file.endsWith(".FITS", Qt::CaseInsensitive))
    {
        loadFITS(m_file, info, &rawImage);
    }
    else if(m_file.endsWith(".XISF", Qt::CaseInsensitive))
    {
        loadXISF(m_file, info, &rawImage);
    }
    else
    {
        QImage img(m_file);
#ifdef COLOR_MANAGMENT
        if(img.colorSpace().isValid() && img.colorSpace() != QColorSpace::SRgb)
            img.convertToColorSpace(QColorSpace::SRgb);
#endif

        ExifData *exif = exif_data_new_from_file(m_file.toLocal8Bit().constData());
        info.info.append({QObject::tr("Width"), QString::number(img.width())});
        info.info.append({QObject::tr("Height"), QString::number(img.height())});
        if(exif)
        {
            loadExifEntry(info, exif->ifd[EXIF_IFD_EXIF], EXIF_TAG_ISO_SPEED_RATINGS);
            loadExifEntry(info, exif->ifd[EXIF_IFD_EXIF], EXIF_TAG_SHUTTER_SPEED_VALUE);
            exif_data_free(exif);
        }
        rawImage = new RawImage(img);
    }

    if(rawImage && m_analyzeLevel >= Statistics && !m_thumbnail)
    {
        double mean, median, min, max, mad;
        double stdDev;
        timer.start();
        rawImage->imageStats(&mean, &stdDev, &median, &min, &max, &mad);
        qDebug() << "image stats" << timer.restart();
        info.info.append({QObject::tr("Mean"), QString::number(mean)});
        info.info.append({QObject::tr("Standart deviation"), QString::number(stdDev)});
        info.info.append({QObject::tr("Median"), QString::number(median)});
        info.info.append({QObject::tr("Minimum"), QString::number(min)});
        info.info.append({QObject::tr("Maximum"), QString::number(max)});
        info.info.append({QObject::tr("MAD"), QString::number(mad)});

        if(m_analyzeLevel >= Peaks)
        {
            std::vector<Peak> peaks;
            if(raw) {
                rawImage->quarter();
                qDebug() << "quarter" << timer.restart();
            }
            RawImage *medianImage = rawImage->medianFilter();
            qDebug() << "median" << timer.restart();
            int numPeaks = medianImage->findPeaks(median+stdDev*2, 20, peaks);
            delete medianImage;
            qDebug() << "peaks" << timer.restart();
            //if(m_analyzeLevel == Peaks)
            //    drawPeaks(img, peaks);
            qDebug() << "draw peaks" << timer.restart();
            info.info.append({QObject::tr("Peaks"), QString::number(numPeaks)});
            info.info.append({QObject::tr("Peaks draw"), QString::number(peaks.size())});

            if(m_analyzeLevel>= Stars)
            {
                double fwhmX = 0;
                double fwhmY = 0;
                const int radius = 13;
                StarFit starFit(radius);
                std::vector<Star> stars;
                for(uint i=0; i<peaks.size(); i++)
                {
                    Peak p = peaks[i];
                    std::vector<double> r;
                    int x = p.x();
                    int y = p.y();
                    rawImage->rect(x, y, radius, radius, r);
                    Star star = starFit.fitStar(r, false);
                    if(star.valid())
                    {
                        //printStarModel(radius, r, star);
                        star.m_x += x;
                        star.m_y += y;
                        fwhmX += star.fwhmX();
                        fwhmY += star.fwhmY();
                        stars.push_back(star);
                    }
                }
                //drawStars(img, stars);
                info.info.append({QObject::tr("FWHM X"), QString::number(fwhmX/stars.size())});
                info.info.append({QObject::tr("FWHM Y"), QString::number(fwhmY/stars.size())});
            }
            qDebug() << "Star fit" << timer.restart();
        }
    }

    if(m_thumbnail)
    {
        if(rawImage)
        {
            rawImage->convertToThumbnail();
            QMetaObject::invokeMethod(m_receiver, "thumbnailLoadFinish", Qt::QueuedConnection, Q_ARG(void*, rawImage));
        }
        else
        {
            qDebug() << "failed";
        }
    }
    else
        QMetaObject::invokeMethod(m_receiver, "imageLoaded", Qt::QueuedConnection, Q_ARG(void*, rawImage), Q_ARG(ImageInfoData, info));
}

bool readFITSHeader(const QString &path, ImageInfoData &info)
{
    fitsfile *fr;
    int status = 0;
    fits_open_diskfile(&fr, path.toLocal8Bit().data(), READONLY, &status);

    if(fr && status == 0)
    {
        status = loadFITSHeader(fr, info);
        fits_close_file(fr, &status);
    }
    return status == 0;
}

bool readXISFHeader(const QString &path, ImageInfoData &info)
{
    try
    {
        pcl::String pclPath = path.utf16();
        pcl::XISFReader xisf;
        xisf.Open(pclPath);

        auto fitskeywords = xisf.ReadFITSKeywords();
        for(auto fits : fitskeywords)
        {
            info.fitsHeader.append(fits);
        }
        info.wcs = std::make_shared<WCSData>(xisf.ImageInfo().width, xisf.ImageInfo().height, info.fitsHeader);
        if(!info.wcs->valid())info.wcs.reset();
    }
    catch (pcl::Error err)
    {
        qDebug() << err.FormatInfo().ToUTF8().c_str();
        return false;
    }
    return true;
}

ConvertRunable::ConvertRunable(const QString &in, const QString &out, const QString &format) :
    m_infile(in),
    m_outfile(out),
    m_format(format)
{
}

template<typename T>
void writeXISFImage(pcl::XISFWriter &writer, RawImage *rawimg)
{
    const cv::Mat &cvmat = rawimg->mat();
    T pclimg(rawimg->width(), rawimg->height(), cvmat.channels() == 1 ? pcl::ColorSpace::Gray : pcl::ColorSpace::RGB);
    if(cvmat.channels() == 1)
    {
        memcpy(pclimg.PixelData(0), rawimg->data(), rawimg->size()*sizeof(typename T::sample));
    }
    if(cvmat.channels() == 3)
    {
        std::vector<cv::Mat> channels;
        cv::split(cvmat, channels);
        memcpy(pclimg.PixelData(0), channels[0].data, rawimg->size()*sizeof(typename T::sample));
        memcpy(pclimg.PixelData(1), channels[1].data, rawimg->size()*sizeof(typename T::sample));
        memcpy(pclimg.PixelData(2), channels[2].data, rawimg->size()*sizeof(typename T::sample));
    }
    writer.WriteImage(pclimg);
}

void writeFITSImage(fitsfile *fw, 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->mat().channels();
    int naxis = channels == 1 ? 2 : 3;
    long naxes[3] = {(int)rawimage->width(), (int)rawimage->height(), rawimage->mat().channels()};

    std::vector<cv::Mat> mat;
    if(channels == 1)
        mat.push_back(rawimage->mat());
    else
        cv::split(rawimage->mat(), mat);

    switch(CV_MAT_DEPTH(rawimage->dataType()))
    {
    case CV_8U:
        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(), mat[i].data, &status);
        }
        break;
    case CV_16U:
        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(), mat[i].data, &status);
        }
        break;
    case CV_32F:
        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(), mat[i].data, &status);
        }
        break;
    }
    for(const FITSRecord &record : imageinfo.fitsHeader)
    {
        if(skipKeys.contains(record.key))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()
{
    ImageInfoData imageinfo;
    RawImage *rawimage = nullptr;
    if(m_infile.endsWith(".FITS", Qt::CaseInsensitive) || m_infile.endsWith(".FIT", Qt::CaseInsensitive))
        loadFITS(m_infile, imageinfo, &rawimage);
    if(m_infile.endsWith(".XISF", Qt::CaseInsensitive))
        loadXISF(m_infile, imageinfo, &rawimage);

    if(rawimage)
    {
        if(m_format == "XISF")
        {
            pcl::XISFOptions options;
            pcl::FITSKeywordArray fitskeywords;
            for(auto &record : imageinfo.fitsHeader)
            {
                pcl::FITSHeaderKeyword key(pcl::IsoString(record.key.data()), pcl::IsoString(record.valueToByteArray().data()), pcl::IsoString(record.comment.data()));
                fitskeywords.Append(key);
            }
            pcl::XISFWriter xisf;
            xisf.Create(m_outfile.utf16(), 1);
            xisf.WriteFITSKeywords(fitskeywords);
            switch(CV_MAT_DEPTH(rawimage->dataType()))
            {
            case CV_8U:
                writeXISFImage<pcl::UInt8Image>(xisf, rawimage);
                break;
            case CV_16U:
                writeXISFImage<pcl::UInt16Image>(xisf, rawimage);
                break;
            case CV_32F:
                writeXISFImage<pcl::Image>(xisf, rawimage);
                break;
            }

        }

        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);
            writeFITSImage(fw, rawimage, imageinfo);
            fits_close_file(fw, &status);
        }
        delete rawimage;
    }
}