--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local BLAS = require("../ops/BLAS")
local Types = require("../Types")
local Util = require("../Util")
local assert = Util.Assert

local function copyShape(shape: { number }): { number }
	local out = table.create(#shape)
	for i = 1, #shape do
		out[i] = shape[i]
	end
	return out
end

local function makeMatrixView(tensor: Types.Tensor, rows: number, cols: number): Types.Tensor
	return {
		_storage = tensor._storage,
		_shape = { rows, cols },
		_dtype = tensor._dtype,
		_requiresGrad = tensor._requiresGrad,
		_parents = nil,
		_backprop = nil,
		_grad = nil,
		_size = nil,
	}
end

local function reshapeInPlace(tensor: Types.Tensor, shape: { number }): Types.Tensor
	tensor._shape = copyShape(shape)
	return tensor
end

local function transposeMatrix(matrix: Types.Tensor): Types.Tensor
	local rows, cols = matrix._shape[1], matrix._shape[2]
	local out = Tensor.zeros({ cols, rows }, matrix._dtype, matrix._requiresGrad)
	task.desynchronize()
	local os, ms = out._storage, matrix._storage
	for r = 1, rows do
		local rowOffset = (r - 1) * cols
		for c = 1, cols do
			os[(c - 1) * rows + r] = ms[rowOffset + c]
		end
	end
	task.synchronize()
	return out
end

local function col2im(
	cols: Types.Tensor,
	C: number,
	H: number,
	Wd: number,
	B: number,
	KH: number,
	KW: number,
	Ho: number,
	Wo: number,
	dtype: Types.DType
): Types.Tensor
	local patch = C * KH * KW
	local out = Tensor.zeros({ C, H, Wd, B }, dtype, false)
	local xs = out._storage
	local cs = cols._storage
	local HWd_B = H * Wd * B
	local Wd_B = Wd * B
	task.desynchronize()
	local col = 0
	for i = 1, Ho do
		for j = 1, Wo do
			for b = 1, B do
				local base = col * patch
				col += 1
				local offset = 0
				local b_minus_1 = b - 1
				for c = 1, C do
					local c_offset = (c - 1) * HWd_B
					for ki = 0, KH - 1 do
						local row_part = (i + ki - 1) * Wd_B
						for kj = 0, KW - 1 do
							local col_part = (j + kj - 1) * B
							local xi = c_offset + row_part + col_part + b_minus_1

							xs[xi + 1] += cs[base + offset + 1]
							offset += 1
						end
					end
				end
			end
		end
	end
	task.synchronize()
	return out
end

local function im2col(
	X: Types.Tensor,
	C: number,
	H: number,
	Wd: number,
	B: number,
	KH: number,
	KW: number,
	Ho: number,
	Wo: number
)
	local patch = C * KH * KW
	local cols = Tensor.zeros({ patch, Ho * Wo * B }, X._dtype, false)
	local xs = X._storage
	local cs = cols._storage
	local HWd_B = H * Wd * B
	local Wd_B = Wd * B
	task.desynchronize()
	local col = 0
	for i = 1, Ho do
		for j = 1, Wo do
			for b = 1, B do
				local base = col * patch
				col += 1
				local offset = 0
				local b_minus_1 = b - 1
				for c = 1, C do
					local c_offset = (c - 1) * HWd_B
					for ki = 0, KH - 1 do
						local row_part = (i + ki - 1) * Wd_B
						for kj = 0, KW - 1 do
							local col_part = (j + kj - 1) * B
							local xi = c_offset + row_part + col_part + b_minus_1
							cs[base + offset + 1] = xs[xi + 1]
							offset += 1
						end
					end
				end
			end
		end
	end
	task.synchronize()
	return cols
end

local function addInto(dst: Types.Tensor, src: Types.Tensor)
	task.desynchronize()
	local ds, ss = dst._storage, src._storage
	for i = 1, #ds do
		ds[i] += ss[i]
	end
	task.synchronize()
end

return function(X: Types.Tensor, W: Types.Tensor): Types.Tensor
	local C_in, H, Wd, B = X._shape[1], X._shape[2], X._shape[3], X._shape[4]
	local Cin_w, C_out, KH, KW = W._shape[1], W._shape[2], W._shape[3], W._shape[4]
	assert(C_in == Cin_w, "convTranspose2d: in_channels mismatch")
	task.desynchronize()
	local Ho = H + KH - 1
	local Wo = Wd + KW - 1
	local X_flat = makeMatrixView(X, C_in, H * Wd * B)
	local W_flat = makeMatrixView(W, C_in, C_out * KH * KW)
	local W_T = transposeMatrix(W_flat) -- {C_out * KH * KW, C_in}
	local cols = BLAS.matmul(W_T, X_flat) -- {C_out * KH * KW, H * Wd * B}
	local Y = col2im(cols, C_out, Ho, Wo, B, KH, KW, H, Wd, X._dtype)
	task.synchronize()
	return Tape.makeNode(Y, { X, W }, function(gy: Types.Tensor)
		local gy_cols = im2col(gy, C_out, Ho, Wo, B, KH, KW, H, Wd)
		if X._requiresGrad then
			if not X._grad then
				X._grad = Tensor.zeros({ C_in, H, Wd, B }, X._dtype, false)
			end
			local dX_flat = BLAS.matmul(W_flat, gy_cols)
			local dX = reshapeInPlace(dX_flat, { C_in, H, Wd, B })
			addInto(X._grad :: Types.Tensor, dX)
		end
		if W._requiresGrad then
			if not W._grad then
				W._grad = Tensor.zeros({ C_in, C_out, KH, KW }, W._dtype, false)
			end
			local gy_cols_T = transposeMatrix(gy_cols)
			local dW_flat = BLAS.matmul(X_flat, gy_cols_T)
			local dW = reshapeInPlace(dW_flat, { C_in, C_out, KH, KW })
			addInto(W._grad :: Types.Tensor, dW)
		end
	end)
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Types = require("../Types")
local Math = {}

local exp = math.exp

local function checkShapes(A: Types.Tensor, B: Types.Tensor, op: string)
	if #A._storage ~= #B._storage then
		error(
			`Math.{op}: Shape mismatch. A has {#A._storage} elements, B has {#B._storage}. Broadcasting is not currently supported.`
		)
	end
	return
end

--@param A Tensor
--@param B Tensor
--@return Tensor
function Math.add(A: Types.Tensor, B: Types.Tensor): Types.Tensor
	checkShapes(A, B, "add")
	task.desynchronize()
	local C = Tensor.empty(A._shape, A._dtype, A._requiresGrad or B._requiresGrad)
	local as, bs, cs = A._storage, B._storage, C._storage
	for i = 1, #as do
		cs[i] = as[i] + bs[i]
	end
	task.synchronize()
	return C
end

--@param A Tensor
--@param B Tensor
--@return Tensor
function Math.sub(A: Types.Tensor, B: Types.Tensor): Types.Tensor
	checkShapes(A, B, "sub")
	task.desynchronize()
	local C = Tensor.empty(A._shape, A._dtype, A._requiresGrad or B._requiresGrad)
	local as, bs, cs = A._storage, B._storage, C._storage
	for i = 1, #as do
		cs[i] = as[i] - bs[i]
	end
	task.synchronize()
	return C
end

--@param A Tensor
--@param B Tensor
--@return Tensor
function Math.mul(A: Types.Tensor, B: Types.Tensor): Types.Tensor
	checkShapes(A, B, "mul")
	task.desynchronize()
	local C = Tensor.empty(A._shape, A._dtype, A._requiresGrad or B._requiresGrad)
	local as, bs, cs = A._storage, B._storage, C._storage
	for i = 1, #as do
		cs[i] = as[i] * bs[i]
	end
	task.synchronize()
	return C
end

--@param A Tensor
--@return Tensor
function Math.exp(A: Types.Tensor): Types.Tensor
	task.desynchronize()
	local C = Tensor.empty(A._shape, A._dtype, A._requiresGrad)
	local as, cs = A._storage, C._storage
	for i = 1, #as do
		cs[i] = exp(as[i])
	end
	task.synchronize()
	return C
end

--@param A Tensor
--@param s number
--@return Tensor
function Math.scalarAdd(A: Types.Tensor, s: number): Types.Tensor
	task.desynchronize()
	local C = Tensor.empty(A._shape, A._dtype, A._requiresGrad)
	local as, cs = A._storage, C._storage
	for i = 1, #as do
		cs[i] = as[i] + s
	end
	task.synchronize()
	return C
end

--@param A Tensor
--@param s number
--@return Tensor
function Math.scalarMul(A: Types.Tensor, s: number): Types.Tensor
	task.desynchronize()
	local C = Tensor.empty(A._shape, A._dtype, A._requiresGrad)
	local as, cs = A._storage, C._storage
	for i = 1, #as do
		cs[i] = as[i] * s
	end
	task.synchronize()
	return C
end

return Math


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

-- MaxPool2d kernel KHxKW stride s; input {C,H,W,B}
--@param X Tensor
--@param KH number
--@param KW number
--@param stride number
--@return Tensor
return function(X: Types.Tensor, KH: number, KW: number, stride: number): Types.Tensor
	local C, H, W, B = X._shape[1], X._shape[2], X._shape[3], X._shape[4]
	local Ho = math.floor((H - KH) / stride) + 1
	local Wo = math.floor((W - KW) / stride) + 1
	local Y = Tensor.zeros({ C, Ho, Wo, B }, X._dtype, X._requiresGrad)
	task.desynchronize()
	local argmax = table.create(C * Ho * Wo * B, 0)
	local HW_B = H * W * B
	local W_B = W * B
	local Wo_B = Wo * B
	local HoWoB = Ho * Wo * B
	local xStorage = X._storage
	local yStorage = Y._storage
	for c = 1, C do
		local c_offset = (c - 1) * HW_B
		local c_out_offset = (c - 1) * HoWoB
		for b = 1, B do
			local b_offset = b - 1
			for i = 1, Ho do
				local i_start = (i - 1) * stride
				local i_out_offset = (i - 1) * Wo_B
				for j = 1, Wo do
					local j_start = (j - 1) * stride
					local best = -math.huge
					local bestXi = 0
					for ki = 0, KH - 1 do
						local row_offset = (i_start + ki) * W_B
						for kj = 0, KW - 1 do
							local col_part = (j_start + kj) * B
							local xi = c_offset + row_offset + col_part + b_offset
							local v = xStorage[xi + 1]
							if v > best then
								best = v
								bestXi = xi
							end
						end
					end
					local yi = c_out_offset + i_out_offset + (j - 1) * B + b_offset + 1
					yStorage[yi] = best
					argmax[yi] = bestXi
				end
			end
		end
	end
	task.synchronize()
	return Tape.makeNode(Y, { X }, function(gy: Types.Tensor)
		if X._requiresGrad then
			X._grad = X._grad or Tensor.zeros(X._shape, X._dtype, false)
			task.desynchronize()
			local xGrad = X._grad :: Types.Tensor
			local xGradStorage = xGrad._storage
			local gyStorage = gy._storage
			local n = #gyStorage
			for idx = 1, n do
				local xi = argmax[idx]
				xGradStorage[xi + 1] += gyStorage[idx]
			end
			task.synchronize()
		end
	end)
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Types = require("../Types")
local Reduce = {}

-- Sum per column (shape {D,B} -> {1,B})
--@param X Tensor
--@return Tensor
function Reduce.sumCols(X: Types.Tensor): Types.Tensor
	task.desynchronize()
	local D, B = X._shape[1], X._shape[2]
	local out = Tensor.empty({ 1, B }, X._dtype, X._requiresGrad)
	local xs, ys = X._storage, out._storage
	for j = 1, B do
		local acc = 0
		local idx = j
		for _ = 1, D do
			acc += xs[idx]
			idx += B
		end
		ys[j] = acc
	end
	task.synchronize()
	return out
end

-- Mean per column (shape {D,B} -> {1,B})
--@param X Tensor
--@return Tensor
function Reduce.meanCols(X: Types.Tensor): Types.Tensor
	task.desynchronize()
	local D, B = X._shape[1], X._shape[2]
	local out = Tensor.empty({ 1, B }, X._dtype, X._requiresGrad)
	local xs, ys = X._storage, out._storage
	local invD = 1 / D
	for j = 1, B do
		local acc = 0
		local idx = j
		for _ = 1, D do
			acc += xs[idx]
			idx += B
		end
		ys[j] = acc * invD
	end
	task.synchronize()
	return out
end

return Reduce


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Types = require("../Types")
local Tape = require("../autograd/Tape")

local Scan = {}

local function ensureGrad(t: Types.Tensor): Types.Tensor?
	if not t._grad then
		t._grad = Tensor.zeros(t._shape, t._dtype, false)
	end
	return t._grad
end

function Scan.selectiveScan(
	u: Types.Tensor,
	delta: Types.Tensor,
	A: Types.Tensor,
	B: Types.Tensor,
	C: Types.Tensor,
	D: Types.Tensor?
): Types.Tensor
	local batch, seqLen, dim = u._shape[1], u._shape[2], u._shape[3]
	local stateDim = A._shape[2] -- N
	local requiresGrad = u._requiresGrad
		or delta._requiresGrad
		or A._requiresGrad
		or B._requiresGrad
		or C._requiresGrad
		or (D and D._requiresGrad)
	local out = Tensor.zeros({ batch, seqLen, dim }, u._dtype, requiresGrad)
	local us = u._storage
	local ds = delta._storage
	local As = A._storage
	local Bs = B._storage
	local Cs = C._storage
	local outs = out._storage
	local Ds = D and D._storage
	local x_storage
	if requiresGrad then
		x_storage = table.create(batch * seqLen * dim * stateDim, 0)
	end
	task.desynchronize()
	for b = 1, batch do
		local x_state = table.create(dim * stateDim, 0)
		local b_offset = (b - 1) * seqLen * dim

		for t = 1, seqLen do
			local t_offset = (t - 1) * dim
			-- B/C: (b-1)*L*N + (t-1)*N
			local bc_base = (b - 1) * seqLen * stateDim + (t - 1) * stateDim
			for d = 1, dim do
				local idx = b_offset + t_offset + d
				local u_val = us[idx]
				local delta_val = ds[idx]
				local y_val = 0
				for n = 1, stateDim do
					local state_idx = (d - 1) * stateDim + n
					local a_val = As[state_idx] -- A is {D, N}
					local b_val = Bs[bc_base + n] -- B is {B, L, N}
					local dA = math.exp(delta_val * a_val)
					local dB = delta_val * b_val
					local prev_x = x_state[state_idx]
					local new_x = dA * prev_x + dB * u_val
					x_state[state_idx] = new_x
					if x_storage then
						local store_idx = ((b - 1) * seqLen + (t - 1)) * dim * stateDim + (d - 1) * stateDim + n
						x_storage[store_idx] = new_x
					end
					local c_val = Cs[bc_base + n]
					y_val += c_val * new_x
				end
				if Ds then
					y_val += Ds[d] * u_val
				end
				outs[idx] = y_val
			end
		end
	end
	task.synchronize()
	if not requiresGrad then
		return out
	end
	local parents = { u, delta, A, B, C }
	if D then
		table.insert(parents, D)
	end
	return Tape.makeNode(out, parents, function(gy)
		local gu = u._requiresGrad and ensureGrad(u)
		local gdelta = delta._requiresGrad and ensureGrad(delta)
		local gA = A._requiresGrad and ensureGrad(A)
		local gB = B._requiresGrad and ensureGrad(B)
		local gC = C._requiresGrad and ensureGrad(C)
		local gD = (D and D._requiresGrad) and ensureGrad(D)
		local gus = gu and gu._storage
		local gds = gdelta and gdelta._storage
		local gAs = gA and gA._storage
		local gBs = gB and gB._storage
		local gCs = gC and gC._storage
		local gDs = gD and gD._storage
		local gys = gy._storage
		task.desynchronize()
		for b = 1, batch do
			local dx_next = table.create(dim * stateDim, 0)
			for t = seqLen, 1, -1 do
				local b_offset = (b - 1) * seqLen * dim
				local t_offset = (t - 1) * dim
				local bc_base = (b - 1) * seqLen * stateDim + (t - 1) * stateDim
				for d = 1, dim do
					local idx = b_offset + t_offset + d
					local gy_val = gys[idx]
					local u_val = us[idx]
					local delta_val = ds[idx]
					local d_delta_acc = 0
					local d_u_acc = 0
					if gDs then
						gDs[d] += gy_val * u_val
						d_u_acc += gy_val * Ds[d]
					elseif Ds then
						d_u_acc += gy_val * Ds[d]
					end
					for n = 1, stateDim do
						local state_idx = (d - 1) * stateDim + n
						local a_val = As[state_idx]
						local b_val = Bs[bc_base + n]
						local c_val = Cs[bc_base + n]
						local store_idx = ((b - 1) * seqLen + (t - 1)) * dim * stateDim + (d - 1) * stateDim + n
						local x_curr = x_storage[store_idx]
						local x_prev = 0
						if t > 1 then
							local prev_store_idx = ((b - 1) * seqLen + (t - 2)) * dim * stateDim
								+ (d - 1) * stateDim
								+ n
							x_prev = x_storage[prev_store_idx]
						end
						local dA = math.exp(delta_val * a_val)
						local dB = delta_val * b_val
						local d_x_curr = gy_val * c_val + dx_next[state_idx]
						if gCs then
							gCs[bc_base + n] += gy_val * x_curr
						end
						dx_next[state_idx] = d_x_curr * dA
						local d_dA = d_x_curr * x_prev
						local d_argA = d_dA * dA -- deriv of exp is exp
						if gAs then
							gAs[state_idx] += d_argA * delta_val
						end
						d_delta_acc += d_argA * a_val
						local d_dB = d_x_curr * u_val
						if gBs then
							gBs[bc_base + n] += d_dB * delta_val
						end
						d_delta_acc += d_dB * b_val
						-- u: dL/du += dL/dx_curr * dB
						d_u_acc += d_x_curr * dB
					end
					if gds then
						gds[idx] = d_delta_acc
					end
					if gus then
						gus[idx] = d_u_acc
					end
				end
			end
		end
		task.synchronize()
	end)
end

return Scan


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")

local SoftmaxOps = {}

local function stableSoftmax(logits: Types.Tensor, out: Types.Tensor)
	local src = logits._storage
	local dst = out._storage
	local C, B = logits._shape[1], logits._shape[2]
	task.desynchronize()
	for j = 1, B do
		local colStart = (1 - 1) * B + j
		local maxVal = -math.huge
		local idx = colStart -- pointer
		for i = 1, C do
			local v = src[idx]
			if v > maxVal then
				maxVal = v
			end
			idx += B
		end
		local sum = 0
		idx = colStart
		for i = 1, C do
			local e = math.exp(src[idx] - maxVal)
			dst[idx] = e
			sum += e
			idx += B
		end
		-- normalize
		local invSum = 1 / sum
		idx = colStart
		for i = 1, C do
			dst[idx] *= invSum
			idx += B
		end
	end
	task.synchronize()
end

function SoftmaxOps.forward(logits: Types.Tensor): Types.Tensor
	local P = Tensor.zeros(logits._shape, logits._dtype, logits._requiresGrad)
	stableSoftmax(logits, P)
	if not logits._requiresGrad then
		return P
	end
	return Tape.makeNode(P, { logits }, function(gy: Types.Tensor)
		if not logits._requiresGrad then
			return
		end
		if not logits._grad then
			logits._grad = Tensor.zeros(logits._shape, logits._dtype, false)
		end
		task.desynchronize()
		local grads = logits._grad._storage
		local probs = P._storage
		local gys = gy._storage
		local C, B = logits._shape[1], logits._shape[2]
		-- dx_i = y_i * (dy_i - dot(y, dy))
		for j = 1, B do
			local colStart = (1 - 1) * B + j
			local dot = 0
			local idx = colStart
			for i = 1, C do
				dot += probs[idx] * gys[idx]
				idx += B
			end
			idx = colStart
			for i = 1, C do
				grads[idx] += probs[idx] * (gys[idx] - dot)
				idx += B
			end
		end
		task.synchronize()
	end)
end

return SoftmaxOps


--!native
--!optimize 2
local Util = require("../Util")
local Types = require("../Types")
local assert = Util.Assert

local function scaleGrads(params: { Types.Tensor }, factor: number)
	task.desynchronize()
	for _, p in params do
		local g = p._grad
		if g then
			local gs = g._storage
			for i = 1, #gs do
				gs[i] *= factor
			end
		end
	end
	task.synchronize()
end

return function(opt: Types.Optimizer, steps: number, params: { Types.Tensor }?, normalize: boolean?)
	assert(steps >= 1, "Accumulated: steps must be >= 1")
	normalize = normalize == nil and true or normalize
	local counter: number = 0
	local invSteps: number = 1 / steps
	local doScale: boolean = normalize and params ~= nil
	return {
		step = function(_)
			counter += 1
			if counter < steps then
				return
			end
			if doScale then
				scaleGrads(params, invSteps)
			end
			opt:step()
			opt:zeroGrad()
			counter = 0
		end,
		flush = function(_)
			if counter == 0 then
				return
			end
			if doScale then
				scaleGrads(params, 1 / counter)
			end
			opt:step()
			opt:zeroGrad()
			counter = 0
		end,
		zeroGrad = function(_)
			counter = 0
			opt:zeroGrad()
		end,
		setLr = function(_, newLr: number)
			if opt.setLr then
				opt:setLr(newLr)
			end
		end,
		getState = function(_)
			local innerState = opt.getState and opt:getState() or nil
			return {
				steps = steps,
				counter = counter,
				normalize = normalize,
				optState = innerState,
			}
		end,
		setState = function(
			_,
			state: {
				steps: number,
				counter: number,
				normalize: boolean,
				optState: { any },
			}
		)
			if state.steps ~= nil then
				steps = state.steps
				invSteps = 1 / steps
			end
			if state.counter ~= nil then
				counter = state.counter
			end
			if state.normalize ~= nil then
				normalize = state.normalize
				doScale = normalize and params ~= nil
			end
			if state.optState ~= nil and opt.setState then
				opt:setState(state.optState)
			end
		end,
	}
end


--!native
--!optimize 2
local Types = require("../Types")

local function flatShape(shape: { number }): (number, number)
	local rank = #shape
	local rows = 1
	for i = 1, rank - 1 do
		rows *= shape[i]
	end
	return rows, shape[rank]
end

local function ensureZeros(buf: { number }?, size: number): { number }
	if not buf or #buf ~= size then
		return table.create(size, 0)
	end
	for i = 1, size do
		buf[i] = 0
	end
	return buf
end

local function ensureBuffer(buf: { number }?, size: number): { number }
	if not buf or #buf ~= size then
		return table.create(size, 0)
	end
	return buf
end

local function clipScaleOf(updates: { number }, count: number)
	if clipThreshold == 0 then
		return 1
	end
	local acc = 0
	for i = 1, count do
		local u = updates[i]
		acc += u * u
	end
	local rms = math.sqrt(acc / count)
	if rms > clipThreshold then
		return clipThreshold / rms
	end
	return 1
end

return function(
	params: { Types.Tensor },
	lr: number,
	beta2: number?,
	eps: number?,
	clipThreshold: number?,
	weightDecay: number?
)
	beta2 = beta2 or 0.999
	eps = eps or 1e-30
	clipThreshold = clipThreshold or 0
	weightDecay = weightDecay or 0

	local factored: {
		[Types.Tensor]: {
			row: { number },
			col: { number },
			rowMean: { number },
			colMean: { number },
			updates: { number },
		},
	} =
		{}
	local full: { [Types.Tensor]: { v: { number }, updates: { number } } } = {}
	local self = {}

	function self:step()
		task.desynchronize()
		for _, p in params do
			local g = p._grad
			if g then
				local ps, gs = p._storage, g._storage
				local count = #ps
				local shape = p._shape
				local rank = #shape
				if rank >= 2 then
					local rows, cols = flatShape(shape)
					local s = factored[p]
					if not s then
						s = {
							row = table.create(rows, 0),
							col = table.create(cols, 0),
							rowMean = table.create(rows, 0),
							colMean = table.create(cols, 0),
							updates = table.create(count, 0),
						}
						factored[p] = s
					elseif #s.row ~= rows or #s.col ~= cols then
						-- reshape
						s.row = table.create(rows, 0)
						s.col = table.create(cols, 0)
						s.rowMean = table.create(rows, 0)
						s.colMean = table.create(cols, 0)
						s.updates = table.create(rows * cols, 0)
					end
					local rowMean = ensureZeros(s.rowMean, rows)
					s.rowMean = rowMean
					local colMean = ensureZeros(s.colMean, cols)
					s.colMean = colMean
					local updates = ensureBuffer(s.updates, count)
					s.updates = updates
					for i = 1, count do
						local g2 = gs[i] * gs[i]
						local r = math.floor((i - 1) / cols) + 1
						local c = ((i - 1) % cols) + 1
						rowMean[r] += g2
						colMean[c] += g2
					end
					local invCols = 1 / cols
					local invRows = 1 / rows
					for r = 1, rows do
						rowMean[r] *= invCols
					end
					for c = 1, cols do
						colMean[c] *= invRows
					end
					local rowState = s.row
					local colState = s.col
					for r = 1, rows do
						rowState[r] = (beta2 :: number) * rowState[r] + (1 - (beta2 :: number)) * rowMean[r]
					end
					for c = 1, cols do
						colState[c] = (beta2 :: number) * colState[c] + (1 - (beta2 :: number)) * colMean[c]
					end
					local meanRow = 0
					for r = 1, rows do
						meanRow += rowState[r]
					end
					meanRow /= rows
					for i = 1, count do
						local r = math.floor((i - 1) / cols) + 1
						local c = ((i - 1) % cols) + 1
						local vHat = (rowState[r] * colState[c]) / (meanRow + (eps :: number))
						updates[i] = gs[i] / (math.sqrt(vHat) + (eps :: number))
					end
					local scale = clipScaleOf(updates, count)
					for i = 1, count do
						ps[i] -= lr * (updates[i] * scale + (weightDecay :: number) * ps[i])
					end
				else
					local v = full[p]
					if not v then
						v = { v = table.create(count, 0), updates = table.create(count, 0) }
						full[p] = v
					elseif #v.v ~= count then
						v.v = table.create(count, 0)
						v.updates = table.create(count, 0)
					end
					local updates = ensureBuffer(v.updates, count)
					v.updates = updates
					local vv = v.v
					for i = 1, count do
						local gval = gs[i]
						vv[i] = (beta2 :: number) * vv[i] + (1 - (beta2 :: number)) * gval * gval
						updates[i] = gval / (math.sqrt(vv[i]) + (eps :: number))
					end
					local scale = clipScaleOf(updates, count)
					for i = 1, count do
						ps[i] -= lr * (updates[i] * scale + (weightDecay :: number) * ps[i])
					end
				end
			end
		end
		task.synchronize()
	end
	function self:zeroGrad()
		for _, p in params do
			p._grad = nil
		end
	end
	function self:setLr(newLr: number)
		lr = newLr
	end
	function self:getState()
		return {
			lr = lr,
			beta2 = beta2,
			eps = eps,
			clipThreshold = clipThreshold,
			weightDecay = weightDecay,
			factored = factored,
			full = full,
		}
	end
	function self:setState(state: {
		lr: number,
		beta2: number,
		eps: number,
		clipThreshold: number,
		weightDecay: number,
		factored: { [Types.Tensor]: { row: { number }, col: { number }, rowMean: { number }, colMean: { number }, updates: { number } } },
		full: { [Types.Tensor]: { v: { number }, updates: { number } } },
	})
		if state.lr ~= nil then
			lr = state.lr
		end
		if state.beta2 ~= nil then
			beta2 = state.beta2
		end
		if state.eps ~= nil then
			eps = state.eps
		end
		if state.clipThreshold ~= nil then
			clipThreshold = state.clipThreshold
		end
		if state.weightDecay ~= nil then
			weightDecay = state.weightDecay
		end
		if state.factored ~= nil then
			factored = state.factored
		end
		if state.full ~= nil then
			full = state.full
		end
	end
	return self
end


--!native
--!optimize 2
export type Optimizer = { step: (self: Optimizer) -> (), zeroGrad: (self: Optimizer) -> () }

return function(params: { any }, lr: number, eps: number?)
	eps = eps or 1e-8
	local G: { [any]: { [number]: number } } = {}
	local self: Optimizer = {
		step = function(_)
			task.desynchronize()
			for _, p in params do
				if p._grad then
					local ps, gs = p._storage, p._grad._storage
					local Gi = G[p]
					if not Gi then
						Gi = table.create(#ps, 0)
						G[p] = Gi
					end
					for i = 1, #ps do
						Gi[i] += gs[i] * gs[i]
						ps[i] -= lr * gs[i] / (math.sqrt(Gi[i]) + (eps :: number))
					end
				end
			end
			task.synchronize()
		end,
		zeroGrad = function(_)
			for _, p in params do
				p._grad = nil
			end
		end,
	}
	return self
end


--!native
--!optimize 2
local Types = require("../Types")

return function(params: { Types.Tensor }, lr: number, b1: number?, b2: number?, eps: number?, weight_decay: number?)
	b1 = b1 or 0.9
	b2 = b2 or 0.999
	eps = eps or 1e-8
	local wd = weight_decay or 0
	local m: { [Types.Tensor]: { [number]: number } } = {}
	local v: { [Types.Tensor]: { [number]: number } } = {}
	local t = 0
	local self = {
		step = function(_)
			t += 1
			local b1t = (b1 :: number) ^ t
			local b2t = (b2 :: number) ^ t
			task.desynchronize()
			for _, p in params do
				if p._grad then
					local ps, gs = p._storage, p._grad._storage
					local mi = m[p]
					if not mi then
						mi = table.create(#ps, 0)
						m[p] = mi
					end
					local vi = v[p]
					if not vi then
						vi = table.create(#ps, 0)
						v[p] = vi
					end
					for i = 1, #ps do
						local g = gs[i]
						if wd ~= 0 then
							g = g + wd * ps[i]
						end
						mi[i] = (b1 :: number) * mi[i] + (1 - (b1 :: number)) * g
						vi[i] = (b2 :: number) * vi[i] + (1 - (b2 :: number)) * g * g
						local mHat = mi[i] / (1 - b1t)
						local vHat = vi[i] / (1 - b2t)
						ps[i] -= lr * (mHat / (math.sqrt(vHat) + (eps :: number)))
					end
				end
			end
			task.synchronize()
		end,
		zeroGrad = function(_)
			for _, p in params do
				p._grad = nil
			end
		end,
		setLr = function(_, newLr: number)
			lr = newLr
		end,
		getState = function(): {
			lr: number,
			b1: number,
			b2: number,
			eps: number,
			t: number,
			wd: number,
			m: { [Types.Tensor]: { [number]: number } },
			v: { [Types.Tensor]: { [number]: number } },
		}
			return { lr = lr, b1 = b1, b2 = b2, eps = eps, t = t, wd = wd, m = m, v = v }
		end,
		setState = function(
			_,
			state: {
				lr: number,
				b1: number,
				b2: number,
				eps: number,
				t: number,
				wd: number,
				m: { [Types.Tensor]: { [number]: number } },
				v: { [Types.Tensor]: { [number]: number } },
			}
		)
			lr = state.lr or lr
			b1 = state.b1 or b1
			b2 = state.b2 or b2
			eps = state.eps or eps
			t = state.t or t
			wd = state.wd or wd
			m = state.m or m
			v = state.v or v
		end,
	}
	return self
end


--!native
--!optimize 2
local Types = require("../Types")

return function(
	params: { Types.Tensor },
	lr: number,
	wd: number,
	b1: number?,
	b2: number?,
	eps: number?
)
	b1 = b1 or 0.9
	b2 = b2 or 0.999
	eps = eps or 1e-8
	local m: { [Types.Tensor]: { [number]: number } } = {}
	local v: { [Types.Tensor]: { [number]: number } } = {}
	local t = 0
	local self = {
		step = function(_)
			t += 1
			local b1t = (b1 :: number) ^ t
			local b2t = (b2 :: number) ^ t
			task.desynchronize()
			for _, p in params do
				if p._grad then
					local ps, gs = p._storage, p._grad._storage
					local mi = m[p]
					if not mi then
						mi = table.create(#ps, 0)
						m[p] = mi
					end
					local vi = v[p]
					if not vi then
						vi = table.create(#ps, 0)
						v[p] = vi
					end
					for i = 1, #ps do
						local g = gs[i]
						mi[i] = (b1 :: number) * mi[i] + (1 - (b1 :: number)) * g
						vi[i] = (b2 :: number) * vi[i] + (1 - (b2 :: number)) * g * g
						local mHat = mi[i] / (1 - b1t)
						local vHat = vi[i] / (1 - b2t)
						ps[i] -= lr * (mHat / (math.sqrt(vHat) + (eps :: number)) + wd * ps[i])
					end
				end
			end
			task.synchronize()
		end,
		zeroGrad = function(_)
			for _, p in params do
				p._grad = nil
			end
		end,
		getState = function()
			return {
				lr = lr,
				wd = wd,
				b1 = b1,
				b2 = b2,
				eps = eps,
				t = t,
			}
		end,
		setState = function(_, state: { lr: number, wd: number, b1: number, b2: number, eps: number, t: number })
			if state.lr ~= nil then
				lr = state.lr
			end
			if state.wd ~= nil then
				wd = state.wd
			end
			if state.b1 ~= nil then
				b1 = state.b1
			end
			if state.b2 ~= nil then
				b2 = state.b2
			end
			if state.eps ~= nil then
				eps = state.eps
			end
			if state.t ~= nil then
				t = state.t
			end
		end,
	}
	return self
end


--!native
--!optimize 2
local Tensor = require("../Tensor")

local function update(self, params)
	task.desynchronize()
	local m = self.decay
	for i = 1, #params do
		local t = params[i]
		local s = self.shadow[i]
		for k = 1, #t._storage do
			s._storage[k] = m * s._storage[k] + (1 - m) * t._storage[k]
		end
	end
	task.synchronize()
end

local function apply(self, params)
	self.backup = {}
	for i = 1, #params do
		local t = params[i]
		local b = Tensor.zeros(t._shape, t._dtype, false)
		for k = 1, #t._storage do
			b._storage[k] = t._storage[k]
			t._storage[k] = self.shadow[i]._storage[k]
		end
		self.backup[i] = b
	end
end

local function restore(self, params)
	if not self.backup then
		return
	end
	for i = 1, #params do
		local t = params[i]
		for k = 1, #t._storage do
			t._storage[k] = self.backup[i]._storage[k]
		end
	end
	self.backup = nil
end

local function new(params, decay: number?)
	local shadow = {}
	for i = 1, #params do
		local t = params[i]
		local buf = Tensor.zeros(t._shape, t._dtype, false)
		for k = 1, #t._storage do
			buf._storage[k] = t._storage[k]
		end
		shadow[i] = buf
	end
	return {
		decay = decay or 0.999,
		shadow = shadow,
		backup = nil,
		update = update,
		apply = apply,
		restore = restore,
	}
end

return function(params, decay)
	return new(params, decay)
end


--!native
--!optimize 2
-- Chen et al., 2023. Update uses sign of momentum (no per-param v).

--@param params {any}
--@param lr number
--@param beta1 number? @default 0.9
--@param beta2 number? @default 0.99
--@param weightDecay number? @default 0.0 (decoupled)
return function(params, lr: number, beta1: number?, beta2: number?, weightDecay: number?)
	local state = { m = {} }
	beta1, beta2 = beta1 or 0.9, beta2 or 0.99
	weightDecay = weightDecay or 0.0
	for i = 1, #params do
		local t = params[i]
		local m = table.create(#t._storage, 0.0)
		state.m[i] = m
	end
	local opt = {}
	opt.lr = lr
	function opt:zeroGrad()
		for i = 1, #params do
			local g = params[i]._grad
			if g then task.desynchronize(); for k = 1, #g._storage do g._storage[k] = 0 end; task.synchronize() end
		end
	end
	function opt:step()
		task.desynchronize()
		for i = 1, #params do
			local p = params[i]
			local g = p._grad
			if g then
				local mi = state.m[i]
				for k = 1, #p._storage do
					if weightDecay ~= 0 then
						p._storage[k] -= opt.lr * weightDecay * p._storage[k]
					end
					-- momentum
					mi[k] = beta1 * mi[k] + (1 - beta1) * g._storage[k]
					-- lookahead step
					local update = math.sign(mi[k]) * opt.lr
					p._storage[k] -= update
					-- second momentum blend
					mi[k] = beta2 * mi[k] + (1 - beta2) * g._storage[k]
				end
			end
		end
		task.synchronize()
	end
	return opt
end


--!native
--!optimize 2

--@param params { Tensor }
--@param base any -- base optimizer already constructed on params
--@param k integer? -- steps between sync (default 5)
--@param alpha number? -- interpolation factor (default 0.5)
return function(params, base, k: number?, alpha: number?)
	k = k or 5; alpha = alpha or 0.5
	local step = 0
	local slow = {}
	for i = 1, #params do
		local t = params[i]
		local buf = table.create(#t._storage, 0.0)
		for j = 1, #t._storage do buf[j] = t._storage[j] end
		slow[i] = buf
	end

	local opt = {}
	opt.base = base
	function opt:zeroGrad()
		base:zeroGrad()
	end
	function opt:step()
		base:step()
		step += 1
		if step % k == 0 then
			task.desynchronize()
			for i = 1, #params do
				local p = params[i]
				local s = slow[i]
				for j = 1, #p._storage do
					s[j] = s[j] + alpha * (p._storage[j] - s[j])
					p._storage[j] = s[j]
				end
			end
			task.synchronize()
		end
	end
	return opt
end


--!native
--!optimize 2
local Types = require("../Types")
local Tensor = require("../Tensor")
local BLAS = require("../ops/BLAS")

-- X_{k+1} = 1.5 * X_k - 0.5 * X_k @ X_k^T @ X_k
local function newtonSchulz(G: Types.Tensor, steps: number): Types.Tensor
	local X = Tensor.detach(G)
	local _rows, _cols = X._shape[1], X._shape[2]
	task.desynchronize()
	local sumSq = 0
	local xs, xs_len = X._storage, #X._storage
	for i = 1, xs_len do
		sumSq += xs[i] * xs[i]
	end
	local frob = math.sqrt(sumSq)
	if frob > 1e-6 then
		for i = 1, xs_len do
			xs[i] /= frob
		end
	end
	task.synchronize()
	for k = 1, steps do
		-- A = X @ X.T
		-- B = A @ X
		-- X = 1.5 * X - 0.5 * B
		local XT = Tensor.transpose(X)
		local A = BLAS.matmul(X, XT)
		local B = BLAS.matmul(A, X)
		local bs = B._storage
		local xs_curr = X._storage
		task.desynchronize()
		for i = 1, #xs_curr do
			xs_curr[i] = 1.5 * xs_curr[i] - 0.5 * bs[i]
		end
		task.synchronize()
	end
	return X
end

return function(params: { Types.Tensor }, lr: number, momentum: number?, nesterov: boolean?, nsSteps: number?)
	momentum = momentum or 0.95
	nsSteps = nsSteps or 3
	local state: { [Types.Tensor]: Types.Tensor } = {}
	local step = 0
	local self = {
		step = function(_)
			step += 1
			for _, p in params do
				if p._grad then
					local g = p._grad
					local m = state[p]
					if not m then
						m = Tensor.zeros(p._shape, p._dtype, false)
						state[p] = m
					end
					task.desynchronize()
					local ms, gs = m._storage, g._storage
					for i = 1, #ms do
						ms[i] = (momentum :: number) * ms[i] + gs[i]
					end
					task.synchronize()
					local rank = #p._shape
					if rank == 2 then
						local orth = newtonSchulz(m, nsSteps :: number)
						local sumW = 0
						local ps, ps_len = p._storage, #p._storage
						local os = orth._storage
						task.desynchronize()
						for i = 1, ps_len do
							sumW += ps[i] * ps[i]
						end
						local rmsW = math.sqrt(sumW / ps_len)
						local scale = math.max(1, rmsW)
						for i = 1, ps_len do
							ps[i] -= lr * scale * os[i]
						end
						task.synchronize()
					else
						task.desynchronize()
						local ps = p._storage
						for i = 1, #ps do
							ps[i] -= lr * ms[i]
						end
						task.synchronize()
					end
				end
			end
		end,
		zeroGrad = function(_)
			for _, p in params do
				p._grad = nil
			end
		end,
		getState = function()
			return {
				lr = lr,
				momentum = momentum,
				nsSteps = nsSteps,
				step = step,
			}
		end,
	}
	return self
end


--!native
--!optimize 2
export type Optimizer = { step: (self: Optimizer) -> (), zeroGrad: (self: Optimizer) -> () }

return function(params: { any }, lr: number, decay: number?, eps: number?)
	decay = decay or 0.99
	eps = eps or 1e-8
	local avg2: { [any]: { [number]: number } } = {}
	local self: Optimizer = {
		step = function(_)
			task.desynchronize()
			for _, p in params do
				if p._grad then
					local ps, gs = p._storage, p._grad._storage
					local a2 = avg2[p]
					if not a2 then
						a2 = table.create(#ps, 0)
						avg2[p] = a2
					end
					for i = 1, #ps do
						a2[i] = (decay :: number) * a2[i] + (1 - (decay :: number)) * gs[i] * gs[i]
						ps[i] -= lr * gs[i] / (math.sqrt(a2[i]) + (eps :: number))
					end
				end
			end
			task.synchronize()
		end,
		zeroGrad = function(_)
			for _, p in params do
				p._grad = nil
			end
		end,
	}
	return self
end


--!native
--!optimize 2
local Types = require("../Types")

return function(params: { Types.Tensor }, lr: number, momentum: number?, nesterov: boolean?)
	momentum = momentum or 0
	local vel: { [Types.Tensor]: { [number]: number } } = {}
	local step = 0
	local self = {
		step = function(_)
			step += 1
			task.desynchronize()
			for _, p in params do
				if p._grad then
					local ps, gs = p._storage, p._grad._storage
					local v = vel[p]
					if not v then
						v = table.create(#ps, 0)
						vel[p] = v
					end
					for i = 1, #ps do
						v[i] = (momentum :: number) * v[i] + gs[i]
						local upd = nesterov and ((momentum :: number) * v[i] + gs[i]) or v[i]
						ps[i] -= lr * upd
					end
				end
			end
			task.synchronize()
		end,
		zeroGrad = function(_)
			for _, p in params do
				p._grad = nil
			end
		end,
		setLr = function(_, newLr: number)
			lr = newLr
		end,
		getState = function()
			return {
				lr = lr,
				momentum = momentum,
				step = step,
				velocity = vel,
			}
		end,
		setState = function(
			_,
			state: {
				lr: number,
				momentum: number,
				step: number,
				velocity: { [Types.Tensor]: { [number]: number } },
			}
		)
			if state.lr ~= nil then
				lr = state.lr
			end
			if state.momentum ~= nil then
				momentum = state.momentum
			end
			if state.step ~= nil then
				step = state.step
			end
			if state.velocity ~= nil then
				vel = state.velocity
			end
		end,
	}
	return self
end


--!native
--!optimize 2
local S = {}

function S.step(lr: number, gamma: number, stepSize: number)
	return function(step: number)
		return lr * (gamma ^ math.floor(step / stepSize))
	end
end

function S.cosine(lr: number, lrMin: number, T: number)
	return function(step: number)
		local t = math.min(step / T, 1)
		return lrMin + 0.5 * (lr - lrMin) * (1 + math.cos(math.pi * t))
	end
end

function S.warmupCosine(lr: number, warmup: number, T: number)
	local base = S.cosine(lr, 0, T)
	return function(step: number)
		if step < warmup then
			return lr * (step / warmup)
		end
		return base(step - warmup)
	end
end

function S.linearWarmupThenCosine(lr, warmupSteps, totalSteps, lrMin)
	return function(step)
		if step < warmupSteps then
			return lr * (step / math.max(1, warmupSteps))
		end
		local t = (step - warmupSteps) / math.max(1, totalSteps - warmupSteps)
		return lrMin + 0.5 * (lr - lrMin) * (1 + math.cos(math.pi * t))
	end
end

function S.oneCycle(lrMax: number, lrMin: number, T_up: number, T_total: number)
	return function(step: number)
		if step < T_up then
			return lrMax * (step / math.max(1, T_up))
		end
		local t = math.min(math.max(step - T_up, 0) / math.max(1, T_total - T_up), 1)
		return lrMin + 0.5 * (lrMax - lrMin) * (1 + math.cos(math.pi * t))
	end
end

function S.cosineRestarts(lr: number, T0: number, Tmult: number)
	local Tcur = T0
	local base = 0
	return function(step: number)
		local t = step - base
		task.desynchronize()
		while t >= Tcur do
			t -= Tcur
			base += Tcur
			Tcur *= (Tmult > 0 and Tmult or 2)
		end
		task.synchronize()
		local frac = t / math.max(1, math.floor(Tcur))
		return 0.5 * lr * (1 + math.cos(math.pi * frac))
	end
end

return S


--!native
--!optimize 2
local Types = require("../Types")

--@param params {Tensor}
--@param lr number
--@param betas {number}? @default {0.965, 0.99}
--@param rho number? @default 0.04
--@param weightDecay number? @default 0.0
return function(params: { Types.Tensor }, lr: number, betas: { number }?, rho: number?, weightDecay: number?)
	local b1 = (betas and betas[1]) or 0.965
	local b2 = (betas and betas[2]) or 0.99
	local rho_val = rho or 0.04
	local wd = weightDecay or 0.0
	local m: { [Types.Tensor]: { number } } = {}
	local h: { [Types.Tensor]: { number } } = {}
	local step = 0
	return {
		step = function(_)
			step += 1
			task.desynchronize()
			for _, p in params do
				if p._grad then
					local ps, gs = p._storage, p._grad._storage
					local ps_len = #ps
					if not m[p] then
						m[p] = table.create(ps_len, 0)
						h[p] = table.create(ps_len, 0)
					end
					local mi = m[p]
					local hi = h[p]
					for i = 1, ps_len do
						local g = gs[i]
						if wd ~= 0 then
							ps[i] -= lr * wd * ps[i]
						end
						local new_hi, new_mi = b2 * hi[i] + (1 - b2) * (g * g), b1 * mi[i] + (1 - b1) * g
						hi[i] = new_hi
						mi[i] = new_mi
						local hessian = math.max(math.sqrt(new_hi), rho_val)
						ps[i] -= lr * (new_mi / hessian)
					end
				end
			end
			task.synchronize()
		end,
		zeroGrad = function(_)
			for _, p in params do
				p._grad = nil
			end
		end,
		getState = function(): {
			lr: number,
			b1: number,
			b2: number,
			rho: number,
			step: number,
			m: { [Types.Tensor]: { number } },
			h: { [Types.Tensor]: { number } },
		}
			return { lr = lr, b1 = b1, b2 = b2, rho = rho_val, step = step, m = m, h = h }
		end,
		setState = function(
			_,
			state: {
				lr: number,
				b1: number,
				b2: number,
				rho: number,
				step: number,
				m: { [Types.Tensor]: { number } },
				h: { [Types.Tensor]: { number } },
			}
		)
			lr = state.lr or lr
			b1 = state.b1 or b1
			b2 = state.b2 or b2
			rho_val = state.rho or rho_val
			step = state.step or step
			m = state.m or m
			h = state.h or h
		end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Util = require("../Util")
local assert = Util.Assert

return function(a: number, b: number)
	a = a or 0
	b = b or 1
	local lo: { number }? = nil
	local hi: { number }? = nil
	local self
	self = {
		fit = function(_, X)
			task.desynchronize()
			local D, N = X._shape[1], X._shape[2]
			lo = table.create(D, 1 / 0)
			hi = table.create(D, -1 / 0)
			for d = 1, D do
				for n = 1, N do
					local v = X._storage[(d - 1) * N + n]
					if v < (lo :: { number })[d] then
						(lo :: { number })[d] = v
					end
					if v > (hi :: { number })[d] then
						(hi :: { number })[d] = v
					end
				end
			end
			task.synchronize()
		end,
		transform = function(_, X)
			assert(lo and hi, "MinMaxScaler: fit first")
			task.desynchronize()
			local D, N = X._shape[1], X._shape[2]
			local Y = Tensor.zeros(X._shape)
			for d = 1, D do
				local l = (lo :: { number })[d]
				local h = (hi :: { number })[d]
				local range = (h - l ~= 0) and (h - l) or 1
				for n = 1, N do
					Y._storage[(d - 1) * N + n] = a + (X._storage[(d - 1) * N + n] - l) * (b - a) / range
				end
			end
			task.synchronize()
			return Y
		end,
	}
	return self
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
return function(C: number)
	return function(idx: { number })
		task.desynchronize()
		local B = #idx
		local Y = Tensor.zeros({ C, B })
		for j = 1, B do
			Y._storage[(idx[j] - 1) * B + j] = 1
		end
		task.synchronize()
		return Y
	end
end


--!native
--!optimize 2
local Tensor = require("../Tensor")

local function matvec(D: number, C: Types.Tensor, vec: {number}): {number} -- C * v
	local out = table.create(D, 0.0)
	for r = 1, D do
		local s = 0.0
		for c = 1, D do
			s += C._storage[(r - 1) * D + c] * vec[c]
		end
		out[r] = s
	end
	return out
end

local function normalize(D: number, v: {number})
	local n2 = 0.0
	for i = 1, D do
		n2 += v[i] * v[i]
	end
	local inv = 1 / math.sqrt(math.max(n2, 1e-12))
	for i = 1, D do
		v[i] *= inv
	end
end

local function transform(self, X)
	task.desynchronize()
	local D, B = X._shape[1], X._shape[2]
	local K = self.k
	local Y = Tensor.zeros({ K, B }, X._dtype, X._requiresGrad)
	for j = 1, B do
		for k = 1, K do
			local s = 0.0
			for i = 1, D do
				local xij = X._storage[(i - 1) * B + j]
				local mu_i = self.mean._storage[i]
				local w = self.components._storage[(i - 1) * K + k]
				s += (xij - mu_i) * w
			end
			Y._storage[(k - 1) * B + j] = s
		end
	end
	task.synchronize()
	return Y
end

local function new(X, K: number, iters: number?)
	local D, N = X._shape[1], X._shape[2]
	iters = iters or 50

	task.desynchronize()
	local mu = Tensor.zeros({ D, 1 }, X._dtype, false)
	for i = 1, D do
		local s = 0.0
		for j = 1, N do
			s += X._storage[(i - 1) * N + j]
		end
		mu._storage[i] = s / N
	end

	local Xc = Tensor.zeros({ D, N }, X._dtype, false)
	for i = 1, D do
		local m = mu._storage[i]
		for j = 1, N do
			Xc._storage[(i - 1) * N + j] = X._storage[(i - 1) * N + j] - m
		end
	end

	local C = Tensor.zeros({ D, D }, X._dtype, false)
	for i = 1, D do
		for k = i, D do
			local s = 0.0
			for j = 1, N do
				s += Xc._storage[(i - 1) * N + j] * Xc._storage[(k - 1) * N + j]
			end
			s /= N
			C._storage[(i - 1) * D + k] = s
			C._storage[(k - 1) * D + i] = s
		end
	end

	local comps = Tensor.zeros({ D, K }, X._dtype, false)

	for k = 1, K do
		local v = table.create(D, 0.0)
		for i = 1, D do
			v[i] = math.random() * 2 - 1
		end
		normalize(D, v)

		for _ = 1, iters do
			local u = matvec(D, C, v)
			normalize(D, u)
			v = u
		end

		for i = 1, D do
			comps._storage[(i - 1) * K + k] = v[i]
		end

		local Cv = matvec(D, C, v)
		local num = 0.0
		for i = 1, D do
			num += v[i] * Cv[i]
		end
		local lambda = num
		for i = 1, D do
			for j = 1, D do
				C._storage[(i - 1) * D + j] -= lambda * v[i] * v[j]
			end
		end
	end
	task.synchronize()

	return {
		mean = mu,
		components = comps,
		k = K,
		transform = transform,
	}
end

return new


--!native
--!optimize 2
local Tensor = require("../Tensor")

--@param x Tensor -- {D, B*T}
--@param T number -- sequence length
--@return Tensor -- new tensor {D, B*T} with PE added
return function(x, T: number)
	local D, BT = x._shape[1], x._shape[2]
	local B = math.floor(BT / T)
	local out = Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
	task.desynchronize()
	local pe = {}
	for pos = 1, T do
		pe[pos] = table.create(D, 0.0)
		for i = 1, D, 2 do
			local div = math.exp(((i - 1) / D) * math.log(10000))
			pe[pos][i] = math.sin((pos - 1) / div)
			if i + 1 <= D then
				pe[pos][i + 1] = math.cos((pos - 1) / div)
			end
		end
	end
	for b = 1, B do
		local base = (b - 1) * T
		for t = 1, T do
			for i = 1, D do
				local col = base + t
				out._storage[(i - 1) * BT + col] = x._storage[(i - 1) * BT + col] + pe[t][i]
			end
		end
	end
	task.synchronize()
	return out
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Util = require("../Util")
local assert = Util.Assert

return function()
	local mu: { number } = nil
	local sigma: { number } = nil
	local self
	self = {
		fit = function(_, X)
			task.desynchronize()
			local D, N = X._shape[1], X._shape[2]
			mu = table.create(D, 0)
			sigma = table.create(D, 1)
			for d = 1, D do
				local s, s2 = 0, 0
				for n = 1, N do
					local v = X._storage[(d - 1) * N + n]
					s += v
					s2 += v * v
				end
				local m = s / N
				local v = math.max(s2 / N - m * m, 1e-9)
				mu[d] = m
				sigma[d] = math.sqrt(v)
			end
			task.synchronize()
		end,
		transform = function(_, X)
			assert(mu and sigma, "StandardScaler: fit first")
			task.desynchronize()
			local D, N = X._shape[1], X._shape[2]
			local Y = Tensor.zeros(X._shape)
			for d = 1, D do
				local m = (mu :: { number })[d]
				local s = (sigma :: { number })[d]
				for n = 1, N do
					Y._storage[(d - 1) * N + n] = (X._storage[(d - 1) * N + n] - m) / s
				end
			end
			task.synchronize()
			return Y
		end,
	}
	return self
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Losses = require("../nn/Losses")
local Softmax = require("../nn/Softmax")
local Util = require("../Util")
local assert = Util.Assert

return function(cfg)
	local gamma = cfg.gamma or 0.99
	local minBatch = cfg.minBatch or 8

	local policy = assert(cfg.policy, "A2C: cfg.policy (actor network) required")
	local value = assert(cfg.value, "A2C: cfg.value (critic network) required")

	local policyOpt
	local valueOpt

	if cfg.policyOptimizerFactory or cfg.valueOptimizerFactory then
		assert(cfg.policyOptimizerFactory, "A2C: cfg.policyOptimizerFactory required when using separate optimizers")
		assert(cfg.valueOptimizerFactory, "A2C: cfg.valueOptimizerFactory required when using separate optimizers")
		policyOpt = cfg.policyOptimizerFactory(policy:parameters())
		valueOpt = cfg.valueOptimizerFactory(value:parameters())
	else
		assert(
			cfg.optimizerFactory,
			"A2C: provide optimizerFactory or (policyOptimizerFactory + valueOptimizerFactory)"
		)
		policyOpt = cfg.optimizerFactory(policy:parameters())
		valueOpt = cfg.optimizerFactory(value:parameters())
	end

	local S_buf = {}
	local A_buf = {}
	local R_buf = {}
	local D_buf = {}

	local function clearBuf()
		table.clear(S_buf)
		table.clear(A_buf)
		table.clear(R_buf)
		table.clear(D_buf)
	end

	local function packStates()
		local T = #S_buf
		assert(T > 0, "A2C: empty buffer")
		local s0 = S_buf[1]
		local sd = s0._shape[1]
		local B = T
		local S = Tensor.zeros({ sd, B }, s0._dtype, true)

		task.desynchronize()
		for j = 1, B do
			local s = S_buf[j]
			assert(s._shape[1] == sd and s._shape[2] == 1, "A2C: expected state {D,1}")
			for i = 1, sd do
				S._storage[(i - 1) * B + j] = s._storage[i]
			end
		end
		task.synchronize()
		return S
	end

	local function computeReturns()
		local B = #R_buf
		local G = table.create(B, 0.0)
		local ret = 0.0
		task.desynchronize()
		for t = B, 1, -1 do
			if D_buf[t] then
				ret = 0.0
			end
			ret = R_buf[t] + gamma * ret
			G[t] = ret
		end
		task.synchronize()
		return G
	end

	local function trainFromBuffer()
		local T = #S_buf
		if T == 0 then
			return nil
		end

		local S = packStates()
		local B = T

		-- critic update: V(s) -> returns
		local V = value:forward(S) -- {1,B}
		local G = computeReturns()
		local target = Tensor.zeros({ 1, B }, V._dtype, false)

		task.desynchronize()
		for j = 1, B do
			target._storage[j] = G[j]
		end
		task.synchronize()

		valueOpt:zeroGrad()
		local vLoss, gV = Losses.mse_backward(V, target)
		Tape.backwardFrom(V, gV)
		valueOpt:step()

		-- compute advantages A_t = G_t - V(s_t)
		local advantages = table.create(B, 0.0)
		task.desynchronize()
		for j = 1, B do
			local vj = V._storage[j]
			advantages[j] = G[j] - vj
		end
		task.synchronize()

		-- actor update with policy gradient
		policyOpt:zeroGrad()
		local logits = policy:forward(S) -- {A,B}
		local C = logits._shape[1]
		local P = Softmax.forward(logits)
		local grad = Tensor.zeros(logits._shape, logits._dtype, false)

		local lossSum = 0.0
		task.desynchronize()
		for j = 1, B do
			local a = A_buf[j]
			local adv = advantages[j]
			local idx = (a - 1) * B + j
			local prob = P._storage[idx]
			local logp = math.log(math.max(prob, 1e-12))
			lossSum += -adv * logp

			for i = 1, C do
				local pij = P._storage[(i - 1) * B + j]
				local indicator = (i == a) and 1 or 0
				-- dL/dz_i = A * (p_i - 1_{i=a})
				grad._storage[(i - 1) * B + j] = adv * (pij - indicator)
			end
		end
		for k = 1, #grad._storage do
			grad._storage[k] /= B
		end
		task.synchronize()

		local pLoss = lossSum / B

		Tape.backwardFrom(logits, grad)
		policyOpt:step()

		clearBuf()
		return {
			policyLoss = pLoss,
			valueLoss = vLoss,
		}
	end

	local Agent = {}

	--@param state {D,1}
	--@return number
	function Agent:act(state)
		local logits = policy:forward(state)
		local B = logits._shape[2]
		assert(B == 1, "A2C: act() expects state with batch=1")
		local C = logits._shape[1]
		local P = Softmax.forward(logits)

		local r = math.random()
		local cum = 0.0
		-- j = 1 column, so index = i
		for i = 1, C do
			cum += P._storage[i]
			if r <= cum then
				return i
			end
		end
		return C
	end

	--@param t { state = Tensor, action = number, reward = number, done = boolean }
	function Agent:observe(t)
		table.insert(S_buf, t.state)
		table.insert(A_buf, t.action)
		table.insert(R_buf, t.reward)
		table.insert(D_buf, t.done == true)
	end

	--@return {policyLoss:number, valueLoss:number} or nil
	function Agent:trainStep()
		if #S_buf < minBatch then
			return nil
		end
		return trainFromBuffer()
	end

	--@return nn.Module
	function Agent:getPolicy()
		return policy
	end

	--@return nn.Module
	function Agent:getValue()
		return value
	end

	--@return {Tensor}
	function Agent:parameters()
		local params = {}
		local pp = policy:parameters()
		local vp = value:parameters()
		for i = 1, #pp do
			table.insert(params, pp[i])
		end
		for i = 1, #vp do
			table.insert(params, vp[i])
		end
		return params
	end

	function Agent:resetBuffer()
		clearBuf()
	end

	return Agent
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Types = require("../Types")
local Util = require("../Util")
local assert = Util.Assert

local Stabilizers_ok, Stabilizers = pcall(require, "../Utils/Stabilizers")

local function softUpdateMaybe(tp, fp, tau: number?)
	if tau and Stabilizers_ok then
		Stabilizers.softUpdate(tp, fp, tau)
	end
end

local function softmaxAtoms(logits: Types.Tensor, A: number, Z: number)
	local B = logits._shape[2]
	local out = Tensor.zeros({ A * Z, B }, logits._dtype, false)
	task.desynchronize()
	for j = 1, B do
		for a = 1, A do
			local base = ((a - 1) * Z)
			local m = -1 / 0
			for z = 1, Z do
				local v = logits._storage[(base + z - 1) * B + j]
				if v > m then
					m = v
				end
			end
			local es = 0
			for z = 1, Z do
				local e = math.exp(logits._storage[(base + z - 1) * B + j] - m)
				out._storage[(base + z - 1) * B + j] = e
				es += e
			end
			local inv = 1 / es
			for z = 1, Z do
				out._storage[(base + z - 1) * B + j] *= inv
			end
		end
	end
	task.synchronize()
	return out
end

local function qExpect(P: Types.Tensor, A: number, Z: number, zVals: { number })
	local B = P._shape[2]
	local Q = Tensor.zeros({ A, B }, P._dtype, false)
	task.desynchronize()
	for j = 1, B do
		for a = 1, A do
			local base = (a - 1) * Z
			local acc = 0
			for z = 1, Z do
				acc += P._storage[(base + z - 1) * B + j] * zVals[z]
			end
			Q._storage[(a - 1) * B + j] = acc
		end
	end
	task.synchronize()
	return Q
end

local function project(
	zNext: { number },
	pNext: Types.Tensor,
	rewards: { number },
	dones: { number },
	gamma: number,
	A: number,
	Z: number,
	B: number
)
	local vmin, vmax = zNext[1], zNext[#zNext]
	local dz = (vmax - vmin) / (Z - 1)
	local Tproj = Tensor.zeros({ A * Z, B })
	return function(j: number, aNext: number)
		local r, done = rewards[j], dones[j]
		for z = 1, Z do
			local tz = done == 1 and r or (r + gamma * zNext[z])
			local tzClamped = math.max(vmin, math.min(vmax, tz))
			local b = (tzClamped - vmin) / dz
			local l = math.floor(b)
			local u = math.ceil(b)
			for m = l, u do
				if m >= 0 and m <= (Z - 1) then
					local w = (m == u and (b - m)) or (m == l and (1 - (b - m))) or 0
					if w ~= 0 then
						local idx = ((aNext - 1) * Z + (m + 1) - 1) * B + j
						Tproj._storage[idx] += w * pNext._storage[((aNext - 1) * Z + z - 1) * B + j]
					end
				end
			end
		end
	end,
		Tproj
end

return function(cfg)
	local A = assert(cfg.actionDim, "actionDim required")
	local Z = cfg.numAtoms or 51
	local vmin = cfg.vmin or -10
	local vmax = cfg.vmax or 10
	local gamma = cfg.gamma or 0.99
	local dz = (vmax - vmin) / (Z - 1)
	local zVals = table.create(Z)
	for i = 1, Z do
		zVals[i] = vmin + (i - 1) * dz
	end

	local policy = cfg.modelFactory()
	local target = cfg.modelFactory()
	local function copyParams()
		local fp, tp = policy:parameters(), target:parameters()
		task.desynchronize()
		for i = 1, #fp do
			for k = 1, #fp[i]._storage do
				tp[i]._storage[k] = fp[i]._storage[k]
			end
		end
		task.synchronize()
	end
	copyParams()
	local opt = cfg.optimizerFactory(policy:parameters())
	local step = 0

	local function epsAt(n: number)
		local d = cfg.epsilonDecay or 15000
		local e0, e1 = cfg.epsilonStart or 1.0, cfg.epsilonEnd or 0.10
		local k = math.exp(-n / d)
		return e1 + (e0 - e1) * k
	end

	local function pickAction(state)
		-- logits {A*Z,1} -> probs -> expected Q per action -> argmax
		local logits = policy:forward(state)
		local P = softmaxAtoms(logits, A, Z)
		local Q = qExpect(P, A, Z, zVals)
		-- argmax
		local best, ai = -1 / 0, 1
		for a = 1, A do
			local v = Q._storage[(a - 1) * 1 + 1]
			if v > best then
				best = v
				ai = a
			end
		end
		return ai
	end

	return {
		getPolicy = function(_)
			return policy
		end,

		parameters = function(_)
			return policy:parameters()
		end,

		act = function(_, state, n)
			local eps = epsAt(n)
			if math.random() < eps then
				return math.random(1, A)
			end
			return pickAction(state)
		end,

		observe = function(_, t)
			if cfg.replay then
				cfg.replay:push(t)
			end
		end,

		trainStep = function(_)
			if not cfg.replay or not cfg.replay:canSample(cfg.batchSize) then
				return nil
			end
			opt:zeroGrad()
			local batch = cfg.replay:sample(cfg.batchSize)
			local S, Aidx, R, NS, D = batch.S, batch.A, batch.R, batch.NS, batch.D
			local B = cfg.batchSize

			-- logits {A*Z,B}
			local logits = policy:forward(S)
			local P = softmaxAtoms(logits, A, Z)

			local logitsNext_policy = policy:forward(NS)
			local Pnext_policy = softmaxAtoms(logitsNext_policy, A, Z)
			local Qnext_policy = qExpect(Pnext_policy, A, Z, zVals)
			local aNext = table.create(B, 1)
			task.desynchronize()
			for j = 1, B do
				local best, ai = -1 / 0, 1
				for a = 1, A do
					local v = Qnext_policy._storage[(a - 1) * B + j]
					if v > best then
						best = v
						ai = a
					end
				end
				aNext[j] = ai
			end
			task.synchronize()

			local logitsNext_target = target:forward(NS)
			local Pnext_target = softmaxAtoms(logitsNext_target, A, Z)

			local buildOne, Tproj = project(zVals, Pnext_target, R, D, gamma, A, Z, B)
			for j = 1, B do
				buildOne(j, aNext[j])
			end

			local Psel = Tensor.zeros({ Z, B })
			task.desynchronize()
			for j = 1, B do
				local a = Aidx[j]
				for z = 1, Z do
					Psel._storage[(z - 1) * B + j] = P._storage[((a - 1) * Z + z - 1) * B + j]
				end
			end
			task.synchronize()

			local Tsel = Tensor.zeros({ Z, B })
			task.desynchronize()
			for j = 1, B do
				local a = Aidx[j]
				for z = 1, Z do
					Tsel._storage[(z - 1) * B + j] = Tproj._storage[((a - 1) * Z + z - 1) * B + j]
				end
			end
			task.synchronize()

			local grad = Tensor.zeros({ A * Z, B })
			local loss = 0
			task.desynchronize()
			for j = 1, B do
				for z = 1, Z do
					local p = math.max(Psel._storage[(z - 1) * B + j], 1e-12)
					local t = Tsel._storage[(z - 1) * B + j]
					loss += -t * math.log(p)
				end
			end
			for j = 1, B do
				local a = Aidx[j]
				for z = 1, Z do
					local p = Psel._storage[(z - 1) * B + j]
					local t = Tsel._storage[(z - 1) * B + j]
					grad._storage[((a - 1) * Z + z - 1) * B + j] = (p - t) / B
				end
			end
			task.synchronize()

			Tape.backwardFrom(logits, grad)
			if Stabilizers_ok then
				Stabilizers.clipGradGlobalNorm(policy:parameters(), 3.0)
			end
			opt:step()

			step += 1
			if cfg.tau then
				softUpdateMaybe(target:parameters(), policy:parameters(), cfg.tau)
			elseif cfg.targetSyncInterval and (step % cfg.targetSyncInterval == 0) then
				copyParams()
			end
			return loss / B
		end,

		loadParameters = function(_, snapshot, strict)
			local fp = policy:parameters()
			if strict ~= false then
				assert(#fp == #snapshot, "snapshot size mismatch")
				for i = 1, #fp do
					assert(#fp[i]._storage == #snapshot[i], "tensor size mismatch at " .. i)
				end
			end
			local n = math.min(#fp, #snapshot)
			task.desynchronize()
			for i = 1, n do
				local p, s = fp[i], snapshot[i]
				local m = math.min(#p._storage, #s)
				for k = 1, m do
					p._storage[k] = s[k]
				end
			end
			task.synchronize()
			local fp2, tp = policy:parameters(), target:parameters()
			task.desynchronize()
			for i = 1, #fp2 do
				for k = 1, #fp2[i]._storage do
					tp[i]._storage[k] = fp2[i]._storage[k]
				end
			end
			task.synchronize()
		end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Losses = require("../nn/Losses")
local Util = require("../Util")
local assert = Util.Assert

-- argmax over actions per column (batch)
local function maxPerColumn(Q)
	task.desynchronize()
	local A, B = Q._shape[1], Q._shape[2]
	local maxv = table.create(B, -1/0)
	local arg = table.create(B, 1)
	for j = 1, B do
		local best, ai = -1/0, 1
		for i = 1, A do
			local v = Q._storage[(i - 1) * B + j]
			if v > best then best, ai = v, i end
		end
		maxv[j], arg[j] = best, ai
	end
	task.synchronize()
	return maxv, arg
end

-- Pack a list of transitions into a batch {S,A,R,NS,D}
local function packBatch(transitions)
	task.desynchronize()
	local n = #transitions
	local sd = transitions[1].state._shape[1]
	local S = Tensor.zeros({ sd, n })
	local NS = Tensor.zeros({ sd, n })
	local A = table.create(n, 1)
	local R = table.create(n, 0.0)
	local D = table.create(n, false)
	for j = 1, n do
		local t = transitions[j]
		A[j], R[j], D[j] = t.action, t.reward, t.done
		for k = 1, sd do
			S._storage[(k - 1) * n + j] = t.state._storage[k]
			NS._storage[(k - 1) * n + j] = t.nextState._storage[k]
		end
	end
	task.synchronize()
	return S, A, R, NS, D
end

local function asArray(entry)
	if type(entry) == "table" then
		if entry.data then return entry.data end
		if entry[1] ~= nil then return entry end
	end
	error("Bad snapshot entry; expected {data=...} or numeric array")
end

return function(cfg)
	local batchSize = cfg.batchSize or 1
	local gamma = cfg.gamma or 0.99
	local epsStart = cfg.epsilonStart or 1.0
	local epsEnd = cfg.epsilonEnd or 0.05
	local epsDecay = cfg.epsilonDecay or 1000
	local targetSyncEvery = cfg.targetSyncInterval or 200
	local tau = cfg.tau
	local actionDim = assert(cfg.actionDim, "cfg.actionDim required")

	local policy = cfg.modelFactory()
	local target = cfg.modelFactory()

	local function copyParams()
		task.desynchronize()
		local fps = policy:parameters()
		local tps = target:parameters()
		for i = 1, #fps do
			for k = 1, #fps[i]._storage do
				tps[i]._storage[k] = fps[i]._storage[k]
			end
		end
		task.synchronize()
	end
	copyParams()

	local function softUpdate(alpha: number)
		if not alpha or alpha <= 0 then return end
		local fp, tp = policy:parameters(), target:parameters()
		task.desynchronize()
		for i = 1, #fp do
			for k = 1, #fp[i]._storage do
				local tv = tp[i]._storage[k]
				local pv = fp[i]._storage[k]
				tp[i]._storage[k] = tv + alpha * (pv - tv)
			end
		end
		task.synchronize()
	end

	local opt = cfg.optimizerFactory(policy:parameters())
	local step = 0

	local onlineBuf = {}
	local onlineMax = cfg.onlineMaxBatch or math.max(1, batchSize)

	local function epsAt(n)
		local k = math.exp(-n / epsDecay)
		return epsEnd + (epsStart - epsEnd) * k
	end

	local function tdUpdate(S, A, R, NS, D)
		task.desynchronize()
		local B = #A
		local Q = policy:forward(S) -- {A,B}
		local Qn = target:forward(NS) -- {A,B}
		local maxn = (maxPerColumn(Qn))

		local Y = Tensor.zeros({ actionDim, B })
		for i = 1, #Q._storage do
			Y._storage[i] = Q._storage[i]
		end
		for j = 1, B do
			local target_j = R[j] + gamma * (D[j] and 0 or maxn[j])
			Y._storage[(A[j] - 1) * B + j] = target_j
		end

		local loss, grad = Losses.mse_backward(Q, Y)
		Tape.backwardFrom(Q, grad)
		opt:step()

		step += 1
		if tau and tau > 0 then
			softUpdate(tau)
		elseif targetSyncEvery > 0 and step % targetSyncEvery == 0 then
			copyParams()
		end
		task.synchronize()
		return loss
	end

	return {
		getPolicy = function(_)
			return policy
		end,

		parameters = function(_)
			return policy:parameters()
		end,

		act = function(_, state, n)
			local eps = epsAt(n or step)
			if math.random() < eps then
				return math.random(1, actionDim)
			end
			task.desynchronize()
			local q = policy:forward(state)
			local _, arg = maxPerColumn(q)
			task.synchronize()
			return arg[1]
		end,

		observe = function(_, t)
			if cfg.replay then
				cfg.replay:push(t)
				return
			end
			table.insert(onlineBuf, t)
			if #onlineBuf > onlineMax then
				table.remove(onlineBuf, 1)
			end
		end,

		trainStep = function(_)
			if cfg.replay then
				if not cfg.replay:canSample(batchSize) then
					return nil
				end
				local B = cfg.replay:sample(batchSize)
				return tdUpdate(B.S, B.A, B.R, B.NS, B.D)
			end
			local lenBuf = #onlineBuf
			if lenBuf == 0 then
				return nil
			end
			task.desynchronize()
			local n = math.min(batchSize, lenBuf)
			local tmp = table.create(n)
			for i = lenBuf - n + 1, lenBuf do
				table.insert(tmp, onlineBuf[i])
			end
			local S, A, R, NS, D = packBatch(tmp)
			if not cfg.onlineKeep then
				table.clear(onlineBuf)
			end
			task.synchronize()
			return tdUpdate(S, A, R, NS, D)
		end,

		loadParameters = function(_, snapshot, strict: boolean?)
			local fp = policy:parameters()
			local n = #fp
			assert(n > 0, "No parameters in policy")

			local arr = table.create(n)
			for i = 1, n do arr[i] = asArray(snapshot[i]) end

			if strict ~= false then
				assert(#arr == n, "snapshot size mismatch")
				for i = 1, n do
					assert(#fp[i]._storage == #arr[i],
						("tensor size mismatch at %d, got %d:%d")
							:format(i, #fp[i]._storage, #arr[i]))
				end
			end

			task.desynchronize()
			for i = 1, n do
				local p, s = fp[i], arr[i]
				local m = math.min(#p._storage, #s)
				for k = 1, m do p._storage[k] = s[k] end
			end
			task.synchronize()

			-- refresh target (hard sync)
			local tp = target:parameters()
			task.desynchronize()
			for i = 1, n do
				local ps, ts = fp[i]._storage, tp[i]._storage
				for k = 1, #ps do ts[k] = ps[k] end
			end
			task.synchronize()
		end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Losses = require("../nn/Losses")
local Util = require("../Util")
local assert = Util.Assert

local function maxPerColumn(Q)
	local A, B = Q._shape[1], Q._shape[2]
	local maxv = table.create(B, -1 / 0)
	local arg = table.create(B, 1)
	task.desynchronize()
	for j = 1, B do
		local best, ai = -1 / 0, 1
		for i = 1, A do
			local v = Q._storage[(i - 1) * B + j]
			if v > best then
				best, ai = v, i
			end
		end
		maxv[j], arg[j] = best, ai
	end
	task.synchronize()
	return maxv, arg
end

local function packBatch(transitions)
	local n = #transitions
	local sd = transitions[1].state._shape[1]
	local S = Tensor.zeros({ sd, n })
	local NS = Tensor.zeros({ sd, n })
	local A = table.create(n, 1)
	local R = table.create(n, 0.0)
	local D = table.create(n, false)
	task.desynchronize()
	for j = 1, n do
		local t = transitions[j]
		A[j], R[j], D[j] = t.action, t.reward, t.done
		for k = 1, sd do
			S._storage[(k - 1) * n + j] = t.state._storage[k]
			NS._storage[(k - 1) * n + j] = t.nextState._storage[k]
		end
	end
	task.synchronize()
	return S, A, R, NS, D
end

local function asArray(entry)
	if type(entry) == "table" then
		if entry.data then
			return entry.data
		end
		if entry[1] ~= nil then
			return entry
		end
	end
	error("Bad snapshot entry; expected {data=...} or numeric array")
end

return function(cfg)
	local actionDim = assert(cfg.actionDim, "cfg.actionDim required")
	local batchSize = cfg.batchSize or 1
	local gamma = cfg.gamma or 0.99
	local epsStart = cfg.epsilonStart or 1.0
	local epsEnd = cfg.epsilonEnd or 0.05
	local epsDecay = cfg.epsilonDecay or 1000
	local syncEvery = cfg.targetSyncInterval or 200
	local losses_fn = cfg.losses_fn or Losses.mse_backward
	local losses_args = cfg.losses_args or nil
	local tau = cfg.tau

	local policy = cfg.modelFactory()
	local target = cfg.modelFactory()

	local function copyParams()
		local fps = policy:parameters()
		local tps = target:parameters()
		task.desynchronize()
		for i = 1, #fps do
			for k = 1, #fps[i]._storage do
				tps[i]._storage[k] = fps[i]._storage[k]
			end
		end
		task.synchronize()
	end
	copyParams()

	local function softUpdateParams(alpha: number)
		if not alpha or alpha <= 0 then
			return
		end
		local fp, tp = policy:parameters(), target:parameters()
		task.desynchronize()
		for i = 1, #fp do
			local p, t = fp[i], tp[i]
			for k = 1, #p._storage do
				local tk = t._storage[k]
				t._storage[k] = tk + alpha * (p._storage[k] - tk)
			end
		end
		task.synchronize()
	end

	local opt = cfg.optimizerFactory(policy:parameters())
	local step = 0

	local onlineBuf = {}
	local onlineMax = cfg.onlineMaxBatch or math.max(1, batchSize)

	local function epsAt(n)
		local k = math.exp(-(n or step) / epsDecay)
		return epsEnd + (epsStart - epsEnd) * k
	end

	local function tdUpdate(S, A, R, NS, D)
		local B = #A
		local Q = policy:forward(S) -- {A,B}
		local Qn_policy = policy:forward(NS) -- selection net
		local _, argNext = maxPerColumn(Qn_policy) -- a* = argmax_a' Q_online(s', a')
		local Qn_target = target:forward(NS) -- evaluation net

		-- Y := copy(Q); overwrite chosen action with TD target
		local Y = Tensor.zeros({ actionDim, B })
		task.desynchronize()
		for i = 1, #Q._storage do
			Y._storage[i] = Q._storage[i]
		end
		for j = 1, B do
			local an = argNext[j]
			local qnext = Qn_target._storage[(an - 1) * B + j]
			local target_j = R[j] + gamma * (D[j] and 0 or qnext)
			Y._storage[(A[j] - 1) * B + j] = target_j
		end
		task.synchronize()

		local loss, grad = losses_fn(Q, Y, losses_args and table.unpack(losses_args) or nil)
		Tape.backwardFrom(Q, grad)
		opt:step()

		step += 1
		if tau and tau > 0 then
			softUpdateParams(tau)
		elseif syncEvery > 0 and step % syncEvery == 0 then
			copyParams()
		end
		return loss
	end

	return {
		getPolicy = function(_)
			return policy
		end,

		parameters = function(_)
			return policy:parameters()
		end,

		act = function(_, state, n)
			local eps = epsAt(n)
			if math.random() < eps then
				return math.random(1, actionDim)
			end
			local q = policy:forward(state)
			local _, arg = maxPerColumn(q)
			return arg[1]
		end,

		observe = function(_, t)
			if cfg.replay then
				cfg.replay:push(t)
			else
				table.insert(onlineBuf, t)
				if #onlineBuf > onlineMax then
					table.remove(onlineBuf, 1)
				end
			end
		end,

		trainStep = function(_)
			if cfg.replay then
				if not cfg.replay:canSample(batchSize) then
					return nil
				end
				local B = cfg.replay:sample(batchSize)
				return tdUpdate(B.S, B.A, B.R, B.NS, B.D)
			end
			local lenBuf = #onlineBuf
			if lenBuf == 0 then
				return nil
			end
			local n = math.min(batchSize, lenBuf)
			local tmp = table.create(n)
			for i = lenBuf - n + 1, lenBuf do
				table.insert(tmp, onlineBuf[i])
			end
			local S, A, R, NS, D = packBatch(tmp)
			if not cfg.onlineKeep then
				table.clear(onlineBuf)
			end
			return tdUpdate(S, A, R, NS, D)
		end,

		loadParameters = function(_, snapshot, strict: boolean?)
			local fp = policy:parameters()
			local n = #fp
			assert(n > 0, "No parameters in policy")

			local arr = table.create(n)
			for i = 1, n do
				if not strict and not snapshot[i] then
					continue
				end
				arr[i] = asArray(snapshot[i])
			end

			if strict ~= false then
				assert(#arr == n, "snapshot size mismatch")
				for i = 1, n do
					assert(
						#fp[i]._storage == #arr[i],
						("tensor size mismatch at %d, got %d:%d"):format(i, #fp[i]._storage, #arr[i])
					)
				end
			end

			task.desynchronize()
			for i = 1, n do
				local p, s = fp[i], arr[i]
				if not s then
					continue
				end
				local m = math.min(#p._storage, #s)
				for k = 1, m do
					p._storage[k] = s[k]
				end
			end
			task.synchronize()

			-- refresh target (hard sync)
			local tp = target:parameters()
			task.desynchronize()
			for i = 1, n do
				local ps, ts = fp[i]._storage, tp[i]._storage
				for k = 1, #ps do
					ts[k] = ps[k]
				end
			end
			task.synchronize()
		end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")

return function(V, A)
	local ADim, B = A._shape[1], A._shape[2]
	local Q = Tensor.zeros({ ADim, B }, A._dtype, V._requiresGrad or A._requiresGrad)
	task.desynchronize()
	for j = 1, B do
		local meanA = 0
		for i = 1, ADim do
			meanA += A._storage[(i - 1) * B + j]
		end
		meanA /= ADim
		for i = 1, ADim do
			local idx = (i - 1) * B + j
			Q._storage[idx] = V._storage[j] + A._storage[idx] - meanA
		end
	end
	task.synchronize()
	return Tape.makeNode(Q, { V, A }, function(gQ)
		if V._requiresGrad then
			if not V._grad then
				V._grad = Tensor.zeros({ 1, B })
			end
			task.desynchronize()
			for j = 1, B do
				local acc = 0
				for i = 1, ADim do
					acc += gQ._storage[(i - 1) * B + j]
				end
				V._grad._storage[j] += acc
			end
			task.synchronize()
		end
		if A._requiresGrad then
			if not A._grad then
				A._grad = Tensor.zeros({ ADim, B })
			end
			task.desynchronize()
			for j = 1, B do
				local acc = 0
				for i = 1, ADim do
					acc += gQ._storage[(i - 1) * B + j]
				end
				local meanGrad = acc / ADim
				for i = 1, ADim do
					local idx = (i - 1) * B + j
					A._grad._storage[idx] += gQ._storage[idx] - meanGrad
				end
			end
			task.synchronize()
		end
	end)
end


--!native
--!optimize 2

local function push(self, t)
	local out = {}
	local b = self.buf
	table.insert(b, t)
	if t.done then
		while #b > 0 do
			local G, w, k = 0, 1, 0
			task.desynchronize()
			for i = 1, #b do
				k += 1
				local ti = b[i]
				G += w * self.gamma ^ (i - 1) * ti.reward
				w *= self.lambda
				if ti.done then
					break
				end
			end
			task.synchronize()
			table.insert(
				out,
				{ state = b[1].state, action = b[1].action, reward = G, nextState = b[k].nextState, done = true }
			)
			table.remove(b, 1)
		end
	else
		if #b > 20 then
			table.remove(b, 1)
		end
	end
	return out
end

local function new(gamma: number, lambda: number)
	return {
		gamma = gamma,
		lambda = lambda,
		buf = {},
		push = push,
	}
end

return { new = new }


--!native
--!optimize 2
local Types = require("../Types")
local Util = require("../Util")
local assert = Util.Assert

local function getAgent(self, i: number): Types.Module
	assert(i >= 1 and i <= self._numAgents, ("MultiAgent: index %d out of range [1,%d]"):format(i, self._numAgents))
	return self._agents[i]
end

local function size(self): number
	return self._numAgents
end

local function act(self, i: number, state: Types.Tensor, step: number): Types.Tensor
	return (getAgent(self, i) :: any):act(state, step)
end

local function actAll(self, states: { Types.Tensor }, step: number): { Types.Tensor }
	assert(#states == self._numAgents, "MultiAgent: actAll() expects one state per agent")
	local actions = table.create(self._numAgents)
	for idx = 1, self._numAgents do
		actions[idx] = self._agents[idx]:act(states[idx], step)
	end
	return actions
end

local function observe(self, i: number, t: Types.Tensor)
	local agent = getAgent(self, i)
	if agent.observe then
		agent:observe(t)
	end
end

local function observeAll(self, ts: { Types.Tensor })
	for idx = 1, self._numAgents do
		local t = ts[idx]
		if t then
			local agent = self._agents[idx]
			if agent.observe then
				agent:observe(t)
			end
		end
	end
end

local function trainStepAll(self): { Types.Tensor? }
	local out = table.create(self._numAgents)
	for idx = 1, self._numAgents do
		local agent = self._agents[idx]
		if agent.trainStep then
			out[idx] = agent:trainStep()
		else
			out[idx] = nil
		end
	end
	return out
end

local function trainStep(self, i: number?): Types.Tensor?
	if i ~= nil then
		local agent = getAgent(self, i)
		if agent.trainStep then
			return agent:trainStep()
		end
		return nil
	end
	return trainStepAll(self)
end

local function parameters(self): { Types.Tensor }
	local params = {}
	for idx = 1, self._numAgents do
		local agent = self._agents[idx]
		if agent.parameters then
			local ps = agent:parameters()
			if ps then
				for j = 1, #ps do
					params[#params + 1] = ps[j]
				end
			end
		end
	end
	return params
end

local function map(self, fn: (Types.Module, number) -> ())
	for idx = 1, self._numAgents do
		fn(self._agents[idx], idx)
	end
end

return function(cfg: { numAgents: number, makeAgent: (number) -> Types.Module })
	local n = assert(cfg.numAgents, "MultiAgent: cfg.numAgents required")
	assert(n >= 1, "MultiAgent: numAgents must be >= 1")
	local makeAgent = cfg.makeAgent
	assert(makeAgent, "MultiAgent: cfg.makeAgent required")
	local agents = table.create(n)
	for i = 1, n do
		local agent = makeAgent(i)
		assert(agent ~= nil, `MultiAgent: makeAgent({i}) returned nil`)
		agents[i] = agent
	end
	return {
		_agents = agents,
		_numAgents = n,
		getAgent = getAgent,
		size = size,
		act = act,
		actAll = actAll,
		observe = observe,
		observeAll = observeAll,
		trainStep = trainStep,
		trainStepAll = trainStepAll,
		parameters = parameters,
		map = map,
	}
end


--!native
--!optimize 2
local NStep = {}
function NStep.returns(rewards: { number }, dones: { boolean }, gamma: number, n: number): { number }
	task.desynchronize()
	local out = table.create(#rewards, 0)
	for t = 1, #rewards do
		local G, g = 0, 1
		for k = 0, n - 1 do
			local i = t + k
			if i > #rewards or dones[i] then
				G += g * (rewards[i] or 0)
				break
			end
			G += g * rewards[i]
			g *= gamma
		end
		out[t] = G
	end
	task.synchronize()
	return out
end
return NStep


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Losses = require("../nn/Losses")
local Softmax = require("../nn/Softmax")
local Util = require("../Util")
local assert = Util.Assert

return function(cfg)
	assert(cfg.policy, "PPO: cfg.policy (policy network) required")
	assert(cfg.value, "PPO: cfg.value (value network) required")

	local policy = cfg.policy
	local value = cfg.value
	local gamma = cfg.gamma or 0.99
	local lam = cfg.lam or 0.95
	local clipEps = cfg.clip or 0.2
	local minBatch = cfg.minBatch or 32
	local epochs = cfg.epochs or 4
	local maxBuffer = cfg.maxBuffer or (minBatch * 4)

	local policyOpt, valueOpt

	if cfg.policyOptimizerFactory or cfg.valueOptimizerFactory then
		assert(cfg.policyOptimizerFactory, "PPO: cfg.policyOptimizerFactory required when using separate optimizers")
		assert(cfg.valueOptimizerFactory, "PPO: cfg.valueOptimizerFactory required when using separate optimizers")
		policyOpt = cfg.policyOptimizerFactory(policy:parameters())
		valueOpt = cfg.valueOptimizerFactory(value:parameters())
	else
		assert(
			cfg.optimizerFactory,
			"PPO: provide optimizerFactory or (policyOptimizerFactory + valueOptimizerFactory)"
		)
		local p = policy:parameters()
		local v = value:parameters()
		policyOpt = cfg.optimizerFactory(p)
		valueOpt = cfg.optimizerFactory(v)
	end

	local S_buf, NS_buf, A_buf, R_buf, D_buf, logp_buf = {}, {}, {}, {}, {}, {}

	local function clearBuf()
		table.clear(S_buf)
		table.clear(NS_buf)
		table.clear(A_buf)
		table.clear(R_buf)
		table.clear(D_buf)
		table.clear(logp_buf)
	end

	local function maybeTrim()
		local n = #S_buf
		if n <= maxBuffer then
			return
		end
		local drop = n - maxBuffer
		task.desynchronize()
		for i = 1, drop do
			table.remove(S_buf, 1)
			table.remove(NS_buf, 1)
			table.remove(A_buf, 1)
			table.remove(R_buf, 1)
			table.remove(D_buf, 1)
			table.remove(logp_buf, 1)
		end
		task.synchronize()
	end

	local function packStates(states)
		local T = #states
		assert(T > 0, "PPO: empty state buffer")
		local s0 = states[1]
		local sd = s0._shape[1]
		local B = T
		local S = Tensor.zeros({ sd, B }, s0._dtype, true)

		task.desynchronize()
		for j = 1, B do
			local s = states[j]
			assert(s._shape[1] == sd and s._shape[2] == 1, "PPO: expected state {D,1}")
			for i = 1, sd do
				S._storage[(i - 1) * B + j] = s._storage[i]
			end
		end
		task.synchronize()
		return S
	end

	local function computeGAE()
		local T = #R_buf
		assert(T > 0, "PPO: empty reward buffer")

		local S = packStates(S_buf)
		local NS = packStates(NS_buf)
		local B = T

		local V = value:forward(S) -- {1,B}
		local Vn = value:forward(NS) -- {1,B}

		local advantages = table.create(B, 0.0)
		local returns = table.create(B, 0.0)

		task.desynchronize()
		local deltas = table.create(B, 0.0)
		for t = 1, B do
			local r = R_buf[t] or 0.0
			local done = D_buf[t] and 1 or 0
			local vt = V._storage[t]
			local vtn = done == 1 and 0.0 or Vn._storage[t]
			local delta = r + gamma * vtn - vt
			deltas[t] = delta
		end

		local adv = 0.0
		for t = B, 1, -1 do
			local done = D_buf[t] and 1 or 0
			adv = deltas[t] + gamma * lam * (1 - done) * adv
			advantages[t] = adv
			local vt = V._storage[t]
			returns[t] = adv + vt
		end

		local mean, var = 0.0, 0.0
		for t = 1, B do
			mean += advantages[t]
		end
		mean /= B
		for t = 1, B do
			local d = advantages[t] - mean
			var += d * d
		end
		var = var / math.max(B, 1)
		local std = math.sqrt(var) + 1e-8
		for t = 1, B do
			advantages[t] = (advantages[t] - mean) / std
		end
		task.synchronize()

		return {
			S = S,
			V = V,
			advantages = advantages,
			returns = returns,
		}
	end

	local function trainFromBuffer()
		local T = #S_buf
		if T < minBatch then
			return nil
		end

		local info = computeGAE()
		local S = info.S
		local V = info.V
		local advantages = info.advantages
		local returns = info.returns

		local B = #advantages

		local target = Tensor.zeros({ 1, B }, V._dtype, false)
		task.desynchronize()
		for j = 1, B do
			target._storage[j] = returns[j]
		end
		task.synchronize()

		valueOpt:zeroGrad()
		local vLoss, gV = Losses.mse_backward(V, target)
		Tape.backwardFrom(V, gV)
		valueOpt:step()

		local lastPolicyLoss = 0.0

		for _ = 1, epochs do
			policyOpt:zeroGrad()
			local logits = policy:forward(S) -- {A,B}
			local P = Softmax.forward(logits)
			local C = logits._shape[1]

			local grad = Tensor.zeros(logits._shape, logits._dtype, false)
			local lossSum = 0.0

			task.desynchronize()
			for j = 1, B do
				local a = A_buf[j]
				local idx = (a - 1) * B + j
				local prob = P._storage[idx]
				local logp = math.log(math.max(prob, 1e-12))
				local logp_old = logp_buf[j]
				local ratio = math.exp(logp - logp_old)
				local clipped = math.clamp(ratio, 1 - clipEps, 1 + clipEps)

				local adv = advantages[j]
				local w = clipped * adv

				lossSum += -w * logp

				for i = 1, C do
					local pij = P._storage[(i - 1) * B + j]
					local indicator = (i == a) and 1 or 0
					local g = -(w / B) * (indicator - pij)
					grad._storage[(i - 1) * B + j] += g
				end
			end
			task.synchronize()

			lastPolicyLoss = lossSum / B

			Tape.backwardFrom(logits, grad)
			policyOpt:step()
		end

		local avgReturn = 0.0
		for j = 1, B do
			avgReturn += returns[j]
		end
		avgReturn /= B

		clearBuf()
		return {
			policyLoss = lastPolicyLoss,
			valueLoss = vLoss,
			avgReturn = avgReturn,
			batchSize = B,
		}
	end

	local Agent = {}

	function Agent:getPolicy()
		return policy
	end

	function Agent:getValue()
		return value
	end

	function Agent:parameters()
		local params = {}
		local pp = policy:parameters()
		local vp = value:parameters()
		for i = 1, #pp do
			table.insert(params, pp[i])
		end
		for i = 1, #vp do
			table.insert(params, vp[i])
		end
		return params
	end

	function Agent:act(state)
		local logits = policy:forward(state) -- {A,1}
		local P = Softmax.forward(logits)
		local B = logits._shape[2]
		assert(B == 1, "PPO: act() expects state with batch=1")
		local A = logits._shape[1]

		task.desynchronize()
		local r = math.random()
		local cum = 0.0
		local chosen = A
		for i = 1, A do
			cum += P._storage[i]
			if r <= cum then
				chosen = i
				break
			end
		end
		task.synchronize()
		return chosen
	end

	function Agent:observe(t)
		assert(t.state and t.nextState and t.action, "PPO.observe: expected {state, action, reward, nextState, done}")
		table.insert(S_buf, t.state)
		table.insert(NS_buf, t.nextState)
		table.insert(A_buf, t.action)
		table.insert(R_buf, t.reward or 0.0)
		table.insert(D_buf, t.done and true or false)

		local logits = policy:forward(t.state)
		local P = Softmax.forward(logits)
		local B = logits._shape[2]
		assert(B == 1, "PPO.observe: state must have batch=1")
		local a = t.action
		local idx = (a - 1) * B + 1

		task.desynchronize()
		local prob = P._storage[idx]
		local logp = math.log(math.max(prob, 1e-12))
		table.insert(logp_buf, logp)
		task.synchronize()

		maybeTrim()
	end

	function Agent:trainStep()
		return trainFromBuffer()
	end

	function Agent:resetBuffer()
		clearBuf()
	end

	return Agent
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Buffer = require("../Util/Buffer")
local Util = require("../Util")
local assert = Util.Assert

return function(capacity: number, alpha: number, beta0: number)
	alpha = alpha or 0.6
	beta0 = beta0 or 0.4
	local size, pos = 0, 1
	local S, A, R, NS, D = {}, {}, {}, {}, {}
	local N = 1
	while N < capacity do
		N *= 2
	end
	local tree = table.create(2 * N, 0)
	local maxPrio = 1
	local stateDim: number? = nil
	local stateDType: string? = nil

	local function updateLeaf(i, p)
		task.desynchronize()
		local idx = N + i - 1
		tree[idx] = p
		idx //= 2
		while idx >= 1 do
			tree[idx] = tree[idx * 2] + tree[idx * 2 + 1]
			idx //= 2
		end
		task.synchronize()
	end
	local function sampleOne(u)
		task.desynchronize()
		local idx = 1
		while idx < N do
			local left = idx * 2
			if u <= tree[left] then
				idx = left
			else
				u -= tree[left]
				idx = left + 1
			end
		end
		task.synchronize()
		return idx - N + 1
	end
	return {
		push = function(_, t)
			assert(t.state and t.nextState, "prio replay: state and nextState required")
			if not stateDim then
				stateDim = #t.state._storage
				stateDType = t.state._dtype
			end
			local dtype = stateDType
			S[pos] = Buffer.writeDenseArray(t.state._storage, dtype, S[pos])
			NS[pos] = Buffer.writeDenseArray(t.nextState._storage, dtype, NS[pos])
			A[pos], R[pos], D[pos] = t.action, t.reward, t.done
			updateLeaf(pos, maxPrio ^ alpha)
			pos += 1
			if pos > capacity then
				pos = 1
			end
			if size < capacity then
				size += 1
			end
		end,
		canSample = function(_, n)
			return size >= n
		end,
		sample = function(_, n)
			assert(size >= n, "prio replay: not enough samples")
			assert(stateDim, "prio replay: buffer is empty")
			task.desynchronize()
			local sd = stateDim
			local dtype = stateDType
			local SX = Tensor.zeros({ sd, n }, dtype, false)
			local NSX = Tensor.zeros({ sd, n }, dtype, false)
			local AA, RR, DD = table.create(n, 1), table.create(n, 0), table.create(n, false)
			local idxs, w = table.create(n, 1), table.create(n, 0.0)
			local tot = tree[1]
			local segment = tot / n
			for j = 1, n do
				local u = (j - 1) * segment + math.random() * segment
				local i = sampleOne(u)
				if i > size then
					i = size
				end
				idxs[j] = i
				local p = tree[N + i - 1]
				local P = p / tot
				w[j] = ((size * P) ^ -beta0)
				AA[j], RR[j], DD[j] = A[i], R[i], D[i]
				Buffer.copyDenseToColumn(S[i], dtype, SX._storage, j, n, sd)
				Buffer.copyDenseToColumn(NS[i], dtype, NSX._storage, j, n, sd)
			end
			local maxw = 0
			for j = 1, n do
				if w[j] > maxw then
					maxw = w[j]
				end
			end
			for j = 1, n do
				w[j] /= (maxw > 0 and maxw or 1)
			end
			task.synchronize()
			return { indices = idxs, weights = w, S = SX, A = AA, R = RR, NS = NSX, D = DD }
		end,
		updatePriorities = function(_, idxs, prios)
			for k = 1, #idxs do
				local i = idxs[k]
				local p = prios[k]
				if p > maxPrio then
					maxPrio = p
				end
				updateLeaf(i, p ^ alpha)
			end
		end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Buffer = require("../Util/Buffer")
local Util = require("../Util")
local assert = Util.Assert

return function(capacity: number)
	local size, pos = 0, 1
	local S, A, R, NS, D = {}, {}, {}, {}, {}
	local stateDim: number? = nil
	local stateDType: string? = nil

	return {
		push = function(_, t)
			assert(t.state and t.nextState, "replay: state and nextState required")
			local srcState = t.state
			local srcNextState = t.nextState
			assert(srcState._storage and srcNextState._storage, "replay: states must be Tensors")
			if not stateDim then
				stateDim = #srcState._storage
				stateDType = srcState._dtype
			end
			S[pos] = Buffer.writeDenseArray(srcState._storage, stateDType, S[pos])
			NS[pos] = Buffer.writeDenseArray(srcNextState._storage, stateDType, NS[pos])
			A[pos], R[pos], D[pos] = t.action, t.reward, t.done
			pos += 1
			if pos > capacity then
				pos = 1
			end
			if size < capacity then
				size += 1
			end
		end,
		canSample = function(_, n)
			return size >= n
		end,
		sample = function(_, n)
			assert(size >= n, "replay: not enough samples")
			assert(stateDim, "replay: buffer is empty")
			task.desynchronize()
			local sd = stateDim
			local dtype = stateDType
			local SX = Tensor.zeros({ sd, n }, dtype, false)
			local NSX = Tensor.zeros({ sd, n }, dtype, false)
			local AA, RR, DD = table.create(n, 1), table.create(n, 0), table.create(n, false)
			for j = 1, n do
				local i = math.random(1, size)
				AA[j], RR[j], DD[j] = A[i], R[i], D[i]
				Buffer.copyDenseToColumn(S[i], dtype, SX._storage, j, n, sd)
				Buffer.copyDenseToColumn(NS[i], dtype, NSX._storage, j, n, sd)
			end
			task.synchronize()
			return { S = SX, A = AA, R = RR, NS = NSX, D = DD }
		end,
	}
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Buffer = require("../Util/Buffer")
local Util = require("../Util")
local assert = Util.Assert

local function packBatch(self, idxs: { number })
	local B = #idxs
	assert(B > 0, "UniformReplay: empty index list")
	assert(self.stateDim, "UniformReplay: buffer is empty")
	local sd = self.stateDim
	local dtype = self.stateDType

	local S = Tensor.zeros({ sd, B }, dtype, false)
	local NS = Tensor.zeros({ sd, B }, dtype, false)
	local A = table.create(B, 1)
	local R = table.create(B, 0.0)
	local D = table.create(B, false)

	task.desynchronize()
	for j = 1, B do
		local k = idxs[j]
		A[j], R[j], D[j] = self.A[k], self.R[k], self.D[k]
		Buffer.copyDenseToColumn(self.S[k], dtype, S._storage, j, B, sd)
		Buffer.copyDenseToColumn(self.NS[k], dtype, NS._storage, j, B, sd)
	end
	task.synchronize()

	local w = table.create(B, 1.0)
	for j = 1, B do
		w[j] = 1.0
	end

	return { S = S, A = A, R = R, NS = NS, D = D, indices = idxs, weights = w }
end

local function push(self, t)
	assert(t.state and t.nextState, "UniformReplay: state and nextState required")
	local stateTensor = t.state
	local nextStateTensor = t.nextState
	assert(stateTensor._storage and nextStateTensor._storage, "UniformReplay: states must be Tensors")
	if not self.stateDim then
		self.stateDim = #stateTensor._storage
		self.stateDType = stateTensor._dtype
	end
	local i = self.head
	self.S[i] = Buffer.writeDenseArray(stateTensor._storage, self.stateDType, self.S[i])
	self.A[i] = t.action
	self.R[i] = t.reward
	self.NS[i] = Buffer.writeDenseArray(nextStateTensor._storage, self.stateDType, self.NS[i])
	self.D[i] = t.done and true or false

	self.head += 1
	if self.head > self.capacity then
		self.head = 1
	end
	if self.size_ < self.capacity then
		self.size_ += 1
	end
end

local function size(self): number
	return self.size_
end

local function clear(self)
	self.size_, self.head = 0, 1
	self.stateDim, self.stateDType = nil, nil
	table.clear(self.S)
	table.clear(self.A)
	table.clear(self.R)
	table.clear(self.NS)
	table.clear(self.D)
end

local function canSample(self, B: number): boolean
	return self.size_ >= B
end

local function sample(self, B: number)
	assert(self.size_ >= B, "UniformReplay: not enough samples")
	assert(self.stateDim, "UniformReplay: buffer is empty")
	local idxs = table.create(B, 0)
	task.desynchronize()
	for j = 1, B do
		idxs[j] = math.random(1, self.size_)
	end
	task.synchronize()
	return packBatch(self, idxs)
end

--@class UniformReplay
--@field capacity integer
--@field size_ integer
--@field head integer
--@field S {any}
--@field A {number}
--@field R {number}
--@field NS {any}
--@field D {boolean}
--@field stateDim number?
--@field stateDType string?
local function new(capacity: number)
	assert(capacity >= 1, "UniformReplay: capacity must be >= 1")
	return {
		capacity = capacity,
		size_ = 0,
		head = 1,
		S = {},
		A = {},
		R = {},
		NS = {},
		D = {},
		stateDim = nil,
		stateDType = nil,
		push = push,
		size = size,
		clear = clear,
		canSample = canSample,
		sample = sample,
	}
end

return function(capacity: number)
	return new(capacity)
end


--!native
--!optimize 2
local Types = require("./Types")
local NFKC = require("./normalizers/NFKC")
local ByteLevel = require("./pre_tokenizers/ByteLevel")
local ByteLevelProcessor = require("./processors/ByteLevel")
local BPE = require("./models/BPE")
local BPEDecoder = require("./decoders/BPEDecoder")

local Tokenizer = {}
Tokenizer.__index = Tokenizer

function Tokenizer.new(model: Types.Model)
	local self = setmetatable({}, Tokenizer)
	self.model = model
	self.normalizer = nil
	self.pre_tokenizer = nil
	self.decoder = nil
	self.post_processor = nil
	return self
end

function Tokenizer:encode(text: string, pair: string?): Types.Encoding
	if self.normalizer then
		text = self.normalizer:normalize(text)
		if pair then
			pair = self.normalizer:normalize(pair)
		end
	end
	local words = { text }
	local pair_words = if pair then { pair } else nil
	if self.pre_tokenizer then
		words, _ = self.pre_tokenizer:pre_tokenize(text)
		if pair and pair_words then
			pair_words, _ = self.pre_tokenizer:pre_tokenize(pair)
		end
	end
	local function run_model(seq_words)
		local tokens = self.model:tokenize(seq_words)
		local ids = {}
		for _, t in ipairs(tokens) do
			local id = self.model:token_to_id(t)
			table.insert(ids, id or 0)
		end
		return {
			ids = ids,
			tokens = tokens,
			offsets = {},
			attention_mask = table.create(#ids, 1),
			special_tokens_mask = table.create(#ids, 0),
			type_ids = table.create(#ids, 0),
		}
	end
	local encoding = run_model(words)
	local pair_encoding = if pair_words then run_model(pair_words) else nil
	if self.post_processor then
		encoding = self.post_processor:process(encoding, pair_encoding)
	elseif pair_encoding then
		for i = 1, #pair_encoding.ids do
			table.insert(encoding.ids, pair_encoding.ids[i])
			table.insert(encoding.tokens, pair_encoding.tokens[i])
			table.insert(encoding.attention_mask, 1)
			table.insert(encoding.special_tokens_mask, 0)
			table.insert(encoding.type_ids, 1)
		end
	end
	return encoding
end

function Tokenizer:decode(ids: { number }): string
	local tokens = {}
	for _, id in ipairs(ids) do
		local t = self.model:id_to_token(id)
		if t then
			table.insert(tokens, t)
		end
	end
	if self.decoder then
		return self.decoder:decode(tokens)
	end
	return table.concat(tokens, " ")
end

function Tokenizer:set_pre_tokenizer(pt: Types.PreTokenizer)
	self.pre_tokenizer = pt
end

function Tokenizer:set_post_processor(pp: Types.PostProcessor)
	self.post_processor = pp
end

function Tokenizer:set_decoder(d: Types.Decoder)
	self.decoder = d
end

function Tokenizer:set_normalizer(n: Types.Normalizer)
	self.normalizer = n
end

function Tokenizer:train_from_iterator(iterator: () -> string?, trainer: Types.Trainer)
	local incoming_iterator = iterator
	if self.normalizer then
		local previous_iterator = iterator
		local function wrapped_iterator()
			local text = previous_iterator()
			if text then
				return self.normalizer:normalize(text)
			end
			return nil
		end
		incoming_iterator = wrapped_iterator
	end
	local old_pre_tok = nil
	if (trainer :: any).pre_tokenizer ~= nil then
		old_pre_tok = (trainer :: any).pre_tokenizer
	end
	if self.pre_tokenizer then
		(trainer :: any).pre_tokenizer = self.pre_tokenizer
	end
	local model = trainer:train(incoming_iterator)
	if self.pre_tokenizer then
		(trainer :: any).pre_tokenizer = old_pre_tok
	end
	self.model = model
end

function Tokenizer:dump()
	local dump = {
		model = self.model:dump(),
		normalizer = nil,
		pre_tokenizer = nil,
		post_processor = nil,
		decoder = nil,
	}
	if self.normalizer and self.normalizer.dump then
		dump.normalizer = self.normalizer:dump()
	elseif self.normalizer then
		if getmetatable(self.normalizer) == NFKC then
			dump.normalizer = { type = "NFKC" }
		end
	end
	if self.pre_tokenizer and self.pre_tokenizer.dump then
		dump.pre_tokenizer = self.pre_tokenizer:dump()
	elseif self.pre_tokenizer then
		if getmetatable(self.pre_tokenizer) == ByteLevel then
			dump.pre_tokenizer = {
				type = "ByteLevel",
				add_prefix_space = self.pre_tokenizer.add_prefix_space,
				use_regex = self.pre_tokenizer.use_regex,
			}
		end
	end
	if self.post_processor and self.post_processor.dump then
		dump.post_processor = self.post_processor:dump()
	elseif self.post_processor then
		if getmetatable(self.post_processor) == ByteLevelProcessor then
			dump.post_processor = {
				type = "ByteLevel",
				trim_offsets = self.post_processor.trim_offsets,
			}
		end
	end
	if self.decoder and self.decoder.dump then
		dump.decoder = self.decoder:dump()
	elseif self.decoder then
		if getmetatable(self.decoder) == BPEDecoder then
			dump.decoder = {
				type = "BPEDecoder",
				suffix = self.decoder.suffix,
				suffix_replacement = self.decoder.suffix_replacement,
			}
		end
	end
	return dump
end

function Tokenizer.from_dump(data)
	local model
	if data.model.type == "BPE" then
		model = BPE.from_dump(data.model)
	else
		error("Unknown model type: " .. tostring(data.model.type))
	end
	local tok = Tokenizer.new(model)
	if data.normalizer then
		if data.normalizer.type == "NFKC" then
			tok:set_normalizer(NFKC.new())
		end
	end
	if data.pre_tokenizer then
		if data.pre_tokenizer.type == "ByteLevel" then
			tok:set_pre_tokenizer(ByteLevel.new(data.pre_tokenizer.add_prefix_space, data.pre_tokenizer.use_regex))
		end
	end
	if data.post_processor then
		if data.post_processor.type == "ByteLevel" then
			tok:set_post_processor(ByteLevelProcessor.new(data.post_processor.trim_offsets))
		end
	end
	if data.decoder then
		if data.decoder.type == "BPEDecoder" then
			tok:set_decoder(BPEDecoder.new(data.decoder.suffix, data.decoder.suffix_replacement))
		end
	end
	return tok
end

return Tokenizer


--!strict
export type Encoding = {
	ids: { number },
	tokens: { string },
	offsets: { { number } },
	attention_mask: { number },
	special_tokens_mask: { number },
	type_ids: { number },
}

export type Model = {
	tokenize: (self: Model, sequences: { string }) -> { any },
	token_to_id: (self: Model, token: string) -> number?,
	id_to_token: (self: Model, id: number) -> string?,
	get_vocab: (self: Model) -> { [string]: number },
	get_vocab_size: (self: Model) -> number,
	save: (self: Model, folder: string, prefix: string?) -> { string },
}

export type PreTokenizer = {
	pre_tokenize: (self: PreTokenizer, text: string) -> ({ string }, { { number } }),
}

export type Decoder = {
	decode: (self: Decoder, tokens: { string }) -> string,
}

export type PostProcessor = {
	process: (self: PostProcessor, encoding: Encoding, pair_encoding: Encoding?) -> Encoding,
}

export type Normalizer = {
	normalize: (self: Normalizer, text: string) -> string,
}

export type Trainer = {
	train: (self: Trainer, iterator: () -> string?) -> (Model, { string }),
}

return nil


--!optimize 2
local BPEDecoder = {}
BPEDecoder.__index = BPEDecoder

local magic = {
	["^"] = "%^",
	["$"] = "%$",
	["("] = "%(",
	[")"] = "%)",
	["%"] = "%%",
	["."] = "%.",
	["["] = "%[",
	["]"] = "%]",
	["*"] = "%*",
	["+"] = "%+",
	["-"] = "%-",
	["?"] = "%?",
}

function BPEDecoder.new(suffix: string?, suffix_replacement: string?)
	local self = setmetatable({}, BPEDecoder)
	self.suffix = suffix or "</w>"
	self.suffix_replacement = suffix_replacement or " "
	return self
end

function BPEDecoder:decode(tokens: { string })
	local text = table.concat(tokens, "")
	local pattern = string.gsub(self.suffix, ".", magic)
	text = string.gsub(text, pattern, self.suffix_replacement)
	return text:match("^%s*(.-)%s*$") or ""
end

return BPEDecoder


local BPEDecoder = require("@self/BPEDecoder")

return {
	BPEDecoder = BPEDecoder,
}