/*
 * Iris Math Kernels - Implementation
 *
 * Math operations for Iris inference.
 * Uses Metal/MPS on Apple Silicon, BLAS otherwise.
 */

#include "iris_kernels.h"
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* Use Metal for GPU acceleration on Apple Silicon */
#ifdef USE_METAL
#include "iris_metal.h"
#endif

/* Use Vulkan compute for GPU acceleration (cross-platform GEMM offload) */
#ifdef USE_VULKAN
#include "iris_vulkan.h"
#endif

/* Minimum matrix size to use GPU (smaller matrices are faster on CPU) */
#define MIN_GPU_ELEMENTS (512 * 512)

/* fast_expf is defined in iris_kernels.h */

/* Progress callbacks - set by caller before inference */
iris_substep_callback_t       iris_substep_callback       = NULL;
iris_step_callback_t          iris_step_callback          = NULL;
iris_phase_callback_t         iris_phase_callback         = NULL;
iris_step_image_callback_t    iris_step_image_callback    = NULL;
void                         *iris_step_image_vae         = NULL;
iris_text_progress_callback_t iris_text_progress_callback = NULL;
iris_vae_progress_callback_t  iris_vae_progress_callback  = NULL;
int                           iris_verbose                = 0;
int                           iris_circular               = 0;
int                           iris_vae_tiling             = 0;

/* ========================================================================
 * Random Number Generator (xoshiro256**)
 * ======================================================================== */

static uint64_t rng_state[4] = {0x853c49e6748fea9bULL, 0xda3e39cb94b95bdbULL, 0x647c4677a2884327ULL, 0xc6e7918d2e2969f5ULL};

static inline uint64_t rotl(const uint64_t x, int k) {
	return (x << k) | (x >> (64 - k));
}

static uint64_t xoshiro256ss(void) {
	const uint64_t result = rotl(rng_state[1] * 5, 7) * 9;
	const uint64_t t      = rng_state[1] << 17;
	rng_state[2] ^= rng_state[0];
	rng_state[3] ^= rng_state[1];
	rng_state[1] ^= rng_state[2];
	rng_state[0] ^= rng_state[3];
	rng_state[2] ^= t;
	rng_state[3] = rotl(rng_state[3], 45);
	return result;
}

void iris_rng_seed(uint64_t seed) {
	/* SplitMix64 to initialize state from seed */
	for (int i = 0; i < 4; i++) {
		seed += 0x9e3779b97f4a7c15ULL;
		uint64_t z   = seed;
		z            = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
		z            = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
		rng_state[i] = z ^ (z >> 31);
	}
}

float iris_random_uniform(void) {
	return (xoshiro256ss() >> 11) * (1.0 / 9007199254740992.0);
}

float iris_random_normal(void) {
	/* Box-Muller transform */
	float u1 = iris_random_uniform();
	float u2 = iris_random_uniform();
	/* Avoid log(0) */
	while (u1 == 0.0f)
		u1 = iris_random_uniform();
	return sqrtf(-2.0f * logf(u1)) * cosf(2.0f * 3.14159265358979323846f * u2);
}

void iris_randn(float *out, int n) {
	for (int i = 0; i < n; i++) {
		out[i] = iris_random_normal();
	}
}

void iris_rand(float *out, int n) {
	for (int i = 0; i < n; i++) {
		out[i] = iris_random_uniform();
	}
}

/* ========================================================================
 * Basic Element-wise Operations
 * ======================================================================== */

void iris_add(float *out, const float *a, const float *b, int n) {
	for (int i = 0; i < n; i++) {
		out[i] = a[i] + b[i];
	}
}

void iris_add_inplace(float *a, const float *b, int n) {
	for (int i = 0; i < n; i++) {
		a[i] += b[i];
	}
}

void iris_mul_inplace(float *a, const float *b, int n) {
	for (int i = 0; i < n; i++) {
		a[i] *= b[i];
	}
}

void iris_axpy(float *a, float scale, const float *b, int n) {
	for (int i = 0; i < n; i++) {
		a[i] += scale * b[i];
	}
}

/* ========================================================================
 * Matrix Operations
 * ======================================================================== */

/* General matrix multiply C = A @ B. Routes to Metal GPU when the matrix
 * is large enough that GPU compute outweighs the CPU-GPU transfer cost,
 * otherwise falls back to BLAS sgemm or a naive triple loop. This is the
 * backbone operation: every linear projection in the transformer, text
 * encoder, and VAE bottleneck goes through here. */
void iris_matmul(float *C, const float *A, const float *B, int M, int K, int N) {
	/* C[M,N] = A[M,K] @ B[K,N] */

#ifdef USE_METAL
	size_t matrix_elements = (size_t)M * N;
	if (iris_metal_available() && matrix_elements >= MIN_GPU_ELEMENTS) {
		iris_metal_sgemm(0, 0, /* no transpose */
		                 M, N, K, 1.0f, A, K, B, N, 0.0f, C, N);
		return;
	}
#endif

#ifdef USE_VULKAN
	size_t vk_elements = (size_t)M * N;
	if (iris_vulkan_available() && vk_elements >= MIN_GPU_ELEMENTS) {
		iris_vulkan_sgemm(0, 0, /* no transpose */
		                  M, N, K, 1.0f, A, K, B, N, 0.0f, C, N);
		return;
	}
#endif

	/* Naive implementation (CPU fallback when no GPU backend handles this size) */
	for (int m = 0; m < M; m++) {
		for (int n = 0; n < N; n++) {
			float sum = 0.0f;
			for (int k = 0; k < K; k++) {
				sum += A[m * K + k] * B[k * N + n];
			}
			C[m * N + n] = sum;
		}
	}
}

