#ifndef __SD_MODEL_VAE_VAE_HPP__
#define __SD_MODEL_VAE_VAE_HPP__

#include "core/tensor_ggml.hpp"
#include "model/common/block.hpp"
#include "model/vae/vae_tiling.hpp"
#include "model_manager.h"

struct VAE : public GGMLRunner {
protected:
    SDVersion version;
    std::string weight_prefix;
    bool scale_input                                      = true;
    virtual sd::Tensor<float> _compute(const int n_threads,
                                       const sd::Tensor<float>& z,
                                       bool decode_graph) = 0;

    virtual bool supports_temporal_tiling(VAETemporalDirection direction) const {
        SD_UNUSED(direction);
        return false;
    }

    virtual int get_default_temporal_tile_frames(VAETemporalDirection direction) const {
        SD_UNUSED(direction);
        return 4;
    }

    virtual int get_default_temporal_tile_overlap(VAETemporalDirection direction) const {
        SD_UNUSED(direction);
        return 1;
    }

    virtual int get_temporal_tile_output_scale(VAETemporalDirection direction) const {
        SD_UNUSED(direction);
        return 1;
    }

    virtual sd::Tensor<float> _compute_temporal_tiled(const int n_threads,
                                                      const sd::Tensor<float>& input,
                                                      VAETemporalDirection direction,
                                                      const VAETemporalTilingConfig& config) {
        if (direction != VAETemporalDirection::DECODE) {
            return _compute(n_threads, input, false);
        }

        VAETemporalTilingConfig resolved_config = config;
        const int output_scale                  = get_temporal_tile_output_scale(direction);
        if (output_scale > 1 &&
            resolved_config.overlap == 0 &&
            input.shape()[2] > resolved_config.tile_frames) {
            LOG_WARN("%s temporal decode requires at least one overlapping latent frame; using overlap=1",
                     get_desc().c_str());
            resolved_config.overlap = 1;
        }

        auto plan = make_vae_temporal_tile_plan(input.shape()[2], resolved_config);
        LOG_DEBUG("%s temporal tiling: tile_frames=%d, overlap=%d, total_frames=%lld, tiles=%d",
                  get_desc().c_str(),
                  plan.tile_frames,
                  plan.overlap,
                  (long long)input.shape()[2],
                  (int)plan.tiles.size());
        return process_vae_temporal_tiles_blended(
            input,
            plan,
            output_scale,
            [&](const sd::Tensor<float>& input_tile, const VAETemporalTile& tile) {
                LOG_DEBUG("%s temporal tile %d/%d: input frames [%lld, %lld)",
                          get_desc().c_str(),
                          tile.index + 1,
                          (int)plan.tiles.size(),
                          (long long)tile.start,
                          (long long)tile.end);
                return _compute(n_threads, input_tile, true);
            });
    }

    sd::Tensor<float> compute_with_temporal_tiling(const int n_threads,
                                                   const sd::Tensor<float>& input,
                                                   VAETemporalDirection direction,
                                                   const sd_tiling_params_t& tiling_params) {
        if (!tiling_params.temporal_tiling || input.dim() != 5 || input.shape()[2] <= 1) {
            return _compute(n_threads, input, direction == VAETemporalDirection::DECODE);
        }
        if (!supports_temporal_tiling(direction)) {
            LOG_WARN("%s does not support temporal tiling for %s; processing the full temporal dimension",
                     get_desc().c_str(),
                     direction == VAETemporalDirection::DECODE ? "decode" : "encode");
            return _compute(n_threads, input, direction == VAETemporalDirection::DECODE);
        }

        auto config = resolve_vae_temporal_tiling_config(
            tiling_params,
            get_default_temporal_tile_frames(direction),
            get_default_temporal_tile_overlap(direction));
        return _compute_temporal_tiled(n_threads, input, direction, config);
    }

    static inline void scale_tensor_to_minus1_1(sd::Tensor<float>* tensor) {
        GGML_ASSERT(tensor != nullptr);
        for (int64_t i = 0; i < tensor->numel(); ++i) {
            (*tensor)[i] = (*tensor)[i] * 2.0f - 1.0f;
        }
    }

