// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#ifndef JUNGFRAUJOCH_RAWTOCONVERTEDGEOMETRY_H
#define JUNGFRAUJOCH_RAWTOCONVERTEDGEOMETRY_H

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

// Take half of the number, but only if not bad pixel/overload
template<typename T>
T half(T in, T min, T max) {
    T tmp = in;
    if ((in > min) && (in < max)) tmp /= 2;
    return tmp;
}

template<typename T>
T quarter(T in, T min, T max) {
    T tmp = in;
    if ((in > min) && (in < max)) tmp /= 4;
    return tmp;
}

// Copy line, divide everything by 2 and extend multipixels + divide them by additional factor of 2
template<typename Td, typename Ts>
void LineCopyAndAdjustMultipixelMidRow(Td *destination, const Ts *source, Ts min, Ts max,
                                       int64_t fast_dir_step = 1, int64_t offset_0 = 0) {
    for (int chip = 0; chip < 4; chip++) {
        for (int i = 0; i < 256; i++) {
            destination[offset_0 + (i + chip * 258) * fast_dir_step] = half(source[i + chip * 256], min, max);
        }
    }
    for (int i = 0; i < 3; i++) {
        destination[offset_0 + (255 + i * 258) * fast_dir_step] = quarter(source[255 + i * 256], min, max);
        destination[offset_0 + (256 + i * 258) * fast_dir_step] = quarter(source[255 + i * 256], min, max);
        destination[offset_0 + (257 + i * 258) * fast_dir_step] = quarter(source[256 + i * 256], min, max);
        destination[offset_0 + (258 + i * 258) * fast_dir_step] = quarter(source[256 + i * 256], min, max);
    }
}

// Copy line and extend multipixels + divide them by 2
template<typename Td, typename Ts>
void LineCopyAndAdjustMultipixel(Td *destination, const Ts *source, Ts min, Ts max,
                                 int64_t fast_dir_step = 1, int64_t offset_0 = 0) {
    for (int chip = 0; chip < 4; chip++) {
        for (int i = 0; i < 256; i++) {
            destination[offset_0 + (i + chip * 258) * fast_dir_step] = source[i + chip * 256];
        }
    }
    for (int chip = 0; chip < 3; chip++) {
        destination[offset_0 + (255 + chip * 258) * fast_dir_step] = half(source[255 + chip * 256], min, max);
        destination[offset_0 + (256 + chip * 258) * fast_dir_step] = half(source[255 + chip * 256], min, max);
        destination[offset_0 + (257 + chip * 258) * fast_dir_step] = half(source[256 + chip * 256], min, max);
        destination[offset_0 + (258 + chip * 258) * fast_dir_step] = half(source[256 + chip * 256], min, max);
    }
}

// Copy line and copy multipixels (e.g. for mask)
template<typename Td, typename Ts>
void LineCopyAndAddMultipixel(Td *destination, const Ts *source,
                              int64_t fast_dir_step = 1, int64_t offset_0 = 0) {
    for (int chip = 0; chip < 4; chip++) {
        for (int i = 0; i < 256; i++)
            destination[offset_0 + (i + chip * 258) * fast_dir_step] = source[i + chip * 256];
    }
    for (int i = 0; i < 3; i++) {
        destination[offset_0 + (256 + i * 258) * fast_dir_step] = source[255 + i * 256];
        destination[offset_0 + (257 + i * 258) * fast_dir_step] = source[256 + i * 256];
    }
}


template<class T, class Tint = int32_t>
T Bin2x2_sum(T a, T b, T c, T d, T underload, T overload) {
    T ret;
    if ((a <= underload) || (b <= underload) || (c <= underload) || (d <= underload))
        ret = underload;
    else if ((a >= overload) || (b >= overload) || (c >= overload) || (d >= overload))
        ret = overload;
    else {
        Tint sum = a + b + c + d;
        if (sum > overload)
            ret = overload;
        else
            ret = static_cast<T>(sum);
    }
    return ret;
}

template<class T, class Tint = int32_t>
void Bin2x2_sum(T *destination, const T *source, size_t width, size_t height, T underload, T overload) {
    for (int y = 0; y < height / 2; y++) {
        for (int x = 0; x < width / 2; x++)
            destination[y * (width / 2) + x] = Bin2x2_sum<T, Tint>(source[(y * 2) * width + (x * 2)],
                                                                   source[(y * 2 + 1) * width + (x * 2)],
                                                                   source[(y * 2) * width + (x * 2 + 1)],
                                                                   source[(y * 2 + 1) * width + (x * 2 + 1)],
                                                                   underload, overload);
    }
}

template<class Td, class Ts>
void TransferModule(Td *destination, const Ts *source,
                    int64_t slow_dir_step,
                    int64_t fast_dir_step = 1,
 int64_t offset_0 = 0) {
    for (size_t line = 0; line < RAW_MODULE_LINES; line++) {
        if ((line == 255) || (line == 256)) {
            LineCopyAndAddMultipixel<Td, Ts>(destination + slow_dir_step * (line + 1),
                                             source + RAW_MODULE_COLS * line,
                                             fast_dir_step, offset_0);
            LineCopyAndAddMultipixel<Td, Ts>(destination + slow_dir_step * (line + ((line > 255) ? 2 : 0)),
                                             source + RAW_MODULE_COLS * line,
                                             fast_dir_step, offset_0);
        } else {
            LineCopyAndAddMultipixel<Td, Ts>(destination + slow_dir_step * (line + ((line > 255) ? 2 : 0)),
                                             source + RAW_MODULE_COLS * line,
                                             fast_dir_step);
        }
    }
}