void iris_matmul_t(float *C, const float *A, const float *B, int M, int K, int N) {
	/* C[M,N] = A[M,K] @ B[N,K]^T */

#ifdef USE_METAL
	size_t matrix_elements = (size_t)M * N;
	if (iris_metal_available() && matrix_elements >= MIN_GPU_ELEMENTS) {
		iris_metal_sgemm(0, 1, /* no transpose A, transpose B */
		                 M, N, K, 1.0f, A, K, B, K, 0.0f, C, N);
		return;
	}
#endif

#ifdef USE_VULKAN
	size_t vk_elements = (size_t)M * N;
	if (iris_vulkan_available() && vk_elements >= MIN_GPU_ELEMENTS) {
		iris_vulkan_sgemm(0, 1, /* no transpose A, transpose B */
		                  M, N, K, 1.0f, A, K, B, K, 0.0f, C, N);
		return;
	}
#endif

	/* Naive implementation (CPU fallback when no GPU backend handles this size) */
	for (int m = 0; m < M; m++) {
		for (int n = 0; n < N; n++) {
			float sum = 0.0f;
			for (int k = 0; k < K; k++) {
				sum += A[m * K + k] * B[n * K + k];
			}
			C[m * N + n] = sum;
		}
	}
}

void iris_linear(float *y, const float *x, const float *W, const float *b, int seq_len, int in_dim, int out_dim) {
	/* y[seq, out] = x[seq, in] @ W[out, in]^T + b[out] */

#ifdef USE_METAL
	/* Use Metal GPU for large matrices */
	size_t matrix_elements = (size_t)seq_len * out_dim;
	if (iris_metal_available() && matrix_elements >= MIN_GPU_ELEMENTS) {
		/* Metal sgemm: C = alpha * A @ B^T
		 * A[M, K] = x[seq_len, in_dim]
		 * B[N, K] = W[out_dim, in_dim] (transposed)
		 * C[M, N] = y[seq_len, out_dim]
		 */
		iris_metal_sgemm_cached(0, 1, /* no transpose A, transpose B */
		                        seq_len, out_dim, in_dim, 1.0f, x, in_dim, W, in_dim, 0.0f, y, out_dim);

		/* Add bias if present */
		if (b != NULL) {
			for (int s = 0; s < seq_len; s++) {
				for (int o = 0; o < out_dim; o++) {
					y[s * out_dim + o] += b[o];
				}
			}
		}
		return;
	}
#endif

#ifdef USE_VULKAN
	size_t vk_elements = (size_t)seq_len * out_dim;
	if (iris_vulkan_available() && vk_elements >= MIN_GPU_ELEMENTS) {
		iris_vulkan_sgemm_cached(0, 1, /* no transpose A, transpose B */
		                         seq_len, out_dim, in_dim, 1.0f, x, in_dim, W, in_dim, 0.0f, y, out_dim);
		if (b != NULL) {
			for (int s = 0; s < seq_len; s++) {
				for (int o = 0; o < out_dim; o++) {
					y[s * out_dim + o] += b[o];
				}
			}
		}
		return;
	}
#endif

	/* Naive implementation (CPU fallback when no GPU backend handles this size) */
	for (int s = 0; s < seq_len; s++) {
		const float *x_row = x + s * in_dim;
		float       *y_row = y + s * out_dim;
		for (int o = 0; o < out_dim; o++) {
			const float *w_row = W + o * in_dim;
			float        sum   = (b != NULL) ? b[o] : 0.0f;
			for (int i = 0; i < in_dim; i++) {
				sum += x_row[i] * w_row[i];
			}
			y_row[o] = sum;
		}
	}
}

void iris_linear_nobias(float *y, const float *x, const float *W, int seq_len, int in_dim, int out_dim) {
	iris_linear(y, x, W, NULL, seq_len, in_dim, out_dim);
}

void iris_linear_nobias_bf16(float *y, const float *x, const uint16_t *W_bf16, int seq_len, int in_dim, int out_dim) {
	/* y[seq, out] = x[seq, in] @ W[out, in]^T */

#ifdef USE_METAL
	/* Use Metal GPU for bf16 matmul - provides 2x memory bandwidth */
	size_t matrix_elements = (size_t)seq_len * out_dim;
	if (iris_metal_available() && matrix_elements >= MIN_GPU_ELEMENTS) {
		/* Metal bf16 sgemm: C = alpha * A @ B^T
		 * A[M, K] = x[seq_len, in_dim] (f32)
		 * B[N, K] = W[out_dim, in_dim] (bf16, transposed)
		 * C[M, N] = y[seq_len, out_dim] (f32)
		 */
		iris_metal_sgemm_bf16(0, 1, /* no transpose A, transpose B */
		                      seq_len, out_dim, in_dim, 1.0f, x, in_dim, W_bf16, in_dim, 0.0f, y, out_dim);
		return;
	}
#endif

#ifdef USE_VULKAN
	size_t vk_elements = (size_t)seq_len * out_dim;
	if (iris_vulkan_available() && vk_elements >= MIN_GPU_ELEMENTS) {
		iris_vulkan_sgemm_bf16(0, 1, /* no transpose A, transpose B */
		                       seq_len, out_dim, in_dim, 1.0f, x, in_dim, W_bf16, in_dim, 0.0f, y, out_dim);
		return;
	}
#endif

	/* Fallback: convert bf16 to f32 and use regular linear */
	float *W_f32 = (float *)malloc((size_t)out_dim * in_dim * sizeof(float));
	if (!W_f32)
		return;

	/* Convert bf16 to f32 */
	for (int i = 0; i < out_dim * in_dim; i++) {
		uint32_t f32_bits = ((uint32_t)W_bf16[i]) << 16;
		memcpy(&W_f32[i], &f32_bits, sizeof(float));
	}

	iris_linear_nobias(y, x, W_f32, seq_len, in_dim, out_dim);
	free(W_f32);
}