    static inline void scale_tensor_to_0_1(sd::Tensor<float>* tensor) {
        GGML_ASSERT(tensor != nullptr);
        for (int64_t i = 0; i < tensor->numel(); ++i) {
            float value  = ((*tensor)[i] + 1.0f) * 0.5f;
            (*tensor)[i] = std::max(0.0f, std::min(1.0f, value));
        }
    }

    sd::Tensor<float> tiled_compute(const sd::Tensor<float>& input,
                                    int n_threads,
                                    int output_width,
                                    int output_height,
                                    int scale,
                                    int p_tile_size_x,
                                    int p_tile_size_y,
                                    float tile_overlap_factor,
                                    bool circular_x,
                                    bool circular_y,
                                    bool decode_graph,
                                    const sd_tiling_params_t& tiling_params,
                                    const char* error_message,
                                    bool silent = false) {
        auto on_processing = [&](const sd::Tensor<float>& input_tile) {
            auto output_tile = compute_with_temporal_tiling(
                n_threads,
                input_tile,
                decode_graph ? VAETemporalDirection::DECODE : VAETemporalDirection::ENCODE,
                tiling_params);
            if (output_tile.empty()) {
                LOG_ERROR("%s", error_message);
                return sd::Tensor<float>();
            }
            return output_tile;
        };
        return ::process_tiles_2d(input,
                                  output_width,
                                  output_height,
                                  scale,
                                  p_tile_size_x,
                                  p_tile_size_y,
                                  tile_overlap_factor,
                                  circular_x,
                                  circular_y,
                                  on_processing,
                                  silent);
    }

public:
    VAE(SDVersion version,
        ggml_backend_t backend,
        const std::string& weight_prefix                    = "",
        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
        : version(version), weight_prefix(weight_prefix), GGMLRunner(backend, weight_manager) {}

    int get_scale_factor() {
        int scale_factor = 8;
        if (version == VERSION_LTXAV) {
            scale_factor = 32;
        } else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
            scale_factor = 16;
        } else if (sd_version_uses_flux2_vae(version)) {
            scale_factor = 16;
        } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
            scale_factor = 1;
        }
        return scale_factor;
    }

    virtual int get_encoder_output_channels(int input_channels) = 0;

    bool can_temporal_tile_decode() const {
        return supports_temporal_tiling(VAETemporalDirection::DECODE);
    }

    void get_tile_sizes(int& tile_size_x,
                        int& tile_size_y,
                        float& tile_overlap,
                        const sd_tiling_params_t& params,
                        int64_t latent_x,
                        int64_t latent_y,
                        float encoding_factor = 1.0f) {
        tile_overlap       = std::max(std::min(params.target_overlap, 0.5f), 0.0f);
        auto get_tile_size = [&](int requested_size, float factor, int64_t latent_size) {
            const int default_tile_size  = 32;
            const int min_tile_dimension = 4;
            int tile_size                = default_tile_size;
            // factor <= 1 means simple fraction of the latent dimension
            // factor > 1 means number of tiles across that dimension
            if (factor > 0.f) {
                if (factor > 1.0)
                    factor = 1 / (factor - factor * tile_overlap + tile_overlap);
                tile_size = static_cast<int>(std::round(latent_size * factor));
            } else if (requested_size >= min_tile_dimension) {
                tile_size = requested_size;
            }
            tile_size = static_cast<int>(tile_size * encoding_factor);
            return std::max(std::min(tile_size, static_cast<int>(latent_size)), min_tile_dimension);
        };

        tile_size_x = get_tile_size(params.tile_size_x, params.rel_size_x, latent_x);
        tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
    }

    virtual sd::Tensor<float> encode(int n_threads,
                                     const sd::Tensor<float>& x,
                                     sd_tiling_params_t tiling_params,
                                     bool circular_x = false,
                                     bool circular_y = false) {
        int64_t t0              = ggml_time_ms();
        sd::Tensor<float> input = x;
        sd::Tensor<float> output;
        if (scale_input) {
            scale_tensor_to_minus1_1(&input);
        }

        if (tiling_params.enabled) {
            const int scale_factor = get_scale_factor();
            int64_t W              = input.shape()[0] / scale_factor;
            int64_t H              = input.shape()[1] / scale_factor;
            float tile_overlap;
            int tile_size_x, tile_size_y;
            // Image VAE encode is more sensitive to tile boundary context than decode.
            // Keep the smaller legacy factor for video VAEs, but default image encode
            // tiles to 64 latent pixels so a 512px SD image is encoded as one tile.
            const float encode_tile_factor = sd_version_is_minimax_h3(version) ? 1.f : (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f
                                                                                                                                                                                            : 2.0f;
            get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor);
            LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
            output = tiled_compute(input,
                                   n_threads,
                                   static_cast<int>(W),
                                   static_cast<int>(H),
                                   scale_factor,
                                   tile_size_x,
                                   tile_size_y,
                                   tile_overlap,
                                   circular_x,
                                   circular_y,
                                   false,
                                   tiling_params,
                                   "vae encode compute failed while processing a tile");
        } else {
            output = compute_with_temporal_tiling(n_threads,
                                                  input,
                                                  VAETemporalDirection::ENCODE,
                                                  tiling_params);
        }

        runner_end();

        if (output.empty()) {
            LOG_ERROR("vae encode compute failed");
            return {};
        }
        int64_t t1 = ggml_time_ms();
        LOG_DEBUG("computing vae encode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
        return std::move(output);
    }

    virtual sd::Tensor<float> decode(int n_threads,
                                     const sd::Tensor<float>& x,
                                     sd_tiling_params_t tiling_params,
                                     bool decode_video = false,
                                     bool circular_x   = false,
                                     bool circular_y   = false,
                                     bool silent       = false) {
        int64_t t0              = ggml_time_ms();
        sd::Tensor<float> input = x;
        sd::Tensor<float> output;

        if (tiling_params.enabled) {
            const int scale_factor = get_scale_factor();
            int64_t W              = input.shape()[0] * scale_factor;
            int64_t H              = input.shape()[1] * scale_factor;
            float tile_overlap;
            int tile_size_x, tile_size_y;
            get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, input.shape()[0], input.shape()[1]);
            if (!silent) {
                LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
            }
            output = tiled_compute(
                input,
                n_threads,
                static_cast<int>(W),
                static_cast<int>(H),
                scale_factor,
                tile_size_x,
                tile_size_y,
                tile_overlap,
                circular_x,
                circular_y,
                true,
                tiling_params,
                "vae decode compute failed while processing a tile",
                silent);
        } else {
            output = compute_with_temporal_tiling(n_threads,
                                                  input,
                                                  VAETemporalDirection::DECODE,
                                                  tiling_params);
        }

        runner_end();

        if (output.empty()) {
            LOG_ERROR("vae decode compute failed");
            return {};
        }
        if (scale_input) {
            scale_tensor_to_0_1(&output);
        }
        int64_t t1 = ggml_time_ms();
        LOG_DEBUG("computing vae decode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
        return std::move(output);
    }

    virtual sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) = 0;
    virtual sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents)                           = 0;
    virtual sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents)                           = 0;
    virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors)                                   = 0;
    virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
};

struct FakeVAE : public VAE {
    FakeVAE(SDVersion version,
            ggml_backend_t backend,
            std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
        : VAE(version, backend, "", weight_manager) {}

    int get_encoder_output_channels(int input_channels) {
        return input_channels;
    }

    sd::Tensor<float> _compute(const int n_threads,
                               const sd::Tensor<float>& z,
                               bool decode_graph) override {
        SD_UNUSED(n_threads);
        SD_UNUSED(decode_graph);
        return z;
    }

    sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
        SD_UNUSED(rng);
        return vae_output;
    }

    sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
        return latents;
    }

    sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
        return latents;
    }

    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {}

    std::string get_desc() override {
        return "fake_vae";
    }
};

#endif  // __SD_MODEL_VAE_VAE_HPP__