#include "NDData.hpp"

#include "MayaFlux/Transitive/Reflect/EnumReflect.hpp"

#include "MayaFlux/Journal/Archivist.hpp"

namespace MayaFlux::Kakshya {

DataDimension::DataDimension(std::string n, uint64_t s, uint64_t st, Role r)
    : name(std::move(n))
    , size(s)
    , stride(st)
    , role(r)
{
}

DataDimension DataDimension::time(uint64_t samples, std::string name)
{
    return { std::move(name), samples, 1, Role::TIME };
}

DataDimension DataDimension::channel(uint64_t count, uint64_t stride)
{
    return { "channel", count, stride, Role::CHANNEL };
}

DataDimension DataDimension::frequency(uint64_t bins, std::string name)
{
    return { std::move(name), bins, 1, Role::FREQUENCY };
}

DataDimension DataDimension::spatial(uint64_t size, char axis, uint64_t stride, std::string name)
{
    Role r = [axis]() {
        switch (axis) {
        case 'x':
            return Role::SPATIAL_X;
        case 'y':
            return Role::SPATIAL_Y;
        case 'z':
        default:
            return Role::SPATIAL_Z;
        }
    }();

    if (name == "") {
        name = std::string("spatial_") + axis;
    }
    return { name, size, stride, r };
}

DataDimension DataDimension::spatial_1d(uint64_t width)
{
    return spatial(width, 'x');
}

DataDimension DataDimension::spatial_2d(uint64_t width, uint64_t height)
{
    return spatial(height, 'y', width * sizeof(float));
}

DataDimension DataDimension::spatial_3d(uint64_t width, uint64_t height, uint64_t depth)
{
    return spatial(depth, 'z', width * height * sizeof(float));
}

DataDimension DataDimension::mipmap_levels(uint64_t levels)
{
    return { "mip_levels", levels, 1, Role::MIP_LEVEL };
}

DataDimension DataDimension::grouped(
    std::string name,
    uint64_t element_count,
    uint8_t components_per_element,
    Role role)
{
    DataDimension dim { std::move(name), element_count, 1, role };
    dim.grouping = ComponentGroup { components_per_element, 0 };
    return dim;
}

std::string_view modality_to_string(DataModality modality)
{
    return Reflect::enum_to_string(modality);
}

DataDimension DataDimension::vertex_positions(uint64_t count)
{
    return grouped("positions", count, 3, Role::POSITION);
}

DataDimension DataDimension::vertex_normals(uint64_t count)
{
    return grouped("normals", count, 3, Role::NORMAL);
}

DataDimension DataDimension::texture_coords(uint64_t count)
{
    return grouped("uvs", count, 2, Role::UV);
}

DataDimension DataDimension::vertex_colors(uint64_t count, bool has_alpha)
{
    return grouped("colors", count, has_alpha ? 4 : 3, Role::COLOR);
}

std::vector<DataDimension> DataDimension::create_dimensions(
    DataModality modality,
    const std::vector<uint64_t>& shape,
    MemoryLayout layout)
{
    std::vector<DataDimension> dims;
    auto strides = calculate_strides(shape, layout);

    switch (modality) {
    case DataModality::AUDIO_1D:
        if (shape.size() != 1) {
            error<std::invalid_argument>(
                Journal::Component::Kakshya,
                Journal::Context::Runtime,
                std::source_location::current(),
                "AUDIO_1D requires 1D shape");
        }
        dims.push_back(DataDimension::time(shape[0]));
        break;

    case DataModality::AUDIO_MULTICHANNEL:
        if (shape.size() != 2) {
            error<std::invalid_argument>(
                Journal::Component::Kakshya,
                Journal::Context::Runtime,
                std::source_location::current(),
                "AUDIO_MULTICHANNEL requires 2D shape [samples, channels]");
        }
        dims.push_back(DataDimension::time(shape[0]));
        dims.push_back(DataDimension::channel(shape[1], strides[1]));
        break;

    case DataModality::IMAGE_2D:
        if (shape.size() != 2) {
            error<std::invalid_argument>(
                Journal::Component::Kakshya,
                Journal::Context::Runtime,
                std::source_location::current(),
                "IMAGE_2D requires 2D shape [height, width]");
        }
        dims.push_back(DataDimension::spatial(shape[0], 'y', strides[0]));
        dims.push_back(DataDimension::spatial(shape[1], 'x', strides[1]));
        break;

    case DataModality::IMAGE_COLOR:
        if (shape.size() != 3) {
            error<std::invalid_argument>(
                Journal::Component::Kakshya,
                Journal::Context::Runtime,
                std::source_location::current(),
                "IMAGE_COLOR requires 3D shape [height, width, channels]");
        }
        dims.push_back(DataDimension::spatial(shape[0], 'y', strides[0]));
        dims.push_back(DataDimension::spatial(shape[1], 'x', strides[1]));
        dims.push_back(DataDimension::channel(shape[2], strides[2]));
        break;

    case DataModality::SPECTRAL_2D:
        if (shape.size() != 2) {
            error<std::invalid_argument>(
                Journal::Component::Kakshya,
                Journal::Context::Runtime,
                std::source_location::current(),
                "SPECTRAL_2D requires 2D shape [time_windows, frequency_bins]");
        }
        dims.push_back(DataDimension::time(shape[0], "time_windows"));
        dims.push_back(DataDimension::frequency(shape[1]));
        dims[1].stride = strides[1];
        break;

    case DataModality::VOLUMETRIC_3D:
        if (shape.size() != 3) {
            error<std::invalid_argument>(
                Journal::Component::Kakshya,
                Journal::Context::Runtime,
                std::source_location::current(),
                "VOLUMETRIC_3D requires 3D shape [x, y, z]");
        }
        dims.push_back(DataDimension::spatial(shape[0], 'x', strides[0]));
        dims.push_back(DataDimension::spatial(shape[1], 'y', strides[1]));
        dims.push_back(DataDimension::spatial(shape[2], 'z', strides[2]));
        break;

    case DataModality::VIDEO_GRAYSCALE:
        if (shape.size() != 3) {
            error<std::invalid_argument>(
                Journal::Component::Kakshya,
                Journal::Context::Runtime,
                std::source_location::current(),
                "VIDEO_GRAYSCALE requires 3D shape [frames, height, width]");
        }
        dims.push_back(DataDimension::time(shape[0], "frames"));
        dims.push_back(DataDimension::spatial(shape[1], 'y', strides[1]));
        dims.push_back(DataDimension::spatial(shape[2], 'x', strides[2]));
        break;

    case DataModality::VIDEO_COLOR:
        if (shape.size() != 4) {
            error<std::invalid_argument>(
                Journal::Component::Kakshya,
                Journal::Context::Runtime,
                std::source_location::current(),
                "VIDEO_COLOR requires 4D shape [frames, height, width, channels]");
        }
        dims.push_back(DataDimension::time(shape[0], "frames"));
        dims.push_back(DataDimension::spatial(shape[1], 'y', strides[1]));
        dims.push_back(DataDimension::spatial(shape[2], 'x', strides[2]));
        dims.push_back(DataDimension::channel(shape[3], strides[3]));
        break;

    default:
        error<std::invalid_argument>(
            Journal::Component::Kakshya,
            Journal::Context::Runtime,
            std::source_location::current(),
            "Unsupported modality for dimension creation: {}", modality_to_string(modality));
    }

    return dims;
}

std::vector<uint64_t> DataDimension::calculate_strides(
    const std::vector<uint64_t>& shape,
    MemoryLayout layout)
{
    if (shape.empty())
        return {};

    std::vector<uint64_t> strides(shape.size());

    if (layout == MemoryLayout::ROW_MAJOR) {
        auto reversed_shape = shape | std::views::reverse;
        std::exclusive_scan(
            reversed_shape.begin(),
            reversed_shape.end(),
            strides.rbegin(),
            uint64_t(1),
            std::multiplies<uint64_t> {});
    } else {
        std::exclusive_scan(
            shape.begin(),
            shape.end(),
            strides.begin(),
            uint64_t(1),
            std::multiplies<uint64_t> {});
    }

    return strides;
}

}