/* Dequantize one GGML Q8_0 element at flat index e from a raw block stream:
 *   value = qs[e % 32] * fp16_to_f32(scale[e / 32])
 * Mirrors the in-shader load_q8() in iris_vulkan_gemm_q8.comp. */
static float dequant_q8_0(const uint8_t *blocks, size_t e) {
	size_t         block  = e >> 5;  /* e / 32 */
	size_t         within = e & 31u; /* e % 32 */
	const uint8_t *blk    = blocks + block * 34u;
	uint16_t       scale_h;
	memcpy(&scale_h, blk, 2);
	/* fp16 -> f32 */
	uint32_t s    = scale_h;
	uint32_t sign = (s & 0x8000u) << 16;
	uint32_t exp  = (s >> 10) & 0x1fu;
	uint32_t mant = s & 0x3ffu;
	uint32_t f;
	if (exp == 0) {
		f = sign; /* subnormal scales round to 0; negligible for Q8_0 */
	}
	else if (exp == 0x1f) {
		f = sign | 0x7f800000u | (mant << 13);
	}
	else {
		f = sign | ((exp + 112u) << 23) | (mant << 13);
	}
	float scale;
	memcpy(&scale, &f, 4);
	int8_t q = (int8_t)blk[2 + within];
	return (float)q * scale;
}

void iris_linear_nobias_q8(float *y, const float *x, const void *W_q8, int seq_len, int in_dim, int out_dim) {
	/* y[seq, out] = x[seq, in] @ W[out, in]^T, W stored as Q8_0 blocks. */

#ifdef USE_VULKAN
	size_t vk_elements = (size_t)seq_len * out_dim;
	if (iris_vulkan_available() && vk_elements >= MIN_GPU_ELEMENTS) {
		iris_vulkan_sgemm_q8(0, 1, /* no transpose A, transpose B */
		                     seq_len, out_dim, in_dim, 1.0f, x, in_dim, W_q8, in_dim, 0.0f, y, out_dim);
		return;
	}
#endif

	/* Fallback: dequantize the Q8_0 weight to f32 and use regular linear.
	 * Only reached on non-Vulkan builds or for sub-threshold matmuls. */
	const uint8_t *blocks = (const uint8_t *)W_q8;
	float         *W_f32  = (float *)malloc((size_t)out_dim * in_dim * sizeof(float));
	if (!W_f32)
		return;
	for (size_t e = 0; e < (size_t)out_dim * in_dim; e++) {
		W_f32[e] = dequant_q8_0(blocks, e);
	}
	iris_linear_nobias(y, x, W_f32, seq_len, in_dim, out_dim);
	free(W_f32);
}

/* ========================================================================
 * GPU Batch Operations
 * ======================================================================== */

void iris_gpu_begin_batch(void) {
#ifdef USE_METAL
	iris_metal_begin_batch();
#endif
}

void iris_gpu_end_batch(void) {
#ifdef USE_METAL
	iris_metal_end_batch();
#endif
}

/* ========================================================================
 * Convolution Operations
 * ======================================================================== */

/* 2D convolution via im2col + GEMM: reshapes input so each column is a
 * flattened receptive field, then multiplies by the kernel weight matrix.
 * Tiles spatially to bound memory usage for large feature maps. This is the
 * standard approach for BLAS/GPU-friendly convolution, used throughout the
 * VAE encoder and decoder. */
/* Wrap an (possibly out-of-range) coordinate into [0, n) for circular padding. */
static inline int iris_wrap_coord(int c, int n) {
	c %= n;
	if (c < 0)
		c += n;
	return c;
}