template<class Td, class Ts>
void TransferModuleAdjustMultipixels(Td *destination, const Ts *source, int64_t slow_dir_step, Ts min, Ts max,
                                     int64_t fast_dir_step = 1, int64_t offset_0 = 0) {
    for (size_t line = 0; line < RAW_MODULE_LINES; line++) {
        if ((line != 255) && (line != 256))
            LineCopyAndAdjustMultipixel<Td, Ts>(destination + slow_dir_step * (line + ((line > 255) ? 2 : 0)),
                                                source + RAW_MODULE_COLS * line, min, max, fast_dir_step,
                                                offset_0);
        else {
            LineCopyAndAdjustMultipixelMidRow<Td, Ts>(destination + slow_dir_step * (line + ((line > 255) ? 2 : 0)),
                                                      source + RAW_MODULE_COLS * line, min, max,
                                                      fast_dir_step, offset_0);
            LineCopyAndAdjustMultipixelMidRow<Td, Ts>(destination + slow_dir_step * (line + 1),
                                                      source + RAW_MODULE_COLS * line, min, max,
                                                      fast_dir_step, offset_0);
        }
    }
}

// Input coord is column + 1024 * (line + 512 * module)
// Copies result over multipixel - can be used for example for mask calculation
template<class Td, class Ts>
void RawToConvertedGeometry(const DiffractionExperiment &experiment, Td *destination, const Ts *source) {
    for (size_t module_number = 0; module_number < experiment.GetModulesNum(); module_number++)
        TransferModule<Td, Ts>(destination + experiment.GetPixel0OfModule(module_number),
                               source + module_number * RAW_MODULE_SIZE,
                               experiment.GetModuleSlowDirectionStep(module_number),
                               experiment.GetModuleFastDirectionStep(module_number));
}

template<class T>
void LineConvtToRaw(T *destination, const T *source) {
    for (int chip = 0; chip < 4; chip++) {
        for (int i = 0; i < 256; i++)
            destination[i + chip * 256] = source[i + chip * 258];
    }
}

template<class T>
void ConvertedToRawGeometry(const DiffractionExperiment &experiment, T *destination, const T *source) {
    for (size_t module_number = 0; module_number < experiment.GetModulesNum(); module_number++) {
        for (size_t line = 0; line < RAW_MODULE_LINES; line++) {
            LineConvtToRaw<T>(destination + module_number * RAW_MODULE_SIZE + RAW_MODULE_COLS * line,
                              source + experiment.GetPixel0OfModule(module_number) +
                              experiment.GetModuleSlowDirectionStep(module_number) *
                              (line + ((line > 255) ? 2 : 0))
            );
        }
    }
}

template<class T>
void Bin2x2_or(T *destination, const T *source, size_t width, size_t height) {
    for (int y = 0; y < height / 2; y++) {
        for (int x = 0; x < width / 2; x++) {
            T tmp[4];
            tmp[0] = source[(y * 2) * width + (x * 2)];
            tmp[1] = source[(y * 2 + 1) * width + (x * 2)];
            tmp[2] = source[(y * 2) * width + (x * 2 + 1)];
            tmp[3] = source[(y * 2 + 1) * width + (x * 2 + 1)];
            destination[y * (width / 2) + x] = tmp[0] | tmp[1] | tmp[2] | tmp[3];
        }
    }
}

// Input coord is column + 1024 * (line + 512 * module)
template<class Ts>
void RawToConvertedGeometryAdjustMultipixels(const DiffractionExperiment &experiment, Ts *destination, const Ts *source) {
    Ts min = experiment.GetUnderflow() + 1;
    if (min > 0) min = 0;
    Ts max = experiment.GetOverflow() - 1;

    if (experiment.GetBinning2x2()) {
        std::vector<Ts> tmp(experiment.GetPixelsNumFullImage());
        for (size_t module_number = 0; module_number < experiment.GetModulesNum(); module_number++) {
            TransferModuleAdjustMultipixels<Ts, Ts>(tmp.data(),
                                                    source + module_number * RAW_MODULE_SIZE,
                                                    experiment.GetModuleSlowDirectionStep(module_number),
                                                    min, max,
                                                    experiment.GetModuleFastDirectionStep(module_number),
                                                    experiment.GetPixel0OfModule(module_number));
            Bin2x2_sum(destination,
                       tmp.data(),
                       experiment.GetXPixelsNumFullImage(),
                       experiment.GetYPixelsNumFullImage(),
                       min, max);
        }
    } else {
        for (size_t module_number = 0; module_number < experiment.GetModulesNum(); module_number++) {
            TransferModuleAdjustMultipixels<Ts, Ts>(destination,
                                                    source + module_number * RAW_MODULE_SIZE,
                                                    experiment.GetModuleSlowDirectionStep(module_number),
                                                    min, max,
                                                    experiment.GetModuleFastDirectionStep(module_number),
                                                    experiment.GetPixel0OfModule(module_number));
        }
    }
}

#endif