#pragma once
#include <algorithm>
#include <fstream>
#include <memory>
#include <sstream>
#include <string>
#include <vector>

#include <Volume.h>

class Loader {
public:
    virtual ~Loader() = default;

    virtual std::unique_ptr<Volume> loadFile(std::string filename) = 0;
};

class MHDLoader : public Loader {
public:
    std::unique_ptr<Volume> loadFile(std::string filename) override {
        auto pos = filename.find_last_of('/');
        fileName = filename.substr(0, pos + 1);

        std::string line;
        std::ifstream file(filename);
        if (file.is_open()) {
            while (getline(file, line)) {
                evaluateLine(line);
            }
            file.close();
        } else {
            std::cerr << "Error: File not found!" << std::endl;
            return nullptr;
        }

        if (bigEndian) {
            std::cerr << "Error: Big Endian not handled!" << std::endl;
            return nullptr;
        }

        auto bytes = ReadAllBytes(fileName.c_str());
        if (bytes.empty()) {
            std::cerr << "Error: could not read raw file!" << std::endl;
            return nullptr;
        }

        if (type == "MET_UCHAR") {
            std::vector<unsigned char> data = loadVolumeUChar(bytes);
            data = rotateArray(data);
            return std::make_unique<Volume>(dimX, dimY, dimZ, spacing, data.data());
        }
        if (type == "MET_USHORT") {
            std::vector<unsigned short> data = loadVolumeUShort(bytes);
            data = rotateArray(data);
            return std::make_unique<Volume>(dimX, dimY, dimZ, spacing, data.data());
        }
        if (type == "MET_SHORT") {
            std::vector<unsigned short> data = loadVolumeShort(bytes);
            data = rotateArray(data);
            return std::make_unique<Volume>(dimX, dimY, dimZ, spacing, data.data());
        }
        if (type == "MET_FLOAT") {
            std::vector<float> data = loadVolumeFloat(bytes);
            data = rotateArray(data);
            return std::make_unique<Volume>(dimX, dimY, dimZ, spacing, data.data());
        }
        std::cerr << "Error: Format not handled" << std::endl;
        return nullptr;
    }

private:
    template<typename T>
    std::vector<T> rotateArray(const std::vector<T>& array) {
        std::vector<T> result = std::vector<T>((long) dimX * (long) dimY * (long) dimZ);
        float3 vec = matrix * make_float3(dimX, dimY, dimZ);
        float3 newSpacing = matrix * spacing;

        long newDimX = abs(vec.x);
        long newDimY = abs(vec.y);
        long newDimZ = abs(vec.z);

#pragma omp parallel for
        for (int z = 0; z < dimZ; ++z) {
            for (int y = 0; y < dimY; ++y) {
                for (int x = 0; x < dimX; ++x) {
                    float3 pos = make_float3((float) x / (float) (dimX - 1), (float) y / (float) (dimY - 1),
                                             (float) z / (float) (dimZ - 1));
                    pos = (pos * 2.f) - 1.f;

                    float3 rotatedPosition = ((matrix * pos) + 1.f) / 2.f;

                    long newX = static_cast<long>(round(rotatedPosition.x * (newDimX - 1)));
                    long newY = static_cast<long>(round(rotatedPosition.y * (newDimY - 1)));
                    long newZ = static_cast<long>(round(rotatedPosition.z * (newDimZ - 1)));

                    result[((newZ * newDimY) + newY) * newDimX + newX] =
                            array[((z * (long) dimY) + y) * (long) dimX + x];
                }
            }
        }

        dimX = newDimX;
        dimY = newDimY;
        dimZ = newDimZ;
        spacing = abs(newSpacing);
        return result;
    }

    std::vector<char> ReadAllBytes(char const* filename) {
        std::ifstream ifs(filename, std::ios::binary | std::ios::ate);
        std::ifstream::pos_type pos = ifs.tellg();

        if (pos <= 0) {
            return std::vector<char>{};
        }

        std::vector<char> result(pos);

        ifs.seekg(0, std::ios::beg);
        ifs.read(&result[0], pos);

        return result;
    }