void iris_conv2d(float *out, const float *in, const float *weight, const float *bias, int batch, int in_ch, int out_ch, int H, int W, int kH, int kW,
                 int stride, int padding) {
	int outH     = (H + 2 * padding - kH) / stride + 1;
	int outW     = (W + 2 * padding - kW) / stride + 1;
	int circular = iris_circular;

#ifdef USE_VULKAN
	/* im2col + GEMM path, with the matrix multiply dispatched to the Vulkan GPU.
	 * Without an available device the im2col build-up is pure overhead, so fall
	 * straight through to the naive implementation. */
	if (!iris_vulkan_available())
		goto naive_fallback;

	/* im2col + GEMM optimization with tiling for large convolutions */
	size_t col_size     = (size_t)in_ch * kH * kW * outH * outW;
	size_t max_col_size = (size_t)256 * 1024 * 1024; /* 1GB limit */

	/* For large convolutions, process in row tiles */
	int tile_rows = outH;
	if (col_size > max_col_size) {
		/* Calculate how many rows we can process at once */
		size_t row_size = (size_t)in_ch * kH * kW * outW;
		tile_rows       = (int)(max_col_size / row_size);
		if (tile_rows < 1)
			tile_rows = 1;
	}

	size_t tile_col_size = (size_t)in_ch * kH * kW * tile_rows * outW;
	float *col           = malloc(tile_col_size * sizeof(float));
	if (!col) {
		goto naive_fallback;
	}

	for (int b = 0; b < batch; b++) {
		const float *in_b  = in + b * in_ch * H * W;
		float       *out_b = out + b * out_ch * outH * outW;

		/* Process in tiles of rows */
		for (int tile_start = 0; tile_start < outH; tile_start += tile_rows) {
			int tile_end = tile_start + tile_rows;
			if (tile_end > outH)
				tile_end = outH;
			int tile_h      = tile_end - tile_start;
			int tile_pixels = tile_h * outW;

			/* im2col for this tile: col[in_ch*kH*kW, tile_pixels] */
			int col_row = 0;
			for (int ic = 0; ic < in_ch; ic++) {
				for (int kh = 0; kh < kH; kh++) {
					for (int kw = 0; kw < kW; kw++) {
						for (int oh = tile_start; oh < tile_end; oh++) {
							for (int ow = 0; ow < outW; ow++) {
								int ih      = oh * stride - padding + kh;
								int iw      = ow * stride - padding + kw;
								int col_idx = col_row * tile_pixels + (oh - tile_start) * outW + ow;
								if (ih >= 0 && ih < H && iw >= 0 && iw < W) {
									col[col_idx] = in_b[ic * H * W + ih * W + iw];
								}
								else if (circular) {
									int wh       = iris_wrap_coord(ih, H);
									int ww       = iris_wrap_coord(iw, W);
									col[col_idx] = in_b[ic * H * W + wh * W + ww];
								}
								else {
									col[col_idx] = 0.0f;
								}
							}
						}
						col_row++;
					}
				}
			}

			/* GEMM: tmp[out_ch, tile_pixels] = weight[out_ch, K] @ col[K, tile_pixels]
			 * where K = in_ch * kH * kW */
			int K = in_ch * kH * kW;

			/* Write sgemm output directly to out_b using strided ldc.
			 * Row oc of sgemm output goes to out_b[oc * outH*outW + tile_start*outW],
			 * which is exactly the right position in NCHW layout. */
			float *out_tile = out_b + tile_start * outW;
			iris_vulkan_sgemm(0, 0, out_ch, tile_pixels, K, 1.0f, weight, K, col, tile_pixels, 0.0f, out_tile, outH * outW);
		}

		/* Add bias */
		if (bias != NULL) {
			for (int oc = 0; oc < out_ch; oc++) {
				float  b_val      = bias[oc];
				float *out_ch_ptr = out_b + oc * outH * outW;
				for (int i = 0; i < outH * outW; i++) {
					out_ch_ptr[i] += b_val;
				}
			}
		}
	}

	free(col);
	return;

naive_fallback:
#endif
	/* Naive implementation (fallback) */
	for (int b = 0; b < batch; b++) {
		for (int oc = 0; oc < out_ch; oc++) {
			for (int oh = 0; oh < outH; oh++) {
				for (int ow = 0; ow < outW; ow++) {
					float sum = (bias != NULL) ? bias[oc] : 0.0f;

					for (int ic = 0; ic < in_ch; ic++) {
						for (int kh = 0; kh < kH; kh++) {
							for (int kw = 0; kw < kW; kw++) {
								int ih = oh * stride - padding + kh;
								int iw = ow * stride - padding + kw;

								if (circular) {
									ih = iris_wrap_coord(ih, H);
									iw = iris_wrap_coord(iw, W);
								}
								if (ih >= 0 && ih < H && iw >= 0 && iw < W) {
									int in_idx = b * in_ch * H * W + ic * H * W + ih * W + iw;
									int w_idx  = oc * in_ch * kH * kW + ic * kH * kW + kh * kW + kw;
									sum += in[in_idx] * weight[w_idx];
								}
							}
						}
					}

					int out_idx  = b * out_ch * outH * outW + oc * outH * outW + oh * outW + ow;
					out[out_idx] = sum;
				}
			}
		}
	}
}

/* ========================================================================
 * Normalization
 * ======================================================================== */

void iris_rms_norm(float *out, const float *x, const float *weight, int seq_len, int hidden, float eps) {
#ifdef USE_METAL
	/* Use GPU for RMSNorm only for very large tensors
	 * The CPU-GPU sync overhead usually outweighs benefits for smaller ops */
	size_t elements = (size_t)seq_len * hidden;
	if (iris_metal_shaders_available() && elements >= 1024 * 1024) {
		iris_metal_rms_norm(out, x, weight, seq_len, hidden, eps);
		return;
	}
#endif

	for (int s = 0; s < seq_len; s++) {
		const float *x_row   = x + s * hidden;
		float       *out_row = out + s * hidden;

		/* Compute RMS */
		float sum_sq = 0.0f;
		for (int i = 0; i < hidden; i++) {
			sum_sq += x_row[i] * x_row[i];
		}
		float rms     = sqrtf(sum_sq / hidden + eps);
		float rms_inv = 1.0f / rms;

		/* Normalize and scale */
		for (int i = 0; i < hidden; i++) {
			out_row[i] = x_row[i] * rms_inv * weight[i];
		}
	}
}

void iris_group_norm(float *out, const float *x, const float *gamma, const float *beta, int batch, int channels, int H, int W, int num_groups, float eps) {
	int channels_per_group = channels / num_groups;
	int spatial            = H * W;

	for (int b = 0; b < batch; b++) {
		for (int g = 0; g < num_groups; g++) {
			int c_start = g * channels_per_group;
			int c_end   = c_start + channels_per_group;

			float mean  = 0.0f;
			int   count = 0;
			for (int c = c_start; c < c_end; c++) {
				for (int i = 0; i < spatial; i++) {
					int idx = b * channels * spatial + c * spatial + i;
					mean += x[idx];
					count++;
				}
			}
			mean /= count;

			float var = 0.0f;
			for (int c = c_start; c < c_end; c++) {
				for (int i = 0; i < spatial; i++) {
					int   idx  = b * channels * spatial + c * spatial + i;
					float diff = x[idx] - mean;
					var += diff * diff;
				}
			}
			var /= count;

			float std_inv = 1.0f / sqrtf(var + eps);

			for (int c = c_start; c < c_end; c++) {
				for (int i = 0; i < spatial; i++) {
					int   idx  = b * channels * spatial + c * spatial + i;
					float norm = (x[idx] - mean) * std_inv;
					out[idx]   = gamma[c] * norm + beta[c];
				}
			}
		}
	}
}

void iris_batch_norm(float *out, const float *x, const float *running_mean, const float *running_var, const float *gamma, const float *beta, int batch,
                     int channels, int H, int W, float eps) {
	int spatial = H * W;

	for (int c = 0; c < channels; c++) {
		float mean    = running_mean[c];
		float var     = running_var[c];
		float std_inv = 1.0f / sqrtf(var + eps);
		float g       = (gamma != NULL) ? gamma[c] : 1.0f;
		float b_val   = (beta != NULL) ? beta[c] : 0.0f;

		for (int n = 0; n < batch; n++) {
			for (int i = 0; i < spatial; i++) {
				int idx  = n * channels * spatial + c * spatial + i;
				out[idx] = g * (x[idx] - mean) * std_inv + b_val;
			}
		}
	}
}

/* ========================================================================
 * Activation Functions
 * ======================================================================== */

void iris_silu(float *x, int n) {
#ifdef USE_METAL
	/* Use GPU for very large arrays (overhead not worth it for small ones) */
	if (iris_metal_shaders_available() && n >= 4 * 1024 * 1024) {
		iris_metal_silu(x, n);
		return;
	}
#endif

	for (int i = 0; i < n; i++) {
		float val = x[i];
		x[i]      = val / (1.0f + fast_expf(-val));
	}
}

/* Fused SiLU(gate) * up in a single pass - avoids double memory traversal */
void iris_silu_mul(float *gate, const float *up, int n) {
#ifdef USE_METAL
	if (iris_metal_shaders_available() && n >= 4 * 1024 * 1024) {
		iris_metal_silu_mul(gate, up, n);
		return;
	}
#endif

	for (int i = 0; i < n; i++) {
		float val = gate[i];
		gate[i]   = (val / (1.0f + fast_expf(-val))) * up[i];
	}
}

/* CPU-only softmax. Safe to call from worker threads (no Metal dispatch). */
void iris_softmax_cpu(float *x, int rows, int cols) {
	for (int r = 0; r < rows; r++) {
		float *row = x + r * cols;

		/* Find max for numerical stability */
		float max_val = row[0];
		for (int c = 1; c < cols; c++) {
			if (row[c] > max_val)
				max_val = row[c];
		}

		/* Compute exp and sum */
		float sum = 0.0f;
		for (int c = 0; c < cols; c++) {
			row[c] = fast_expf(row[c] - max_val);
			sum += row[c];
		}

		/* Normalize */
		float inv_sum = 1.0f / sum;
		for (int c = 0; c < cols; c++) {
			row[c] *= inv_sum;
		}
	}
}

void iris_softmax(float *x, int rows, int cols) {
#ifdef USE_METAL
	/* Use GPU only for very large softmax operations
	 * Sync overhead usually dominates for smaller ops */
	if (iris_metal_shaders_available() && (size_t)rows * cols >= 4 * 1024 * 1024) {
		iris_metal_softmax(x, rows, cols);
		return;
	}
#endif
	iris_softmax_cpu(x, rows, cols);
}

/* ========================================================================
 * Attention Operations
 * ======================================================================== */

/* Scaled dot-product attention: softmax(Q @ K^T / sqrt(d)) @ V.
 * This is the naive implementation that materializes the full seq_q x seq_k
 * attention matrix. Used only for small sequences; the transformer's main
 * attention path uses iris_flash_attention() or the GPU kernel instead. */
void iris_attention(float *out, const float *Q, const float *K, const float *V, int batch, int heads, int seq_q, int seq_k, int head_dim, float scale) {
	/* Allocate attention scores */
	float *scores = (float *)malloc(seq_q * seq_k * sizeof(float));

	for (int b = 0; b < batch; b++) {
		for (int h = 0; h < heads; h++) {
			const float *q = Q + (b * heads + h) * seq_q * head_dim;
			const float *k = K + (b * heads + h) * seq_k * head_dim;
			const float *v = V + (b * heads + h) * seq_k * head_dim;
			float       *o = out + (b * heads + h) * seq_q * head_dim;

			/* scores = Q @ K^T * scale */
			for (int i = 0; i < seq_q; i++) {
				for (int j = 0; j < seq_k; j++) {
					float dot = 0.0f;
					for (int d = 0; d < head_dim; d++) {
						dot += q[i * head_dim + d] * k[j * head_dim + d];
					}
					scores[i * seq_k + j] = dot * scale;
				}
			}

			/* softmax */
			iris_softmax(scores, seq_q, seq_k);

			/* out = scores @ V */
			for (int i = 0; i < seq_q; i++) {
				for (int d = 0; d < head_dim; d++) {
					float sum = 0.0f;
					for (int j = 0; j < seq_k; j++) {
						sum += scores[i * seq_k + j] * v[j * head_dim + d];
					}
					o[i * head_dim + d] = sum;
				}
			}
		}
	}

	free(scores);
}