    void evaluateLine(const std::string& line) {
        std::stringstream ss(line);
        std::string word;
        std::string equals;
        ss >> word;
        ss >> equals;
        if (word == "NDims") {
            std::string temp;
            ss >> temp;
            dims = std::stoi(temp);
            return;
        }
        if (word == "DimSize") {
            std::string temp;
            ss >> temp;
            dimX = std::stoi(temp);
            ss >> temp;
            dimY = std::stoi(temp);
            ss >> temp;
            dimZ = std::stoi(temp);
            return;
        }
        if (word == "BinaryDataByteOrderMSB") {
            std::string temp;
            ss >> temp;
            bigEndian = temp == "True";
            return;
        }
        if (word == "ElementType") {
            std::string temp;
            ss >> temp;
            type = temp;
            return;
        }
        if (word == "ElementDataFile") {
            std::string temp;
            ss >> temp;
            fileName += temp;
            return;
        }
        if (word == "ElementSpacing") {
            std::string temp;
            ss >> temp;
            float spaceX = std::stof(temp);
            ss >> temp;
            float spaceY = std::stof(temp);
            ss >> temp;
            float spaceZ = std::stof(temp);

            spacing = make_float3(spaceX, spaceY, spaceZ);
            return;
        }
        if (word == "TransformMatrix") {
            std::string temp;
            ss >> temp;
            float XX = std::stof(temp);
            ss >> temp;
            float XY = std::stof(temp);
            ss >> temp;
            float XZ = std::stof(temp);
            ss >> temp;
            float YX = std::stof(temp);
            ss >> temp;
            float YY = std::stof(temp);
            ss >> temp;
            float YZ = std::stof(temp);
            ss >> temp;
            float ZX = std::stof(temp);
            ss >> temp;
            float ZY = std::stof(temp);
            ss >> temp;
            float ZZ = std::stof(temp);

            matrix = make_mat3(make_float3(XX, XY, XZ), make_float3(YX, YY, YZ), make_float3(ZX, ZY, ZZ));
            return;
        }
    }

    [[nodiscard]] std::vector<unsigned char> loadVolumeUChar(const std::vector<char>& bytes) {
        size_t size = bytes.size();
        std::vector<unsigned char> data = std::vector<unsigned char>(size);
        for (size_t i = 0; i < size; ++i) {
            auto temp = static_cast<unsigned char>(bytes[i]);
            data[i] = temp;
        }

        float max = *std::ranges::max_element(data);

        for (size_t i = 0; i < size; ++i) {
            unsigned char val = data[i];
            float temp = (float) val / max;
            unsigned char value = (unsigned char) (temp * 255.f);
            data[i] = value;
        }

        return data;
    }

    [[nodiscard]] std::vector<unsigned short> loadVolumeUShort(const std::vector<char>& bytes) {
        size_t size = bytes.size() / 2;
        std::vector<unsigned short> data = std::vector<unsigned short>(size);
        for (size_t i = 0; i < size; ++i) {
            auto first = static_cast<unsigned char>(bytes[i * 2]);
            auto second = static_cast<unsigned char>(bytes[i * 2 + 1]);
            unsigned short temp = second << 8 | first;
            data[i] = temp;
        }

        float max = *std::ranges::max_element(data);
        for (size_t i = 0; i < size; ++i) {
            unsigned short val = data[i];
            float temp = (float) val / max;
            unsigned short value = (unsigned short) (temp * 65535.f);
            data[i] = value;
        }

        return data;
    }

    [[nodiscard]] std::vector<unsigned short> loadVolumeShort(const std::vector<char>& bytes) {
        unsigned long size = bytes.size() / 2;
        std::vector<short> data = std::vector<short>(size);
        for (size_t i = 0; i < size; ++i) {
            auto first = static_cast<uint8_t>(bytes[i * 2]);
            auto second = static_cast<uint8_t>(bytes[i * 2 + 1]);
            uint16_t temp = second << 8 | first;
            auto s = static_cast<short>(temp);
            data[i] = s;
        }

        short min = *std::ranges::min_element(data);
        short max = *std::ranges::max_element(data);
        max -= min;
        std::vector<unsigned short> values = std::vector<unsigned short>(size);
        for (size_t i = 0; i < size; ++i) {
            short val = data[i] - min;
            float scaled = val / max;
            unsigned short value = (unsigned short) (scaled * 65535.f);
            values[i] = value;
        }

        return values;
    }

    [[nodiscard]] std::vector<float> loadVolumeFloat(std::vector<char> bytes) {
        size_t size = bytes.size() / 4;
        std::vector<float> data = std::vector<float>(size);
        for (size_t i = 0; i < size; ++i) {
            auto first = static_cast<unsigned char>(bytes[i * 4]);
            auto second = static_cast<unsigned char>(bytes[i * 4 + 1]);
            auto third = static_cast<unsigned char>(bytes[i * 4 + 2]);
            auto fourth = static_cast<unsigned char>(bytes[i * 4 + 3]);
            uint32_t temp = fourth << 24 | third << 16 | second << 8 | first;
            float f = 0.0f;
            memcpy(&f, &temp, sizeof(temp));
            data[i] = f;
        }

        bytes.clear();
        bytes.shrink_to_fit();
        float max = *std::ranges::max_element(data);
        float min = *std::ranges::min_element(data);
        max -= min;

        std::vector<float> result = std::vector<float>(size);
        for (size_t i = 0; i < size; ++i) {
            float val = (data[i] - min) / max;
            result[i] = val;
        }

        return result;
    }

    int dims = 0;
    bool bigEndian = false;
    int dimX = 0;
    int dimY = 0;
    int dimZ = 0;
    std::string type;
    std::string fileName;

    float3 spacing = make_float3(1);
    mat3 matrix = make_mat3_identity();
};