/* ========================================================================
 * Flash Attention - Memory-Efficient Tiled Attention
 *
 * Uses online softmax algorithm to compute attention without materializing
 * the full [seq_q, seq_k] attention matrix. Reduces memory from O(n²) to O(n).
 *
 * Algorithm (for each query position):
 * 1. Initialize: max_score = -inf, sum = 0, output = 0
 * 2. For each key/value block:
 *    - Compute local scores = Q @ K^T * scale
 *    - Update running max and sum with correction factors
 *    - Accumulate weighted values into output
 * 3. Normalize: output /= sum
 *
 * Reference: "FlashAttention: Fast and Memory-Efficient Exact Attention"
 * ======================================================================== */

/*
 * Flash attention for a single head.
 * Q: [seq_q, head_dim], K: [seq_k, head_dim], V: [seq_k, head_dim]
 * out: [seq_q, head_dim]
 * Uses O(head_dim) working memory per query instead of O(seq_k).
 */
static void flash_attention_head(float *out, const float *Q, const float *K, const float *V, int seq_q, int seq_k, int head_dim, float scale) {
	/* Process each query position independently */
	for (int i = 0; i < seq_q; i++) {
		const float *q_row = Q + i * head_dim;
		float       *o_row = out + i * head_dim;

		/* Running statistics for online softmax */
		float max_score = -1e30f; /* Large negative value (avoid -INFINITY with -ffast-math) */
		float sum_exp   = 0.0f;

		/* Initialize output to zero */
		for (int d = 0; d < head_dim; d++) {
			o_row[d] = 0.0f;
		}

		/* Iterate over all key/value positions */
		for (int j = 0; j < seq_k; j++) {
			const float *k_row = K + j * head_dim;
			const float *v_row = V + j * head_dim;

			/* Compute attention score: Q[i] · K[j] * scale */
			float score = 0.0f;
			for (int d = 0; d < head_dim; d++) {
				score += q_row[d] * k_row[d];
			}
			score *= scale;

			/* Online softmax update */
			if (score > max_score) {
				/* New maximum found - rescale previous accumulations */
				float correction = fast_expf(max_score - score);
				sum_exp          = sum_exp * correction + 1.0f;
				for (int d = 0; d < head_dim; d++) {
					o_row[d] = o_row[d] * correction + v_row[d];
				}
				max_score = score;
			}
			else {
				/* Score is less than current max */
				float weight = fast_expf(score - max_score);
				sum_exp += weight;
				for (int d = 0; d < head_dim; d++) {
					o_row[d] += weight * v_row[d];
				}
			}
		}

		/* Normalize by sum */
		float inv_sum = 1.0f / sum_exp;
		for (int d = 0; d < head_dim; d++) {
			o_row[d] *= inv_sum;
		}
	}
}

/*
 * Flash attention with BLAS-optimized tiling.
 * Processes queries in tiles for better cache utilization.
 * Uses BLAS for tile-level matrix operations when available.
 *
 * Q: [seq_q, head_dim], K: [seq_k, head_dim], V: [seq_k, head_dim]
 * out: [seq_q, head_dim]
 * tile_scores: scratch buffer of size [q_tile_size, k_tile_size]
 */
static void flash_attention_head_tiled(float *out, const float *Q, const float *K, const float *V, int seq_q, int seq_k, int head_dim, float scale,
                                       float *tile_scores, int q_tile_size, int k_tile_size) {
	/* Per-query running statistics: max_score[seq_q], sum_exp[seq_q] */
	float *max_scores = (float *)malloc(seq_q * sizeof(float));
	float *sum_exps   = (float *)malloc(seq_q * sizeof(float));

	/* Initialize */
	for (int i = 0; i < seq_q; i++) {
		max_scores[i] = -1e30f; /* Large negative value (avoid -INFINITY with -ffast-math) */
		sum_exps[i]   = 0.0f;
	}
	memset(out, 0, seq_q * head_dim * sizeof(float));

	/* Process in tiles over K/V dimension */
	for (int k_start = 0; k_start < seq_k; k_start += k_tile_size) {
		int k_end = (k_start + k_tile_size < seq_k) ? k_start + k_tile_size : seq_k;
		int k_len = k_end - k_start;

		/* Process in tiles over Q dimension */
		for (int q_start = 0; q_start < seq_q; q_start += q_tile_size) {
			int q_end = (q_start + q_tile_size < seq_q) ? q_start + q_tile_size : seq_q;
			int q_len = q_end - q_start;

			const float *Q_tile   = Q + q_start * head_dim;
			const float *K_tile   = K + k_start * head_dim;
			const float *V_tile   = V + k_start * head_dim;
			float       *out_tile = out + q_start * head_dim;

			/* Compute tile scores: Q_tile @ K_tile^T * scale */
			for (int qi = 0; qi < q_len; qi++) {
				for (int ki = 0; ki < k_len; ki++) {
					float dot = 0.0f;
					for (int d = 0; d < head_dim; d++) {
						dot += Q_tile[qi * head_dim + d] * K_tile[ki * head_dim + d];
					}
					tile_scores[qi * k_tile_size + ki] = dot * scale;
				}
			}

			/* Online softmax update for this tile */
			for (int qi = 0; qi < q_len; qi++) {
				int    i         = q_start + qi;
				float *score_row = tile_scores + qi * k_tile_size;
				float *o_row     = out_tile + qi * head_dim;

				/* Find max in this tile */
				float tile_max = score_row[0];
				for (int ki = 1; ki < k_len; ki++) {
					if (score_row[ki] > tile_max)
						tile_max = score_row[ki];
				}

				/* Compute correction factors */
				float old_max = max_scores[i];
				float new_max = (tile_max > old_max) ? tile_max : old_max;

				/* Rescale old accumulations if needed */
				if (old_max > -1e29f) { /* Check if we have prior accumulations */
					float correction = fast_expf(old_max - new_max);
					sum_exps[i] *= correction;
					for (int d = 0; d < head_dim; d++) {
						o_row[d] *= correction;
					}
				}

				/* Accumulate this tile's contribution */
				for (int ki = 0; ki < k_len; ki++) {
					float weight = fast_expf(score_row[ki] - new_max);
					sum_exps[i] += weight;
					const float *v_row = V_tile + ki * head_dim;
					for (int d = 0; d < head_dim; d++) {
						o_row[d] += weight * v_row[d];
					}
				}

				max_scores[i] = new_max;
			}
		}
	}

	/* Final normalization */
	for (int i = 0; i < seq_q; i++) {
		float  inv_sum = 1.0f / sum_exps[i];
		float *o_row   = out + i * head_dim;
		for (int d = 0; d < head_dim; d++) {
			o_row[d] *= inv_sum;
		}
	}

	free(max_scores);
	free(sum_exps);
}

/*
 * Flash attention for multi-head attention.
 * Works on [seq, heads*head_dim] layout (same as transformer tensors).
 *
 * Q: [seq_q, heads * head_dim]
 * K: [seq_k, heads * head_dim]
 * V: [seq_k, heads * head_dim]
 * out: [seq_q, heads * head_dim]
 *
 * Memory usage: O(seq_q + tile_size²) instead of O(seq_q * seq_k)
 */
void iris_flash_attention(float *out, const float *Q, const float *K, const float *V, int seq_q, int seq_k, int heads, int head_dim, float scale) {
	/* Tile sizes for cache efficiency */
	int q_tile_size = 32; /* Process 32 queries at a time */
	int k_tile_size = 64; /* Process 64 keys at a time */

	/* Allocate tile scratch buffer */
	float *tile_scores = (float *)malloc(q_tile_size * k_tile_size * sizeof(float));

	/* Process each head */
	for (int h = 0; h < heads; h++) {
		const float *Q_head   = Q + h * head_dim;
		const float *K_head   = K + h * head_dim;
		const float *V_head   = V + h * head_dim;
		float       *out_head = out + h * head_dim;

		/* Stride between consecutive positions for this head */
		int hidden = heads * head_dim;

		/* For small sequences, use simple non-tiled version */
		if (seq_q <= 64 && seq_k <= 128) {
			/* Extract head data into contiguous buffers */
			float *Q_contig   = (float *)malloc(seq_q * head_dim * sizeof(float));
			float *K_contig   = (float *)malloc(seq_k * head_dim * sizeof(float));
			float *V_contig   = (float *)malloc(seq_k * head_dim * sizeof(float));
			float *out_contig = (float *)malloc(seq_q * head_dim * sizeof(float));

			for (int i = 0; i < seq_q; i++) {
				for (int d = 0; d < head_dim; d++) {
					Q_contig[i * head_dim + d] = Q_head[i * hidden + d];
				}
			}
			for (int j = 0; j < seq_k; j++) {
				for (int d = 0; d < head_dim; d++) {
					K_contig[j * head_dim + d] = K_head[j * hidden + d];
					V_contig[j * head_dim + d] = V_head[j * hidden + d];
				}
			}

			flash_attention_head(out_contig, Q_contig, K_contig, V_contig, seq_q, seq_k, head_dim, scale);

			/* Copy back with stride */
			for (int i = 0; i < seq_q; i++) {
				for (int d = 0; d < head_dim; d++) {
					out_head[i * hidden + d] = out_contig[i * head_dim + d];
				}
			}

			free(Q_contig);
			free(K_contig);
			free(V_contig);
			free(out_contig);
		}
		else {
			/* For larger sequences, use tiled version with strided access */
			/* Extract head data into contiguous buffers for BLAS efficiency */
			float *Q_contig   = (float *)malloc(seq_q * head_dim * sizeof(float));
			float *K_contig   = (float *)malloc(seq_k * head_dim * sizeof(float));
			float *V_contig   = (float *)malloc(seq_k * head_dim * sizeof(float));
			float *out_contig = (float *)malloc(seq_q * head_dim * sizeof(float));

			for (int i = 0; i < seq_q; i++) {
				for (int d = 0; d < head_dim; d++) {
					Q_contig[i * head_dim + d] = Q_head[i * hidden + d];
				}
			}
			for (int j = 0; j < seq_k; j++) {
				for (int d = 0; d < head_dim; d++) {
					K_contig[j * head_dim + d] = K_head[j * hidden + d];
					V_contig[j * head_dim + d] = V_head[j * hidden + d];
				}
			}

			flash_attention_head_tiled(out_contig, Q_contig, K_contig, V_contig, seq_q, seq_k, head_dim, scale, tile_scores, q_tile_size, k_tile_size);

			/* Copy back with stride */
			for (int i = 0; i < seq_q; i++) {
				for (int d = 0; d < head_dim; d++) {
					out_head[i * hidden + d] = out_contig[i * head_dim + d];
				}
			}

			free(Q_contig);
			free(K_contig);
			free(V_contig);
			free(out_contig);
		}
	}

	free(tile_scores);
}

/* Apply precomputed RoPE (Rotary Position Embedding) in-place using the
 * split-half convention: dim d pairs with dim d+half for rotation. This is
 * the Flux convention (4-axis, split-half). RoPE lets the transformer learn
 * relative position from the dot-product structure of Q and K. */
void iris_apply_rope(float *x, const float *freqs, int batch, int seq, int heads, int head_dim) {
	/* x: [batch, seq, heads, head_dim]
	 * freqs: [seq, head_dim/2, 2] (cos, sin)
	 * Apply rotary embedding to pairs of dimensions */

	int half_dim = head_dim / 2;

	for (int b = 0; b < batch; b++) {
		for (int s = 0; s < seq; s++) {
			for (int h = 0; h < heads; h++) {
				float *vec = x + ((b * seq + s) * heads + h) * head_dim;

				for (int d = 0; d < half_dim; d++) {
					float cos_val = freqs[s * half_dim * 2 + d * 2];
					float sin_val = freqs[s * half_dim * 2 + d * 2 + 1];

					float x0 = vec[d];
					float x1 = vec[d + half_dim];

					vec[d]            = x0 * cos_val - x1 * sin_val;
					vec[d + half_dim] = x0 * sin_val + x1 * cos_val;
				}
			}
		}
	}
}

void iris_compute_rope_freqs(float *freqs, const int *pos, int seq, int dim, float theta) {
	int half_dim = dim / 2;

	for (int s = 0; s < seq; s++) {
		float p = (float)pos[s];
		for (int d = 0; d < half_dim; d++) {
			float freq                          = 1.0f / powf(theta, (float)(2 * d) / (float)dim);
			float angle                         = p * freq;
			freqs[s * half_dim * 2 + d * 2]     = cosf(angle);
			freqs[s * half_dim * 2 + d * 2 + 1] = sinf(angle);
		}
	}
}

/* ========================================================================
 * Pooling and Reshape
 * ======================================================================== */

void iris_upsample_nearest(float *out, const float *in, int batch, int channels, int H, int W, int scale_h, int scale_w) {
	int outH = H * scale_h;
	int outW = W * scale_w;

	for (int b = 0; b < batch; b++) {
		for (int c = 0; c < channels; c++) {
			for (int oh = 0; oh < outH; oh++) {
				for (int ow = 0; ow < outW; ow++) {
					int ih       = oh / scale_h;
					int iw       = ow / scale_w;
					int in_idx   = b * channels * H * W + c * H * W + ih * W + iw;
					int out_idx  = b * channels * outH * outW + c * outH * outW + oh * outW + ow;
					out[out_idx] = in[in_idx];
				}
			}
		}
	}
}

/* Convert spatial latent to patch tokens for the diffusion transformer.
 * Groups each ps x ps spatial block into a single token vector:
 * [batch, channels, H, W] -> [batch, channels*ps*ps, H/ps, W/ps].
 * The transformer operates on these patch tokens, not individual spatial
 * positions, reducing sequence length by ps*ps (4x for ps=2). */
void iris_patchify(float *out, const float *in, int batch, int channels, int H, int W, int patch_size) {
	/* [B, C, H, W] -> [B, C*p*p, H/p, W/p] */
	int p      = patch_size;
	int outH   = H / p;
	int outW   = W / p;
	int out_ch = channels * p * p;

	for (int b = 0; b < batch; b++) {
		for (int c = 0; c < channels; c++) {
			for (int ph = 0; ph < outH; ph++) {
				for (int pw = 0; pw < outW; pw++) {
					for (int pi = 0; pi < p; pi++) {
						for (int pj = 0; pj < p; pj++) {
							int ih     = ph * p + pi;
							int iw     = pw * p + pj;
							int in_idx = b * channels * H * W + c * H * W + ih * W + iw;

							int out_c    = c * p * p + pi * p + pj;
							int out_idx  = b * out_ch * outH * outW + out_c * outH * outW + ph * outW + pw;
							out[out_idx] = in[in_idx];
						}
					}
				}
			}
		}
	}
}

void iris_unpatchify(float *out, const float *in, int batch, int channels, int H, int W, int patch_size) {
	/* [B, C*p*p, H, W] -> [B, C, H*p, W*p] */
	int p     = patch_size;
	int in_ch = channels * p * p;
	int outH  = H * p;
	int outW  = W * p;

	for (int b = 0; b < batch; b++) {
		for (int c = 0; c < channels; c++) {
			for (int ph = 0; ph < H; ph++) {
				for (int pw = 0; pw < W; pw++) {
					for (int pi = 0; pi < p; pi++) {
						for (int pj = 0; pj < p; pj++) {
							int in_c   = c * p * p + pi * p + pj;
							int in_idx = b * in_ch * H * W + in_c * H * W + ph * W + pw;

							int oh       = ph * p + pi;
							int ow       = pw * p + pj;
							int out_idx  = b * channels * outH * outW + c * outH * outW + oh * outW + ow;
							out[out_idx] = in[in_idx];
						}
					}
				}
			}
		}
	}
}

/* ========================================================================
 * Utility Functions
 * ======================================================================== */

void iris_copy(float *dst, const float *src, int n) {
	memcpy(dst, src, n * sizeof(float));
}