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

local Workspace = game:GetService("Workspace")
local RunService = game:GetService("RunService")
local HttpService = game:GetService("HttpService")

local Types = require("../Types")

type Tensor = Types.Tensor
type Module = Types.Module

export type RenderOptions = {
	layerSizeHints: { number }?,
	maxNeuronsPerLayer: number?,
	maxBeamsPerEdge: number?,
	layerSpacing: number?,
	boxDepth: number?,
	unitHeight: number?,
	updateInterval: number?,
	mode: string?,
	palette: { Color3 }?,
	name: string?,
	getActivations: ((model: any) -> { { number } })?,
	showNeuronValues: boolean?,
	valueEvery: number?,
	valueDecimals: number?,
	valueColor: Color3?,
	valueSize: Vector2?,
	sparsity: number?,
}

export type RenderHandle = {
	folder: Folder,
	setMode: (self: RenderHandle, mode: string) -> (),
	update: (self: RenderHandle) -> (),
	destroy: (self: RenderHandle) -> (),
	setAutoUpdate: (self: RenderHandle, on: boolean) -> (),
	setSparsity: (self: RenderHandle, th: number) -> (),
	setValuesVisible: (self: RenderHandle, on: boolean) -> (),
}

type BeamObj = {
	beam: Beam,
	iOut: number,
	jIn: number,
	_cWidth: number,
	_cColor: Color3,
	_cEnabled: boolean,
}

type NodeObj = {
	part: BasePart,
	label: TextLabel?,
	idx: number,
	_cColor: Color3,
	_cTrans: number,
	_cSize: Vector3,
	_cText: string?,
}

local function fmtVal(a: number, digits: number): string
	local aa = math.abs(a)
	if aa < 1e-3 and aa ~= 0 then
		return `{string.format("%.2e", a)}`
	end
	return string.format("%." .. digits .. "f", a)
end

local function newPart(parent: Instance, cf: CFrame, size: Vector3, color: Color3, name: string): BasePart
	local p = Instance.new("Part")
	p.Name = name
	p.Anchored = true
	p.CanCollide = false
	p.CastShadow = false
	p.Material = Enum.Material.SmoothPlastic
	p.Color = color
	p.CFrame = cf
	p.Size = size
	p.Parent = parent
	return p
end

local function label(part: BasePart, text: string)
	local bb = Instance.new("BillboardGui")
	bb.Name = "Label"
	bb.AlwaysOnTop = true
	bb.Size = UDim2.fromScale(6, 2)
	bb.StudsOffset = Vector3.new(0, part.Size.Y * 0.6 + 0.75, 0)
	local tl = Instance.new("TextLabel")
	tl.Size = UDim2.fromScale(1, 1)
	tl.BackgroundTransparency = 1
	tl.Text = text
	tl.TextScaled = true
	tl.Font = Enum.Font.GothamMedium
	tl.TextColor3 = Color3.new(1, 1, 1)
	tl.TextStrokeTransparency = 0.6
	tl.Parent = bb
	bb.Parent = part
end

local function labelVal(node: BasePart, valueDigits: number, valueColor: Color3, valueSize: Vector2)
	local bb = Instance.new("BillboardGui")
	bb.Name = "Val"
	bb.AlwaysOnTop = true
	bb.LightInfluence = 0
	bb.Size = UDim2.fromScale(valueSize.X, valueSize.Y)
	bb.StudsOffset = Vector3.new(0, 0, 0.28)
	local tl = Instance.new("TextLabel")
	tl.BackgroundTransparency = 1
	tl.Size = UDim2.fromScale(1, 1)
	tl.TextScaled = true
	tl.Font = Enum.Font.GothamMedium
	tl.TextColor3 = valueColor
	tl.TextStrokeTransparency = 0.35
	tl.Text = fmtVal(0, valueDigits)
	tl.Parent = bb
	bb.Parent = node
	return tl
end

local function ema(prev: number, x: number, alpha: number): number
	return alpha * x + (1 - alpha) * prev
end

local function pickEven(N: number, k: number): { number }
	k = math.min(k, N)
	local out = table.create(k)
	for i = 1, k do
		local idx = math.floor((i - 0.5) * (N / k)) + 1
		if idx > N then
			idx = N
		end
		out[i] = idx
	end
	return out
end

local function collectWeightMats(params: { Tensor }): { Tensor }
	local Ws = {}
	for _, t in ipairs(params) do
		if t._shape and #t._shape == 2 then
			table.insert(Ws, t)
		end
	end
	return Ws
end

local function inferSizes(params: { Tensor }): { number }
	local Ws = collectWeightMats(params)
	local L = #Ws
	if L == 0 then
		return {}
	end
	local sizes = table.create(L + 1)
	sizes[1] = Ws[1]._shape[2]
	for i = 1, L do
		table.insert(sizes, Ws[i]._shape[1])
	end
	return sizes
end

local function sizesFromActs(getActsFn, model)
	local acts = getActsFn(model)
	local L = type(acts) == "table" and #acts or 0
	if L > 1 then
		local s = table.create(L)
		for i = 1, L do
			s[i] = #acts[i]
		end
		return s
	end
	return {}
end

function Visual3D.render(model: any, origin: CFrame, opts: RenderOptions?): RenderHandle
	local opt = opts or {}
	local maxNeurons = opt.maxNeuronsPerLayer or 48
	local maxBeamsPerEdge = opt.maxBeamsPerEdge or 256
	local spacing = opt.layerSpacing or 10
	local boxDepth = opt.boxDepth or 2
	local unitH = opt.unitHeight or 0.12
	local updateEvery = opt.updateInterval or 0.8
	local sparsity = opt.sparsity or 0
	local mode = (opt.mode == "grads") and "grads" or "weights"
	local palette = opt.palette
		or { Color3.fromRGB(80, 160, 255), Color3.fromRGB(120, 120, 120), Color3.fromRGB(255, 160, 80) }
	local name = opt.name or `NNViz_{HttpService:GenerateGUID(false)}`
	local getActs = opt.getActivations
	local showVals = (typeof(getActs) == "function" and opt.showNeuronValues == true)
	local valueEvery = opt.valueEvery or 1
	local valueDigits = opt.valueDecimals or 2
	local valueColor = opt.valueColor or Color3.fromRGB(235, 235, 245)
	local valueSize = opt.valueSize or Vector2.new(0.5, 0.15)

	local folder = Instance.new("Folder")
	folder.Name = name
	folder.Parent = Workspace

	local params = model:parameters() :: { Tensor }
	local sizes = opt.layerSizeHints or (getActs and sizesFromActs(getActs, model)) or inferSizes(params)

	local L = #sizes
	if L < 2 then
		local p = newPart(folder, origin, Vector3.new(5, 1, 2), Color3.fromRGB(255, 90, 90), "Error")
		label(p, "Visual3D: < 2 layers detected")
		return {
			folder = folder,
			setMode = function() end,
			update = function() end,
			destroy = function(self)
				self.folder:Destroy()
			end,
			setAutoUpdate = function() end,
			setSparsity = function() end,
			setValuesVisible = function() end,
		}
	end
	local layerParts: { BasePart } = table.create(L)
	local nodeCache: { { NodeObj } } = table.create(L)
	local neuronPicks: { { number } } = table.create(L)
	local nodePinsByLayer: { [number]: { [number]: Attachment } } = {}
	for li = 1, L do
		local n = sizes[li]
		local height = math.max(1, n * unitH)
		local color = (li == 1) and palette[1] or ((li == L) and palette[3] or palette[2])
		local p = newPart(
			folder,
			origin * CFrame.new((li - 1) * spacing, height / 2, 0),
			Vector3.new(3, height, boxDepth),
			color,
			`Layer_{li}`
		)
		local tag = (li == 1 and `Input {n}`) or (li == L and `Output {n}`) or `Hidden {n}`
		label(p, tag)
		layerParts[li] = p

		local nf = Instance.new("Folder")
		nf.Name = `Nodes_{li}`
		nf.Parent = folder
		local picks = pickEven(n, math.min(maxNeurons, n))
		neuronPicks[li] = picks
		local layerNodes = table.create(#picks)
		nodeCache[li] = layerNodes
		local pins = {}
		nodePinsByLayer[li] = pins

		for k = 1, #picks do
			local j = picks[k]
			local y = -p.Size.Y / 2 + (j - 0.5) * (p.Size.Y / n)
			local node = Instance.new("Part")
			node.Name = `n_{j}`
			node.Shape = Enum.PartType.Ball
			node.Anchored = true
			node.CanCollide = false
			node.CastShadow = false
			node.Material = Enum.Material.SmoothPlastic
			node.Color = Color3.fromRGB(200, 200, 200)
			node.Size = Vector3.new(0.35, 0.35, 0.35)
			node.CFrame = p.CFrame * CFrame.new(0, y, 0.6 + 0.5)
			node.Parent = nf
			local lbl: TextLabel? = nil
			if showVals and ((k - 1) % valueEvery == 0) then
				lbl = labelVal(node, valueDigits, valueColor, valueSize)
			end
			local pin = Instance.new("Attachment")
			pin.Name = `pin_{j}`
			pin.Parent = node
			pins[j] = pin
			layerNodes[k] = {
				part = node,
				label = lbl,
				idx = j,
				_cColor = Color3.new(1, 1, 1),
				_cTrans = 0,
				_cSize = Vector3.new(1, 1, 1),
				_cText = nil,
			}
		end
	end
	local Ws = collectWeightMats(params)
	local beamSpecs = {} :: { { W: Tensor, objects: { BeamObj }, emaMax: number } }
	local linkCount = math.min(#Ws, L - 1)
	for li = 1, linkCount do
		local W = Ws[li]
		local fullInPick = neuronPicks[li] or {}
		local fullOutPick = neuronPicks[li + 1] or {}
		local targetSide = math.floor(math.sqrt(maxBeamsPerEdge))
		local strideIn = math.max(1, math.floor(#fullInPick / math.max(1, targetSide)))
		local strideOut = math.max(1, math.floor(#fullOutPick / math.max(1, targetSide)))
		local bf = Instance.new("Folder")
		bf.Name = `Beams_{li}`
		bf.Parent = folder
		local objects: { BeamObj } = {}
		local total = 0
		for oi = 1, #fullOutPick, strideOut do
			if total >= maxBeamsPerEdge then
				break
			end
			for ij = 1, #fullInPick, strideIn do
				if total >= maxBeamsPerEdge then
					break
				end
				local iOut = fullOutPick[oi]
				local jIn = fullInPick[ij]
				local pinIn = nodePinsByLayer[li] and nodePinsByLayer[li][jIn]
				local pinOut = nodePinsByLayer[li + 1] and nodePinsByLayer[li + 1][iOut]
				if pinIn and pinOut then
					local b = Instance.new("Beam")
					b.FaceCamera = true
					b.Transparency = NumberSequence.new(0.15)
					b.Attachment0 = pinIn
					b.Attachment1 = pinOut
					b.Width0 = 0.01
					b.Width1 = 0.01
					b.Color = ColorSequence.new(Color3.fromRGB(180, 180, 180))
					b.Parent = bf
					table.insert(objects, {
						beam = b,
						iOut = iOut,
						jIn = jIn,
						_cWidth = 0.01,
						_cColor = Color3.new(1, 1, 1),
						_cEnabled = true,
					})
					total += 1
				end
			end
		end
		table.insert(beamSpecs, { W = W, objects = objects, emaMax = 1.0 })
	end
	local function getValueField(t: Tensor, idx: number, useGrad: boolean): number
		if useGrad and t._grad and t._grad._storage then
			return math.abs(t._grad._storage[idx] or 0)
		end
		return t._storage[idx] or 0
	end
	local autoConn: RBXScriptConnection? = nil
	local currentMode = mode
	local handle: RenderHandle
	handle = {
		folder = folder,
		setMode = function(self, m: string)
			if m == "grads" or m == "weights" then
				currentMode = m
			end
		end,
		setSparsity = function(_, th: number)
			sparsity = math.clamp(th, 0, 0.5)
		end,
		setValuesVisible = function(_, on: boolean)
			if not getActs then
				return
			end
			showVals = on
			for _, layerList in nodeCache do
				for _, obj in layerList do
					if obj.label and obj.label.Parent and obj.label.Parent:IsA("BillboardGui") then
						obj.label.Parent.Enabled = on
					end
				end
			end
		end,
		update = function(self)
			task.desynchronize()
			local useGrad = (currentMode == "grads")
			for _, spec in ipairs(beamSpecs) do
				local W = spec.W
				local inDim = W._shape[2]
				local curMax = 1e-12
				for _, obj in ipairs(spec.objects) do
					local idx = (obj.iOut - 1) * inDim + obj.jIn
					local v = getValueField(W, idx, useGrad)
					local a = math.abs(v)
					if a > curMax then
						curMax = a
					end
				end
				spec.emaMax = ema(spec.emaMax, curMax, 0.2)
				local denom = spec.emaMax + 1e-12
				for _, obj in ipairs(spec.objects) do
					local idx = (obj.iOut - 1) * inDim + obj.jIn
					local val = getValueField(W, idx, useGrad)
					local mag = math.clamp(math.abs(val) / denom, 0, 1)
					local keep = (mag >= sparsity)
					obj._cEnabled = keep
					if keep then
						local w = 0.01 + 0.015 * mag
						obj._cWidth = w
						if useGrad then
							obj._cColor = Color3.fromRGB(255, 220, 90)
						else
							if val >= 0 then
								obj._cColor = Color3.fromRGB(90, 180, 255)
							else
								obj._cColor = Color3.fromRGB(255, 120, 120)
							end
						end
					end
				end
			end
			local actResults = nil
			if getActs then
				local acts = getActs(model)
				if type(acts) == "table" then
					actResults = acts
					for li = 1, #nodeCache do
						local layerActs = acts[li]
						local layerNodes = nodeCache[li]
						if layerActs and layerNodes then
							for _, nodeObj in ipairs(layerNodes) do
								local j = nodeObj.idx
								local a = layerActs[j] or 0
								local s = math.clamp(math.abs(a), 0, 1)
								if a >= 0 then
									nodeObj._cColor = Color3.fromRGB(150, 220, 255)
								else
									nodeObj._cColor = Color3.fromRGB(255, 160, 160)
								end
								nodeObj._cTrans = 1 - (0.15 + 0.75 * s)
								local sz = 0.30 + 0.25 * s
								nodeObj._cSize = Vector3.new(sz, sz, sz)
								if showVals and nodeObj.label then
									nodeObj._cText = fmtVal(a, valueDigits)
								end
							end
						end
					end
				end
			end
			task.synchronize()
			for _, spec in ipairs(beamSpecs) do
				for _, obj in ipairs(spec.objects) do
					obj.beam.Enabled = obj._cEnabled
					if obj._cEnabled then
						local w = obj._cWidth
						obj.beam.Width0 = w
						obj.beam.Width1 = w
						obj.beam.Color = ColorSequence.new(obj._cColor)
					end
				end
			end
			if getActs and actResults then
				for li = 1, #nodeCache do
					local layerNodes = nodeCache[li]
					for _, nodeObj in ipairs(layerNodes) do
						nodeObj.part.Color = nodeObj._cColor
						nodeObj.part.Transparency = nodeObj._cTrans
						nodeObj.part.Size = nodeObj._cSize
						if nodeObj._cText and nodeObj.label then
							nodeObj.label.Text = nodeObj._cText
						end
					end
				end
			end
		end,
		setAutoUpdate = function(self, on: boolean)
			if autoConn then
				autoConn:Disconnect()
				autoConn = nil
			end
			if on then
				local t = 0.0
				autoConn = RunService.PreSimulation:Connect(function(dt: number)
					t += dt
					if t >= updateEvery then
						t = 0.0
						self:update()
					end
				end)
			end
		end,
		destroy = function(self)
			if autoConn then
				autoConn:Disconnect()
				autoConn = nil
			end
			self.folder:Destroy()
		end,
	}
	handle:update()
	task.synchronize()
	handle:setAutoUpdate(true)
	return handle
end

return Visual3D


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

local Types = require("./Types")

local function isModule(m: any): boolean
	return type(m) == "table" and type(m.forward) == "function" and type(m.parameters) == "function"
end
local function applyRecursive(m: any, fn: (any) -> ())
	if isModule(m) then
		fn(m)
	end
	if type(m) == "table" then
		if type(m._layers) == "table" then
			for _, sub in m._layers do
				applyRecursive(sub, fn)
			end
		end
		for _, v in m do
			if v ~= m and isModule(v) then
				applyRecursive(v, fn)
			end
		end
	end
end

--@param shape {number}
--@return number
function Util.sizeFromShape(shape: { number }): number
	task.desynchronize()
	local n = 1
	for _, d in shape do
		n *= d
	end
	task.synchronize()
	return n
end

function Util.train(m: Types.Module)
	applyRecursive(m, function(mm)
		(mm :: any)._train = true
	end)
end

function Util.eval(m: Types.Module)
	applyRecursive(m, function(mm)
		(mm :: any)._train = false
	end)
end

function Util.apply(m: Types.Module, fn: (Types.Module) -> ())
	applyRecursive(m, function(mm)
		fn(mm :: Types.Module)
	end)
end

function Util.to(m: Types.Module, dtype: Types.DType)
	applyRecursive(m, function(mm)
		for _, p in mm:parameters() do
			p._dtype = dtype
		end
	end)
end

function Util.shape(x: Types.Tensor): (number, number?)
    return x._shape[1], x._shape[2]
end

Util.ModelStats = require("@self/ModelStats")
Util.Visual2D = require("@self/Visual2D")
Util.Visual3D = require("@self/Visual3D")
Util.RunningNorm = require("@self/RunningNorm")
Util.Anomaly = require("@self/Anomaly")
Util.RNG = require("@self/RNG")
Util.Hooks = require("@self/Hooks")
Util.Parallel = require("@self/Parallel")
Util.GradStats = require("@self/GradStats")
Util.Buffer = require("@self/Buffer")
Util.Profiler = require("@self/Profiler")
Util.Assert = require("@self/Assert")

return Util


--!native
--!optimize 2
local Tensor = require("../Tensor")
local BLAS = require("../ops/BLAS")
local Types = require("../Types")
local Tape = {}
local gradEnabled = true

local function withTape(t: Types.Tensor, parents: { Types.Tensor }, back: (grad: Types.Tensor) -> ()): Types.Tensor
	if not gradEnabled then
		t._parents = nil
		t._backprop = nil
		t._requiresGrad = false
		return t
	end
	t._parents = parents
	t._backprop = back
	t._requiresGrad = false
	for _, p in parents do
		if p._requiresGrad then
			t._requiresGrad = true
			break
		end
	end
	return t
end

--@param out any @Tensor
--@param parents {any}
--@param back fun(grad: any)
function Tape.makeNode(out: Types.Tensor, parents: { Types.Tensor }, back: (grad: Types.Tensor) -> ()): Types.Tensor
	return withTape(out, parents, back)
end

-- y = ReLU(x)
--@param x @Tensor
function Tape.relu(x: Types.Tensor): Types.Tensor
	task.desynchronize()
	local y = Tensor.zeros(x._shape, x._dtype, x._requiresGrad)
	for i = 1, #x._storage do
		local v = x._storage[i]
		y._storage[i] = (v > 0) and v or 0
	end
	task.synchronize()
	return withTape(y, { x }, function(gy)
		if x._requiresGrad then
			task.desynchronize()
			if not x._grad then
				x._grad = Tensor.zeros(x._shape, x._dtype)
			end
			if x._grad then
				for i = 1, #x._grad._storage do
					x._grad._storage[i] += (x._storage[i] > 0 and 1 or 0) * gy._storage[i]
				end
			end
			task.synchronize()
		end
	end)
end

-- y = A @ B
--@param x @Tensor
function Tape.matmul(A: Types.Tensor, B: Types.Tensor): Types.Tensor
	local Y = BLAS.matmul(A, B)
	return withTape(Y, { A, B }, function(gY: Types.Tensor)
		-- dA = gY @ B^T ; dB = A^T @ gY
		task.desynchronize()
		local M, K = A._shape[1], A._shape[2]
		local _, N = B._shape[1], B._shape[2]
		if A._requiresGrad then
			if not A._grad then
				A._grad = Tensor.zeros({ M, K }, A._dtype)
			end
			local BT
			Tape.noGrad(function()
				BT = Tensor.transpose(B)
			end)
			local dA = BLAS.matmul(gY, BT)
			if A._grad then
				local dst = A._grad._storage
				local src = dA._storage
				for idx = 1, #dst do
					dst[idx] += src[idx]
				end
			end
		end
		if B._requiresGrad then
			if not B._grad then
				B._grad = Tensor.zeros({ K, N }, B._dtype)
			end
			local AT
			Tape.noGrad(function()
				AT = Tensor.transpose(A)
			end)
			local dB = BLAS.matmul(AT, gY)
			if B._grad then
				local dst = B._grad._storage
				local src = dB._storage
				for idx = 1, #dst do
					dst[idx] += src[idx]
				end
			end
		end
		task.synchronize()
	end)
end

function Tape.backwardFrom(y: Types.Tensor, grad: Types.Tensor)
	task.desynchronize()
	local pending: { [Types.Tensor]: number } = {}
	local visitStack = { y }
	local seen = {}
	while #visitStack > 0 do
		local node = table.remove(visitStack)
		if not seen[node] then
			seen[node] = true
			if node._parents then
				for _, parent in node._parents do
					if parent._requiresGrad then
						pending[parent] = (pending[parent] or 0) + 1
						table.insert(visitStack, parent)
					end
				end
			end
		end
	end
	y._grad = grad
	local ready = { y }
	while #ready > 0 do
		local node = table.remove(ready)
		if node._requiresGrad then
			local g = node._grad
			if not g then
				g = Tensor.zeros(node._shape, node._dtype)
				node._grad = g
			end
			if node._backprop then
				node._backprop(g)
			end
			if node._parents then
				for _, parent in node._parents do
					if parent._requiresGrad then
						local remaining = pending[parent]
						if remaining then
							remaining -= 1
							if remaining == 0 then
								pending[parent] = nil
								if parent._backprop or parent._parents then
									table.insert(ready, parent)
								end
							else
								pending[parent] = remaining
							end
						else
							if parent._backprop or parent._parents then
								table.insert(ready, parent)
							end
						end
					end
				end
			end
		end
	end

	task.synchronize()
end

function Tape.noGrad(fn: () -> ()): ()
	gradEnabled = false
	fn()
	gradEnabled = true
end

function Tape.grad(f: (...any) -> Types.Tensor, inputs: { Types.Tensor }): { Types.Tensor? }
	local oldReq: { boolean } = {}
	for i, x in inputs do
		oldReq[i] = x._requiresGrad
		x._requiresGrad = true
		x._grad = nil
	end

	local y = f(table.unpack(inputs))
	if y == nil or y._shape == nil then
		error("Autograd.grad: f must return a Tensor")
	end

	Tape.backwardFrom(y, Tensor.ones(y._shape))

	local out: { Types.Tensor? } = {}
	for i, x in inputs do
		out[i] = x._grad
		x._requiresGrad = oldReq[i]
	end
	return out
end

return Tape


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

return function(C: number)
	local mu = {}
	local var = {}
	local prior = table.create(C, 0)
	local fitted = false
	local self
	self = {
		fit = function(_, X, y: { number })
			local D, N = X._shape[1], X._shape[2]
			task.desynchronize()
			for c = 1, C do
				mu[c] = table.create(D, 0)
				var[c] = table.create(D, 0)
				prior[c] = 0
			end
			for n = 1, N do
				prior[y[n]] += 1
			end
			for c = 1, C do
				if prior[c] == 0 then
					prior[c] = 1
				end
			end
			for c = 1, C do
				for d = 1, D do
					local s = 0
					for n = 1, N do
						if y[n] == c then
							s += X._storage[(d - 1) * N + n]
						end
					end
					mu[c][d] = s / prior[c]
				end
			end
			for c = 1, C do
				for d = 1, D do
					local s = 0
					for n = 1, N do
						if y[n] == c then
							local dx = X._storage[(d - 1) * N + n] - mu[c][d]
							s += dx * dx
						end
					end
					var[c][d] = math.max(s / prior[c], 1e-6)
				end
			end
			for c = 1, C do
				prior[c] = prior[c] / N
			end
			task.synchronize()
			fitted = true
		end,
		predict = function(_, X)
			assert(fitted, "GaussianNB: fit first")
			local D, N = X._shape[1], X._shape[2]
			local out = table.create(N, 1)
			task.desynchronize()
			for n = 1, N do
				local best = -1 / 0
				local arg = 1
				for c = 1, C do
					local logp = math.log(prior[c])
					for d = 1, D do
						local m = mu[c][d]
						local v = var[c][d]
						local x = X._storage[(d - 1) * N + n]
						logp += -0.5 * math.log(2 * math.pi * v) - (x - m) * (x - m) / (2 * v)
					end
					if logp > best then
						best = logp
						arg = c
					end
				end
				out[n] = arg
			end
			task.synchronize()
			return out
		end,
	}
	return self
end


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

return function(K: number, iters: number)
	iters = iters or 20
	local centers: any = nil
	local function initPlusPlus(X)
		local D, N = X._shape[1], X._shape[2]
		local C = Tensor.zeros({ D, K })
		local first = math.random(1, N)
		for i = 1, D do
			C._storage[(i - 1) * K + 1] = X._storage[(i - 1) * N + first]
		end
		local dist = table.create(N, math.huge)
		task.desynchronize()
		for c = 2, K do
			for n = 1, N do
				local best = dist[n]
				for cc = 1, c - 1 do
					local s = 0
					for d = 1, D do
						local dx = X._storage[(d - 1) * N + n] - C._storage[(d - 1) * K + cc]
						s += dx * dx
					end
					if s < best then
						best = s
					end
				end
				dist[n] = best
			end
			local sum = 0
			for n = 1, N do
				sum += dist[n]
			end
			local r = math.random() * math.max(sum, 1e-9)
			local acc = 0
			local pick = N
			for n = 1, N do
				acc += dist[n]
				if acc >= r then
					pick = n
					break
				end
			end
			for d = 1, D do
				C._storage[(d - 1) * K + c] = X._storage[(d - 1) * N + pick]
			end
		end
		task.synchronize()
		return C
	end
	local self
	self = {
		fit = function(_, X)
			local D, N = X._shape[1], X._shape[2]
			centers = initPlusPlus(X)
			local labels = table.create(N, 1)
			task.desynchronize()
			for _ = 1, iters do
				for n = 1, N do
					local best = math.huge
					local li = 1
					for k = 1, K do
						local s = 0
						for d = 1, D do
							local dx = X._storage[(d - 1) * N + n] - centers._storage[(d - 1) * K + k]
							s += dx * dx
						end
						if s < best then
							best = s
							li = k
						end
					end
					labels[n] = li
				end
				local counts = table.create(K, 0)
				local newC = Tensor.zeros({ D, K })
				for n = 1, N do
					local k = labels[n]
					counts[k] += 1
					for d = 1, D do
						newC._storage[(d - 1) * K + k] += X._storage[(d - 1) * N + n]
					end
				end
				for k = 1, K do
					local c = math.max(counts[k], 1)
					for d = 1, D do
						newC._storage[(d - 1) * K + k] /= c
					end
				end
				centers = newC
			end
			task.synchronize()
			return centers
		end,
		predict = function(_, X)
			assert(centers, "KMeans: fit first")
			local D, N = X._shape[1], X._shape[2]
			local labels = table.create(N, 1)
			task.desynchronize()
			for n = 1, N do
				local best = math.huge
				local li = 1
				for k = 1, K do
					local s = 0
					for d = 1, D do
						local dx = X._storage[(d - 1) * N + n] - centers._storage[(d - 1) * K + k]
						s += dx * dx
					end
					if s < best then
						best = s
						li = k
					end
				end
				labels[n] = li
			end
			task.synchronize()
			return labels
		end,
		centers = function(_)
			return centers
		end,
	}
	return self
end


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

--@class KNN
--@field X any
--@field y {number}
--@field k integer

local function predict(self, Q: Types.Tensor): { number }
	task.desynchronize()
	local D, N = self.X._shape[1], self.X._shape[2]
	local B = Q._shape[2]
	local out = table.create(B, 1)
	for j = 1, B do
		local distIdx = table.create(N)
		for n = 1, N do
			local s = 0.0
			for i = 1, D do
				local a = self.X._storage[(i - 1) * N + n]
				local b = Q._storage[(i - 1) * B + j]
				local d = a - b
				s += d * d
			end
			distIdx[n] = { s, n }
		end
		table.sort(distIdx, function(a, b)
			return a[1] < b[1]
		end)
		local votes = {}
		local K = math.min(self.k, N)
		for r = 1, K do
			local label = self.y[distIdx[r][2]]
			votes[label] = (votes[label] or 0) + 1
		end
		local best, arg = -1, 1
		for lbl, cnt in votes do
			if cnt > best then
				best, arg = cnt, lbl
			end
		end
		out[j] = arg
	end
	task.synchronize()
	return out
end

local function new(X: Types.Tensor, y: { number }, k: number?)
	assert(#y == X._shape[2], "KNN: y length must equal N")
	return {
		X = X,
		y = y,
		k = k or 5,
		predict = predict,
	}
end

return new


--!native
--!optimize 2
local Tape = require("../autograd/Tape")
local Losses = require("../nn/Losses")
local Linear = require("../nn/Linear")
local SGD = require("../optim/SGD")

return function(inF: number)
	local model = Linear(inF, 1)
	local self
	self = {
		fit = function(_, X, Y, epochs: number, lr: number)
			local opt = SGD(model:parameters(), lr or 1e-2)
			for _ = 1, epochs do
				opt:zeroGrad()
				local pred = model:forward(X)
				local _, grad = Losses.mse_backward(pred, Y)
				Tape.backwardFrom(pred, grad)
				opt:step()
			end
		end,
		predict = function(_, X)
			return model:forward(X)
		end,
		parameters = function(_)
			return model:parameters()
		end,
	}
	return self
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Losses = require("../nn/Losses")
local Linear = require("../nn/Linear")
local Adam = require("../optim/Adam")

return function(inF: number)
	local model = Linear(inF, 1)
	local self
	self = {
		fit = function(_, X, targets: { number }, epochs: number, lr: number)
			local opt = Adam(model:parameters(), lr or 1e-3)
			for _ = 1, epochs do
				opt:zeroGrad()
				local logits = model:forward(X)
				local _, grad = Losses.bceWithLogits_backward(logits, targets)
				Tape.backwardFrom(logits, grad)
				opt:step()
			end
		end,
		predictProba = function(_, X)
			local logits = model:forward(X)
			local out = Tensor.zeros(logits._shape)
			task.desynchronize()
			for i = 1, #logits._storage do
				out._storage[i] = 1 / (1 + math.exp(-logits._storage[i]))
			end
			task.synchronize()
			return out
		end,
		predict = function(self, X, threshold: number?)
			threshold = threshold or 0.5
			local p = self:predictProba(X)
			local r = table.create(p._shape[2], 0)
			task.desynchronize()
			for j = 1, #r do
				r[j] = (p._storage[j] >= threshold) and 1 or 0
			end
			task.synchronize()
			return r
		end,
		parameters = function(_)
			return model:parameters()
		end,
	}
	return self
end


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

return function(C: number, alpha: number)
	alpha = alpha or 1.0
	local classCount = table.create(C, 0)
	local featureCount = {}
	local fitted = false
	local self
	self = {
		fit = function(_, X, y)
			local D, N = X._shape[1], X._shape[2]
			task.desynchronize()
			for c = 1, C do
				featureCount[c] = table.create(D, 0)
				classCount[c] = 0
			end
			for n = 1, N do
				local c = y[n]
				classCount[c] += 1
				for d = 1, D do
					featureCount[c][d] += math.max(0, X._storage[(d - 1) * N + n])
				end
			end
			task.synchronize()
			fitted = true
		end,
		predict = function(_, X)
			assert(fitted, "MultinomialNB: fit first")
			local D, N = X._shape[1], X._shape[2]
			local out = table.create(N, 1)
			task.desynchronize()
			for n = 1, N do
				local best = -1 / 0
				local arg = 1
				for c = 1, C do
					local sum = 0
					for d = 1, D do
						sum += featureCount[c][d]
					end
					local logp = math.log((classCount[c] + alpha) / (N + C * alpha))
					for d = 1, D do
						local p = (featureCount[c][d] + alpha) / (sum + D * alpha)
						local x = math.max(0, X._storage[(d - 1) * N + n])
						if p > 0 then
							logp += x * math.log(p)
						end
					end
					if logp > best then
						best = logp
						arg = c
					end
				end
				out[n] = arg
			end
			task.synchronize()
			return out
		end,
	}
	return self
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Tape = require("../autograd/Tape")
local Losses = require("../nn/Losses")
local Linear = require("../nn/Linear")
local SGD = require("../optim/SGD")

return function(inF: number, wd: number)
	wd = wd or 1e-2
	local model = Linear(inF, 1)
	local self
	self = {
		fit = function(_, X, Y, epochs: number, lr: number)
			local opt = SGD(model:parameters(), lr or 1e-2, 0.9, true)
			for _ = 1, epochs do
				opt:zeroGrad()
				local pred = model:forward(X)
				local _, grad = Losses.mse_backward(pred, Y)
				local W = model.W
				if not W._grad then
					W._grad = Tensor.zeros(W._shape)
				end
				task.desynchronize()
				for i = 1, #W._storage do
					W._grad._storage[i] += 2 * wd * W._storage[i]
				end
				task.synchronize()
				Tape.backwardFrom(pred, grad)
				opt:step()
			end
		end,
		predict = function(_, X)
			return model:forward(X)
		end,
		parameters = function(_)
			return model:parameters()
		end,
	}
	return self
end


--!native
--!optimize 2
local Linear = require("../nn/Linear")
local Losses = require("../nn/Losses")
local Tape = require("../autograd/Tape")
local Tensor = require("../Tensor")
local Types = require("../Types")
local Adam = require("../optim/Adam")
local Util = require("../Util")

local assert = Util.Assert

local function slice(
	N: number,
	X: Types.Tensor,
	D: number,
	y: { number },
	startIdx: number,
	size: number
): (Types.Tensor, { number })
	local B = math.min(size, N - startIdx + 1)
	local XB = Tensor.zeros({ D, B }, X._dtype, X._requiresGrad)
	local yB = table.create(B, 1)
	task.desynchronize()
	for j = 1, B do
		local srcj = startIdx + j - 1
		for i = 1, D do
			XB._storage[(i - 1) * B + j] = X._storage[(i - 1) * N + srcj]
		end
		yB[j] = y[srcj]
	end
	task.synchronize()
	return XB, yB
end

local function forward(self, X: Types.Tensor): Types.Tensor
	return self.fc:forward(X)
end

local function parameters(self): { Types.Tensor }
	return self.fc:parameters()
end

local function fit(self, X: Types.Tensor, y: { number }, lr: number, epochs: number, batch: number, smoothing: number)
	lr, epochs, batch = lr or 1e-2, epochs or 5, batch or 64
	smoothing = smoothing or 0
	local D, N = X._shape[1], X._shape[2]
	assert(#y == N, "y length must equal N")

	task.desynchronize()
	local opt = Adam(parameters(self), lr)
	for _ = 1, epochs do
		for s = 1, N, batch do
			local XB, yB = slice(N, X, D, y, s, batch)
			opt:zeroGrad()
			local logits = forward(self, XB)
			local _, dlogits = Losses.cross_entropy_backward(logits, yB, smoothing)
			Tape.backwardFrom(logits, dlogits)
			opt:step()
		end
	end
	task.synchronize()
end

local function new(inDim: number, classes: number): Types.Module<Types.Tensor, Types.Tensor>
	return {
		fc = Linear(inDim, classes),
		forward = forward,
		parameters = parameters,
		fit = fit,
	}
end

return new


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

return function(
	dataset: Types.Dataset,
	batchSize: number,
	shuffle: boolean,
	generator: { randint: (a: number, b: number) -> number }?,
	drop_last: boolean?
)
	local N = dataset:size()
	generator = generator or {
		randint = function(a: number, b: number)
			return math.random(a, b)
		end,
	}
	local order = table.create(N, 0)
	for i = 1, N do
		order[i] = i
	end
	if shuffle then
		for i = N, 2, -1 do
			local j = generator.randint(1, i)
			order[i], order[j] = order[j], order[i]
		end
	end
	local idx = 1
	return function()
		if idx > N or (drop_last and (N - idx + 1) < batchSize) then
			return nil
		end
		local last = math.min(idx + batchSize - 1, N)
		local count = last - idx + 1
		local batchIdx = table.create(count)
		for i = 1, count do
			batchIdx[i] = order[idx + i - 1]
		end
		idx = last + 1
		if dataset.batch then
			return dataset:batch(batchIdx)
		elseif dataset.slice and not shuffle then
			return dataset:slice(batchIdx[1], batchIdx[#batchIdx])
		end
		return dataset:at(batchIdx[1])
	end
end


--!native
--!optimize 2
local Types = require("../Types")
return function(N: number, K: number): Types.BatchIter
	local order: { number } = {}
	for i = 1, N do
		order[i] = i
	end
	for i = N, 2, -1 do
		local j = math.random(i)
		order[i], order[j] = order[j], order[i]
	end
	local fold = 1
	return function()
		if fold > K then
			return nil
		end
		local size = math.floor(N / K)
		local s = (fold - 1) * size + 1
		local e = (fold == K) and N or (s + size - 1)
		local test: { number } = {}
		for i = s, e do
			table.insert(test, order[i])
		end
		local train: { number } = {}
		for i = 1, N do
			if i < s or i > e then
				table.insert(train, order[i])
			end
		end
		fold += 1
		return train, test
	end
end


--!native
--!optimize 2

--@param N number -- total items
--@param testRatio number? -- fraction in test split (default 0.2)
--@param stratify {number}? -- optional class per item (1..C), length N
--@param rng Random.new? -- optional RNG for deterministic splits
--@return {number}, {number} -- trainIdx, testIdx (1-based indices into your data)
return function(N: number, testRatio: number?, stratify: { number }?, rng: Random?): ({ number }, { number })
	rng = rng or Random.new()
	testRatio = testRatio or 0.2
	local idx = table.create(N, 0)
	for i = 1, N do
		idx[i] = i
	end
	if stratify then
		local by = {}
		for i = 1, N do
			local c = stratify[i]
			if not by[c] then
				by[c] = {}
			end
			table.insert(by[c], i)
		end
		local train, test = {}, {}
		for _, arr in by do
			for i = #arr, 2, -1 do
				local j = rng:NextInteger(1, i)
				arr[i], arr[j] = arr[j], arr[i]
			end
			local cut = math.floor(#arr * (1 - testRatio :: number))
			for i = 1, #arr do
				if i <= cut then
					table.insert(train, arr[i])
				else
					table.insert(test, arr[i])
				end
			end
		end
		return train, test
	end
	for i = N, 2, -1 do
		local j = rng:NextInteger(1, i)
		idx[i], idx[j] = idx[j], idx[i]
	end
	local cut = math.floor(N * (1 - testRatio :: number))
	local train = {}
	local test = {}
	for i = 1, N do
		if i <= cut then
			table.insert(train, idx[i])
		else
			table.insert(test, idx[i])
		end
	end
	return train, test
end


--!native
--!optimize 2
local Tensor = require("../Tensor")
local Types = require("../Types")
return function(X: Types.Tensor, Y: Types.Tensor): Types.Dataset
	local D, N = X._shape[1], X._shape[2]
	local function colSlice(startIdx: number, endIdx: number): (Types.Tensor, Types.Tensor)
		local B = endIdx - startIdx + 1
		local XB = Tensor.zeros({ D, B }, X._dtype, X._requiresGrad)
		for d = 1, D do
			local base = (d - 1) * N
			for b = 1, B do
				XB._storage[(d - 1) * B + b] = X._storage[base + (startIdx + b - 1)]
			end
		end
		local yB = table.create(B, 1)
		for b = 1, B do
			yB[b] = Y[startIdx + b - 1]
		end
		return XB, yB
	end
	local function colGather(indices: { number }): (Types.Tensor, Types.Tensor)
		local B = #indices
		local XB = Tensor.zeros({ D, B }, X._dtype, X._requiresGrad)
		for d = 1, D do
			local base = (d - 1) * N
			for b = 1, B do
				XB._storage[(d - 1) * B + b] = X._storage[base + indices[b]]
			end
		end
		local yB = table.create(B, 1)
		for b = 1, B do
			yB[b] = Y[indices[b]]
		end
		return XB, yB
	end
	return {
		size = function(self): number
			return N
		end,
		at = function(self, i: number): (Types.Tensor, Types.Tensor)
			return colSlice(i, i)
		end,
		slice = function(self, i: number, j: number): (Types.Tensor, Types.Tensor)
			return colSlice(i, j)
		end,
		batch = function(self, indices: { number }): (Types.Tensor, Types.Tensor)
			return colGather(indices)
		end,
	}
end


--!native
--!optimize 2
-- Quantum-Inspired Metaheuristic Neural Network (QIMHNN)
local Types = require("../../Types")
local Tensor = require("../../Tensor")
local Tape = require("../../autograd/Tape")
local Util = require("../../Util")
local Utils = require("../../models/Utils")
local LayerNorm = require("../../nn/LayerNorm")
local Dropout = require("../../nn/Dropout")
local Linear = require("../../nn/Linear")
local Sequential = require("../../nn/Sequential")

local assert = Util.Assert
local sizeFromShape = Util.sizeFromShape

local function sameShape(a: { number }, b: { number }): boolean
	if #a ~= #b then
		return false
	end
	for i = 1, #a do
		if a[i] ~= b[i] then
			return false
		end
	end
	return true
end

local function ensureScratchTensor(buf: Types.Tensor?, shape: { number }): Types.Tensor
	if not buf or not sameShape(buf._shape, shape) then
		return Tensor.zeros(shape, "f64", false)
	end
	task.desynchronize()
	for i = 1, #buf._storage do
		buf._storage[i] = 0
	end
	task.synchronize()
	return buf
end

export type QILayerConfig = {
	inDim: number,
	outDim: number,
	useSuperposition: boolean?,
	useInterference: boolean?,
	useEntanglement: boolean?,
	quantumAmplitude: number?, -- amp for |z|
	initScale: number?,
}

local function QILayer(config: QILayerConfig): Types.Module<Types.Tensor, Types.Tensor>
	assert(config.inDim and config.inDim > 0, "QILayer: inDim required")
	assert(config.outDim and config.outDim > 0, "QILayer: outDim required")

	local useSuperposition = config.useSuperposition ~= false
	local useInterference = config.useInterference ~= false
	local useEntanglement = config.useEntanglement ~= false
	local quantumAmplitude = config.quantumAmplitude or 0.1
	local initScale = config.initScale or 0.1

	local interferenceAlpha = 0.1
	local entangleBeta = 0.1
	local eps = 1e-8

	local W_real = Tensor.zeros({ config.outDim, config.inDim }, "f64", true)
	local W_imag = Tensor.zeros({ config.outDim, config.inDim }, "f64", true) -- used when superposition OR interference enabled
	local b = Tensor.zeros({ config.outDim, 1 }, "f64", true)
	local entanglement = if useEntanglement then Tensor.zeros({ config.outDim, config.outDim }, "f64", true) else nil

	task.desynchronize()
	local fanIn, fanOut = config.inDim, config.outDim
	local bound = initScale * math.sqrt(2 / (fanIn + fanOut))
	for i = 1, #W_real._storage do
		W_real._storage[i] = (math.random() * 2 - 1) * bound
		W_imag._storage[i] = (math.random() * 2 - 1) * bound * 0.5
	end
	if entanglement then
		for i = 1, #entanglement._storage do
			entanglement._storage[i] = (math.random() * 2 - 1) * initScale * 0.1
		end
	end
	task.synchronize()
	return {
		W_real = W_real,
		W_imag = W_imag,
		b = b,
		entanglement = entanglement,
		quantumAmplitude = quantumAmplitude,
		useSuperposition = useSuperposition,
		useInterference = useInterference,
		useEntanglement = useEntanglement,
		_scratch = {},
		forward = function(self, x: Types.Tensor): Types.Tensor
			local B = x._shape[2]
			local outF = config.outDim
			local scratch = self._scratch
			-- linear
			local z_real = Tape.matmul(self.W_real, x)
			local hasImag = (self.useSuperposition or self.useInterference)
			local z_imag: Types.Tensor? = nil
			if hasImag then
				z_imag = Tape.matmul(self.W_imag, x)
			end
			-- y_sup: superposition amplitude mix
			local y_sup = ensureScratchTensor(scratch.y_sup, { outF, B })
			scratch.y_sup = y_sup
			task.desynchronize()
			for i = 1, outF do
				for j = 1, B do
					local idx = (i - 1) * B + j
					local zr = z_real._storage[idx]
					if self.useSuperposition and z_imag then
						local zi = z_imag._storage[idx]
						local mag = math.sqrt(zr * zr + zi * zi + eps)
						y_sup._storage[idx] = zr + self.quantumAmplitude * mag
					else
						y_sup._storage[idx] = zr
					end
				end
			end
			task.synchronize()
			-- scale by 1 + alpha * cos(phase), phase = atan2(zi, zr)
			local y_int = ensureScratchTensor(scratch.y_int, { outF, B })
			scratch.y_int = y_int
			task.desynchronize()
			for i = 1, outF do
				for j = 1, B do
					local idx = (i - 1) * B + j
					local scale = 1
					if self.useInterference then
						local zr = z_real._storage[idx]
						local zi = if z_imag then z_imag._storage[idx] else 0
						local phase = math.atan2(zi, zr)
						scale = 1 + interferenceAlpha * math.cos(phase)
					end
					y_int._storage[idx] = y_sup._storage[idx] * scale
				end
			end
			task.synchronize()
			local stored_y_int: Types.Tensor = nil
			if self.useEntanglement and self.entanglement then
				stored_y_int = ensureScratchTensor(scratch.stored_y_int, { outF, B })
				scratch.stored_y_int = stored_y_int
				task.desynchronize()
				for i = 1, #stored_y_int._storage do
					stored_y_int._storage[i] = y_int._storage[i]
				end
				task.synchronize()
			end
			-- residual E @ y_int
			local y_preBias = y_int
			local entangled: Types.Tensor = nil
			if self.useEntanglement and self.entanglement then
				entangled = Tape.matmul(self.entanglement, y_int)
				task.desynchronize()
				for i = 1, #y_preBias._storage do
					y_preBias._storage[i] = y_preBias._storage[i] + entangleBeta * entangled._storage[i]
				end
				task.synchronize()
			end
			-- bias
			task.desynchronize()
			for i = 1, outF do
				local bi = self.b._storage[i]
				for j = 1, B do
					y_preBias._storage[(i - 1) * B + j] += bi
				end
			end
			task.synchronize()
			local stored_z_real = z_real
			local stored_z_imag = z_imag
			-- local stored_y_sup = y_sup
			local usedImag = hasImag
			local parents: { Types.Tensor } = { self.W_real, self.b, x }
			if self.useSuperposition or self.useInterference then
				table.insert(parents, self.W_imag)
			end
			if self.useEntanglement and self.entanglement then
				table.insert(parents, self.entanglement)
			end
			return Tape.makeNode(y_preBias, parents, function(gy: Types.Tensor)
				-- gy is dL/d(y_preBias)
				local outF = config.outDim
				local inF = config.inDim
				local B = gy._shape[2]
				if self.b._requiresGrad then
					if not self.b._grad then
						self.b._grad = Tensor.zeros({ outF, 1 })
					end
					task.desynchronize()
					for i = 1, outF do
						local acc = 0
						for j = 1, B do
							acc += gy._storage[(i - 1) * B + j]
						end
						self.b._grad._storage[i] += acc
					end
					task.synchronize()
				end
				-- y_preBias = y_int + beta * (E @ y_int)
				local gy_int = gy -- start with gy
				if self.useEntanglement and self.entanglement then
					-- dE += beta * gy @ y_int^T
					if self.entanglement._requiresGrad then
						if not self.entanglement._grad then
							self.entanglement._grad = Tensor.zeros({ outF, outF })
						end
						local yIntT: Types.Tensor = nil
						Tape.noGrad(function()
							yIntT = Tensor.transpose(stored_y_int)
						end)
						local dEnt
						Tape.noGrad(function()
							dEnt = Tape.matmul(gy, yIntT)
						end)
						task.desynchronize()
						for i = 1, #self.entanglement._grad._storage do
							self.entanglement._grad._storage[i] += entangleBeta * dEnt._storage[i]
						end
						task.synchronize()
					end
					-- E: gy_int = gy + beta * (E^T @ gy)
					local ET
					Tape.noGrad(function()
						ET = Tensor.transpose(self.entanglement)
					end)
					local add
					Tape.noGrad(function()
						add = Tape.matmul(ET, gy)
					end)
					local tmp = ensureScratchTensor(scratch.gy_int_tmp, { outF, B })
					scratch.gy_int_tmp = tmp
					task.desynchronize()
					for i = 1, #tmp._storage do
						tmp._storage[i] = gy._storage[i] + entangleBeta * add._storage[i]
					end
					task.synchronize()
					gy_int = tmp
				end
				local gz_real = ensureScratchTensor(scratch.gz_real, { outF, B })
				scratch.gz_real = gz_real
				local gz_imag = nil
				if usedImag then
					gz_imag = ensureScratchTensor(scratch.gz_imag, { outF, B })
					scratch.gz_imag = gz_imag
				end
				task.desynchronize()
				for i = 1, outF do
					for j = 1, B do
						local idx = (i - 1) * B + j
						local zr = stored_z_real._storage[idx]
						local zi = if stored_z_imag then stored_z_imag._storage[idx] else 0
						local denom = zr * zr + zi * zi + eps
						local mag = math.sqrt(denom)

						local ysup = if self.useSuperposition then (zr + self.quantumAmplitude * mag) else zr

						local M = 1
						local phase = 0
						if self.useInterference then
							phase = math.atan2(zi, zr)
							M = 1 + interferenceAlpha * math.cos(phase)
						end
						-- gy wrt y_sup and wrt M
						local gyij = gy_int._storage[idx]
						local gy_sup = gyij * M
						local gM = gyij * ysup
						local gzR_phase, gzI_phase = 0, 0
						if self.useInterference then
							local gphase = -interferenceAlpha * gM * math.sin(phase)
							-- d phase/dzr = -zi/denom ; d phase/dzi = zr/denom
							gzR_phase = gphase * (-zi / denom)
							gzI_phase = gphase * (zr / denom)
						end
						-- superposition path
						local gzR_sup, gzI_sup = 0, 0
						if self.useSuperposition then
							gzR_sup = gy_sup * (1 + self.quantumAmplitude * (zr / mag))
							gzI_sup = gy_sup * (self.quantumAmplitude * (zi / mag))
						else
							gzR_sup = gy_sup
							gzI_sup = 0
						end
						local totalR = gzR_sup + gzR_phase
						gz_real._storage[idx] = gz_real._storage[idx] + totalR
						if gz_imag then
							local totalI = gzI_sup + gzI_phase
							gz_imag._storage[idx] = gz_imag._storage[idx] + totalI
						end
					end
				end
				task.synchronize()
				local xT
				Tape.noGrad(function()
					xT = Tensor.transpose(x)
				end)
				-- dW_real += gz_real @ x^T
				if self.W_real._requiresGrad then
					if not self.W_real._grad then
						self.W_real._grad = Tensor.zeros({ outF, inF })
					end
					local dW
					Tape.noGrad(function()
						dW = Tape.matmul(gz_real, xT)
					end)
					task.desynchronize()
					for i = 1, #self.W_real._grad._storage do
						self.W_real._grad._storage[i] += dW._storage[i]
					end
					task.synchronize()
				end
				-- dW_imag
				if gz_imag and self.W_imag._requiresGrad then
					if not self.W_imag._grad then
						self.W_imag._grad = Tensor.zeros({ outF, inF })
					end
					local dWi
					Tape.noGrad(function()
						dWi = Tape.matmul(gz_imag, xT)
					end)
					task.desynchronize()
					for i = 1, #self.W_imag._grad._storage do
						self.W_imag._grad._storage[i] += dWi._storage[i]
					end
					task.synchronize()
				end
				-- dx = W_real^T @ gz_real + (if used) W_imag^T @ gz_imag
				local dx: Types.Tensor = nil
				if x._requiresGrad then
					dx = ensureScratchTensor(scratch.dx, { inF, B })
					scratch.dx = dx
					local WT
					Tape.noGrad(function()
						WT = Tensor.transpose(self.W_real)
					end)
					local dxr
					Tape.noGrad(function()
						dxr = Tape.matmul(WT, gz_real)
					end)
					task.desynchronize()
					for i = 1, #dx._storage do
						dx._storage[i] = dxr._storage[i]
					end
					task.synchronize()
					if gz_imag then
						local WIT
						Tape.noGrad(function()
							WIT = Tensor.transpose(self.W_imag)
						end)
						local dxi
						Tape.noGrad(function()
							dxi = Tape.matmul(WIT, gz_imag)
						end)
						task.desynchronize()
						for i = 1, #dx._storage do
							dx._storage[i] += dxi._storage[i]
						end
						task.synchronize()
					end
					if not x._grad then
						x._grad = Tensor.zeros({ inF, B })
					end
					task.desynchronize()
					for i = 1, #x._grad._storage do
						x._grad._storage[i] += dx._storage[i]
					end
					task.synchronize()
				end
				return nil -- no grad
			end)
		end,
		parameters = function(self): { Types.Tensor }
			local params = { self.W_real, self.b }
			if self.useSuperposition or self.useInterference then
				table.insert(params, 2, self.W_imag)
			end
			if self.entanglement then
				table.insert(params, self.entanglement)
			end
			return params
		end,
	}
end

export type QIMHNNConfig = {
	inputDim: number,
	outputDim: number,
	hiddenDims: { number }?,
	activation: Utils.ActivationName?,
	finalActivation: Utils.ActivationName?,
	useSuperposition: boolean?,
	useInterference: boolean?,
	useEntanglement: boolean?,
	quantumAmplitude: number?,
	dropout: number?,
	layerNorm: boolean?,
}

return function(config: QIMHNNConfig): Types.Module<Types.Tensor, Types.Tensor>
	assert(config.inputDim, "QIMHNN: inputDim is required")
	assert(config.outputDim, "QIMHNN: outputDim is required")

	local hidden = if config.hiddenDims and #config.hiddenDims > 0
		then table.clone(config.hiddenDims)
		else { math.max(config.inputDim, config.outputDim) }

	local dropProb = math.clamp(config.dropout or 0, 0, 0.95)
	local layers = {}
	local trainables: { any } = {}
	local inDim = config.inputDim

	for _, hiddenDim in ipairs(hidden) do
		assert(hiddenDim > 0, "QIMHNN: hidden dimensions must be positive")

		local qiLayer = QILayer({
			inDim = inDim,
			outDim = hiddenDim,
			useSuperposition = config.useSuperposition,
			useInterference = config.useInterference,
			useEntanglement = config.useEntanglement,
			quantumAmplitude = config.quantumAmplitude or 0.1,
		})
		table.insert(layers, qiLayer)
		table.insert(trainables, qiLayer)

		if config.layerNorm then
			local norm = LayerNorm(hiddenDim)
			table.insert(layers, norm)
			table.insert(trainables, norm)
		end

		table.insert(layers, Utils.makeActivationModule(config.activation))

		if dropProb > 0 then
			local drop = Dropout(dropProb)
			table.insert(layers, drop)
			table.insert(trainables, drop)
		end
		inDim = hiddenDim
	end

	local head = Linear(inDim, config.outputDim)
	table.insert(layers, head)
	table.insert(trainables, head)

	if config.finalActivation then
		table.insert(layers, Utils.makeActivationModule(config.finalActivation))
	end

	local seq = Sequential(layers)

	return {
		forward = function(_, x: Types.Tensor): Types.Tensor
			return seq:forward(x)
		end,
		act = function(_, x: Types.Tensor): Types.Tensor
			return seq:forward(x)
		end,
		parameters = function(_): { Types.Tensor }
			return seq:parameters()
		end,
		train = function(_, mode: boolean?)
			Utils.trainAll(trainables, mode)
		end,
	}
end


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

local function RBF(x: Types.Tensor, grid: Types.Tensor, h: number): Types.Tensor
	-- x: {B, In}
	-- grid: {G} (centers)
	-- out: {B, In, G}
	local B, In = x._shape[1], x._shape[2]
	local G = grid._shape[1]
	local out = Tensor.zeros({ B, In, G }, x._dtype, x._requiresGrad or grid._requiresGrad)
	local xs = x._storage
	local gs = grid._storage
	local os = out._storage
	task.desynchronize()
	for b = 1, B do
		for i = 1, In do
			local idx_x = (b - 1) * In + i
			local val = xs[idx_x]
			for g = 1, G do
				local c = gs[g]
				local dist = (val - c) / h
				local rbf = math.exp(-dist * dist)
				-- out: {B, In, G}
				-- index: (b-1)*In*G + (i-1)*G + g
				local idx_o = ((b - 1) * In + (i - 1)) * G + g
				os[idx_o] = rbf
			end
		end
	end
	task.synchronize()
	return Tape.makeNode(out, { x, grid }, function(gy)
		if x._requiresGrad then
			if not x._grad then
				x._grad = Tensor.zeros(x._shape, x._dtype, false)
			end
			task.desynchronize()
			local x_grad = x._grad :: Types.Tensor
			local dxs = x_grad._storage
			local gys = gy._storage
			for b = 1, B do
				for i = 1, In do
					local idx_x = (b - 1) * In + i
					local val = xs[idx_x]
					local acc = 0
					for g = 1, G do
						local idx_o = ((b - 1) * In + (i - 1)) * G + g
						local gy_val = gys[idx_o]
						local c = gs[g]
						local dist = (val - c) / h
						local rbf = os[idx_o] -- exp term
						-- d/dx exp(-((x-c)/h)^2) = exp(...) * -2((x-c)/h) * 1/h
						local grad_val = rbf * (-2 * dist / h)
						acc += gy_val * grad_val
					end
					dxs[idx_x] += acc
				end
			end
			task.synchronize()
		end
	end)
end

local function KANLayer(in_features: number, out_features: number, grid_size: number?, spline_order: number?)
	grid_size = grid_size or 5
	spline_order = spline_order or 3
	local h = 1.0 / (grid_size :: number)
	local base_linear = Linear(in_features, out_features)
	local spline_linear = Linear(in_features * (grid_size :: number), out_features)
	local grid = Tensor.fromArray({}, { grid_size :: number }, "f64", false)
	local step = 2.0 / ((grid_size :: number) - 1)
	for i = 1, grid_size :: number do
		grid._storage[i] = -1 + (i - 1) * step
	end
	return {
		forward = function(x: Types.Tensor)
			local input_shape = x._shape
			local rank = #input_shape
			local flattened_dim = 1
			for i = 1, rank - 1 do
				flattened_dim *= input_shape[i]
			end
			local x_flat = Tensor.reshape(x, { flattened_dim, in_features })
			local x_silu = Activations.SiLU(x_flat)
			local base_out = base_linear:forward(x_silu)
			local x_rbf = RBF(x_flat, grid, h) -- {N, In, G}
			local x_rbf_flat = Tensor.reshape(x_rbf, { flattened_dim, in_features * (grid_size :: number) })
			local spline_out = spline_linear:forward(x_rbf_flat)
			local out = Tensor.zeros(base_out._shape, base_out._dtype, true)
			task.desynchronize()
			local os, bs, ss = out._storage, base_out._storage, spline_out._storage
			for i = 1, #os do
				os[i] = bs[i] + ss[i]
			end
			task.synchronize()
			out = Tape.makeNode(out, { base_out, spline_out }, function(gy)
				local gys = gy._storage
				if base_out._requiresGrad then
					if not base_out._grad then
						base_out._grad = Tensor.zeros(base_out._shape)
					end
					local gbs = base_out._grad._storage
					task.desynchronize()
					for i = 1, #gys do
						gbs[i] += gys[i]
					end
					task.synchronize()
				end
				if spline_out._requiresGrad then
					if not spline_out._grad then
						spline_out._grad = Tensor.zeros(spline_out._shape)
					end
					local gss = spline_out._grad._storage
					task.desynchronize()
					for i = 1, #gys do
						gss[i] += gys[i]
					end
					task.synchronize()
				end
			end)
			local out_shape = table.clone(input_shape)
			out_shape[rank] = out_features
			return Tensor.reshape(out, out_shape)
		end,
		parameters = function(_)
			local p = {}
			for _, v in base_linear:parameters() do
				table.insert(p, v)
			end
			for _, v in spline_linear:parameters() do
				table.insert(p, v)
			end
			return p
		end,
	}
end

return KANLayer


--!native
--!optimize 2
local Tensor = require("../../Tensor")
local Tape = require("../../autograd/Tape")
local Types = require("../../Types")
local Scan = require("../../ops/Scan")
local Initializer = require("../../Initializer")
local BLAS = require("../../ops/BLAS")
local Math = require("../../ops/Math")
local Activations = require("../../nn/Activations")

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

local function depthwiseConv1d(x: Types.Tensor, weight: Types.Tensor): Types.Tensor
	-- x: {B, L, D}
	if not Tensor.is_contiguous(x) then
		x = Tensor.contiguous(x)
	end
	if not Tensor.is_contiguous(weight) then
		weight = Tensor.contiguous(weight)
	end
	local B, L, D = x._shape[1], x._shape[2], x._shape[3]
	local K = weight._shape[3]
	local y = Tensor.zeros({ B, L, D }, x._dtype, x._requiresGrad or weight._requiresGrad)
	task.desynchronize()
	local xs = x._storage
	local ws = weight._storage
	local ys = y._storage
	for b = 0, B - 1 do
		local b_offset = b * L * D
		for l = 0, L - 1 do
			local l_offset = b_offset + l * D
			for d = 0, D - 1 do
				local sum = 0
				for k = 0, K - 1 do
					if l - k >= 0 then
						local x_idx = b_offset + (l - k) * D + d + 1
						local w_idx = d * K + k + 1
						sum += xs[x_idx] * ws[w_idx]
					end
				end
				ys[l_offset + d + 1] = sum
			end
		end
	end
	task.synchronize()
	return Tape.makeNode(y, { x, weight }, function(gy)
		if not x._requiresGrad and not weight._requiresGrad then
			return
		end
		local gx = x._requiresGrad and ensureGrad(x)
		local gw = weight._requiresGrad and ensureGrad(weight)
		local gys = gy._storage
		task.desynchronize()
		local gxs = gx and gx._storage
		local gws = gw and gw._storage
		for b = 0, B - 1 do
			local b_offset = b * L * D
			for l = 0, L - 1 do
				local l_offset = b_offset + l * D
				for d = 0, D - 1 do
					local gy_val = gys[l_offset + d + 1]
					if gy_val == 0 then
						continue
					end
					for k = 0, K - 1 do
						if l - k >= 0 then
							local x_idx = b_offset + (l - k) * D + d + 1
							local w_idx = d * K + k + 1
							if gxs then
								gxs[x_idx] += gy_val * ws[w_idx]
							end
							if gws then
								gws[w_idx] += gy_val * xs[x_idx]
							end
						end
					end
				end
			end
		end
		task.synchronize()
	end)
end

local function linear(x: Types.Tensor, weight: Types.Tensor, bias: Types.Tensor?): Types.Tensor
	local inFeat = weight._shape[2]
	local outFeat = weight._shape[1]
	local xShape, xShape_len = x._shape, #x._shape
	local lastDim = xShape[xShape_len]
	if lastDim ~= inFeat then
		error(`Linear: input dimension mismatch. Expected {inFeat}, got {lastDim}`)
	end
	local flattenedDim = 1
	for i = 1, xShape_len - 1 do
		flattenedDim *= xShape[i]
	end
	local xFlat = Tensor.reshape(x, { flattenedDim, inFeat })
	local wT = Tensor.transpose(weight, 1, 2)
	local outFlat = BLAS.matmul(xFlat, wT)
	if bias then
		local outWithBias = Tensor.empty(outFlat._shape, outFlat._dtype, outFlat._requiresGrad or bias._requiresGrad)
		local b_storage = bias._storage
		local out_storage = outFlat._storage
		local res_storage = outWithBias._storage
		local outSize = #out_storage
		task.desynchronize()
		for i = 0, outSize - 1 do
			local c = (i % outFeat) + 1
			res_storage[i + 1] = out_storage[i + 1] + b_storage[c]
		end
		task.synchronize()
		if outWithBias._requiresGrad then
			outFlat = Tape.makeNode(outWithBias, { outFlat, bias }, function(gy)
				if outFlat._requiresGrad then
					local gout = ensureGrad(outFlat)
					local gys = gy._storage
					local gos = gout._storage
					task.desynchronize()
					for i = 1, #gys do
						gos[i] += gys[i]
					end
					task.synchronize()
				end
				if bias._requiresGrad then
					local gb = ensureGrad(bias)
					local gys = gy._storage
					local gbs = gb._storage
					task.desynchronize()
					for i = 0, #gys - 1 do
						local c = (i % outFeat) + 1
						gbs[c] += gys[i + 1]
					end
					task.synchronize()
				end
			end)
		else
			outFlat = outWithBias
		end
	end
	local outShape = table.clone(xShape)
	outShape[#outShape] = outFeat
	return Tensor.reshape(outFlat, outShape)
end

local function MambaBlock(dModel: number, dState: number?, dConv: number?, expand: number?)
	dState = dState or 16
	dConv = dConv or 4
	expand = expand or 2
	local dInner = dModel * (expand :: number)
	local dtRank = math.ceil(dModel / 16)
	local inProj = Tensor.empty({ dInner * 2, dModel }, "f64", true)
	Initializer.xavierUniform(inProj)
	local convWeight = Tensor.empty({ dInner, 1, dConv :: number }, "f64", true)
	Initializer.kaimingUniform(convWeight)
	local A_log = Tensor.empty({ dInner, dState :: number }, "f64", true)
	Initializer.kaimingUniform(A_log)
	local D = Tensor.ones({ dInner }, "f64", true)
	local xProj = Tensor.empty({ dtRank + (dState :: number) * 2, dInner }, "f64", true)
	local dtProj = Tensor.empty({ dInner, dtRank }, "f64", true)
	Initializer.xavierUniform(xProj)
	Initializer.xavierUniform(dtProj)
	local outProj = Tensor.empty({ dModel, dInner }, "f64", true)
	Initializer.xavierUniform(outProj)
	return {
		forward = function(_, u: Types.Tensor): Types.Tensor
			local _B_size, _L_size, D_size = u._shape[1], u._shape[2], u._shape[3]
			local xz = linear(u, inProj)
			local x = Tensor.slice(xz, 3, 1, dInner)
			local z = Tensor.slice(xz, 3, dInner + 1, dInner * 2)
			local x_conv = depthwiseConv1d(x, convWeight)
			local x_act = Activations.SiLU(x_conv)
			local x_dbl = linear(x_act, xProj)
			local dt_rank = Tensor.slice(x_dbl, 3, 1, dtRank)
			local B_ssm = Tensor.slice(x_dbl, 3, dtRank + 1, dtRank + (dState :: number))
			local C_ssm = Tensor.slice(x_dbl, 3, dtRank + (dState :: number) + 1, dtRank + (dState :: number) * 2)
			local dt = linear(dt_rank, dtProj)
			dt = Activations.Softplus(dt)
			local A = Math.scalarMul(Math.exp(A_log), -1)
			local y = Scan.selectiveScan(x_act, dt, A, B_ssm, C_ssm, D)
			local z_act = Activations.SiLU(z)
			local y_gated = Math.mul(y, z_act)
			local out = linear(y_gated, outProj)
			if D_size == dModel then
				return Math.add(out, u)
			end
			return out
		end,
		parameters = function()
			return { inProj, convWeight, A_log, D, xProj, dtProj, outProj }
		end,
	}
end

return MambaBlock


--!native
--!optimize 2
-- Metaheuristic Optimizer for QIMHNN
local Types = require("../../Types")
local Tensor = require("../../Tensor")
local Util = require("../../Util")

local assert = Util.Assert

export type MetaheuristicConfig = {
	lr: number,
	swarmSize: number?,
	inertia: number?,
	cognitive: number?,
	social: number?,
	gradScale: number?,
	mutationRate: number?,
	mutationStrength: number?,
}

return function(params: { Types.Tensor }, config: MetaheuristicConfig): Types.Optimizer
	assert(config.lr and config.lr > 0, "MetaheuristicOptimizer: lr required")

	local lr = config.lr
	local _swarmSize = config.swarmSize or 10
	local inertia = config.inertia or 0.9
	local cognitive = config.cognitive or 1.5
	local social = config.social or 1.5
	local gradScale = config.gradScale or 0.1
	local mutationRate = config.mutationRate or 0.05
	local mutationStrength = config.mutationStrength or 0.01

	local velocities: { Types.Tensor } = {}
	local personalBest: { Types.Tensor } = {}
	local globalBest: { Types.Tensor } = {}
	local personalBestScore: { number } = {}
	local globalBestScore: { number } = {}

	for i, param in ipairs(params) do
		local v = Tensor.zeros(param._shape, param._dtype, false)
		local pb = Tensor.zeros(param._shape, param._dtype, false)
		local gb = Tensor.zeros(param._shape, param._dtype, false)

		local pStore = param._storage
		local vStore = v._storage
		local pbStore = pb._storage
		local gbStore = gb._storage

		task.desynchronize()
		for j = 1, #pStore do
			vStore[j] = (math.random() * 2 - 1) * 0.001
			pbStore[j] = pStore[j]
			gbStore[j] = pStore[j]
		end
		task.synchronize()

		velocities[i] = v
		personalBest[i] = pb
		globalBest[i] = gb
		personalBestScore[i] = math.huge
		globalBestScore[i] = math.huge
	end
	return {
		step = function()
			for i, param in ipairs(params) do
				local grad = param._grad
				if grad then
					local gStore = grad._storage
					local gradNormSq = 0

					task.desynchronize()
					for j = 1, #gStore do
						local g = gStore[j]
						gradNormSq += g * g
					end
					task.synchronize()

					local score = gradNormSq
					if score < personalBestScore[i] then
						personalBestScore[i] = score
						local pb = personalBest[i]
						local pStore = param._storage
						local pbStore = pb._storage
						task.desynchronize()
						for j = 1, #pStore do
							pbStore[j] = pStore[j]
						end
						task.synchronize()
					end
					if score < globalBestScore[i] then
						globalBestScore[i] = score
						local gb = globalBest[i]
						local pStore = param._storage
						local gbStore = gb._storage
						task.desynchronize()
						for j = 1, #pStore do
							gbStore[j] = pStore[j]
						end
						task.synchronize()
					end
				end
			end
			for i, param in ipairs(params) do
				local grad = param._grad
				if grad then
					local v = velocities[i]
					local pb = personalBest[i]
					local gb = globalBest[i]

					local pStore = param._storage
					local gStore = grad._storage
					local vStore = v._storage
					local pbStore = pb._storage
					local gbStore = gb._storage
					task.desynchronize()
					for j = 1, #pStore do
						local r1 = math.random()
						local r2 = math.random()
						local g = gStore[j]

						local current = pStore[j]
						local vj = vStore[j]
						local pbj = pbStore[j]
						local gbj = gbStore[j]
						vj = inertia * vj
							+ cognitive * r1 * (pbj - current)
							+ social * r2 * (gbj - current)
							- gradScale * g

						local newVal = current + lr * vj
						if mutationRate > 0 and math.random() < mutationRate then
							newVal += (math.random() * 2 - 1) * mutationStrength
						end
						vStore[j] = vj
						pStore[j] = newVal
					end
					task.synchronize()
				end
			end
		end,
		zeroGrad = function()
			for _, param in ipairs(params) do
				local grad = param._grad
				if grad then
					local gStore = grad._storage
					task.desynchronize()
					for i = 1, #gStore do
						gStore[i] = 0
					end
					task.synchronize()
				end
			end
		end,
		getState = function()
			return nil
		end,
		setState = function(_self, _state)
			return
		end,
	}
end


--!native
--!optimize 2
-- multi-particle PSO Optimizer for QIMHNN
local Types = require("../../Types")
local Tensor = require("../../Tensor")
local Util = require("../../Util")

local assert = Util.Assert

export type SwarmPSOConfig = {
	swarmSize: number,
	inertia: number?,
	cognitive: number?,
	social: number?,
	lr: number?,
	mutationRate: number?,
	mutationStrength: number?,
	evalFn: (number) -> number,
}

local function copyTensorData(src: Types.Tensor, dst: Types.Tensor)
	assert(#src._storage == #dst._storage, "SwarmPSO: tensor size mismatch in copyTensorData")
	task.desynchronize()
	for i = 1, #src._storage do
		dst._storage[i] = src._storage[i]
	end
	task.synchronize()
end
local function swapTensorStorage(a: Types.Tensor, b: Types.Tensor)
	local tmp = a._storage
	a._storage = b._storage
	b._storage = tmp
end

return function(params: { Types.Tensor }, config: SwarmPSOConfig): Types.Optimizer
	assert(config and type(config) == "table", "SwarmPSO: config table required")
	assert(config.swarmSize and config.swarmSize > 0, "SwarmPSO: swarmSize > 0 required")
	assert(config.evalFn, "SwarmPSO: evalFn callback is required")

	local swarmSize = config.swarmSize
	local inertia = config.inertia or 0.7
	local cognitive = config.cognitive or 1.5
	local social = config.social or 1.5
	local lr = config.lr or 1.0
	local mutationRate = config.mutationRate or 0.0
	local mutationStrength = config.mutationStrength or 0.0
	local evalFn = config.evalFn

	local positions: { { Types.Tensor } } = {}
	local velocities: { { Types.Tensor } } = {}
	local pbest: { { Types.Tensor } } = {}
	local pbestScore: { number } = {}

	local gbest: { Types.Tensor } = {}
	local gbestScore = math.huge

	for p = 1, swarmSize do
		local posForParticle: { Types.Tensor } = {}
		local velForParticle: { Types.Tensor } = {}
		local pbestForParticle: { Types.Tensor } = {}
		for _, param in ipairs(params) do
			local pos = Tensor.zeros(param._shape, param._dtype, false)
			copyTensorData(param, pos)

			local vel = Tensor.zeros(param._shape, param._dtype, false)
			task.desynchronize()
			for j = 1, #vel._storage do
				vel._storage[j] = (math.random() * 2 - 1) * 0.001
			end
			task.synchronize()

			local pb = Tensor.zeros(param._shape, param._dtype, false)
			copyTensorData(param, pb)

			table.insert(posForParticle, pos)
			table.insert(velForParticle, vel)
			table.insert(pbestForParticle, pb)
		end
		positions[p] = posForParticle
		velocities[p] = velForParticle
		pbest[p] = pbestForParticle
		pbestScore[p] = math.huge
	end
	for _, param in ipairs(params) do
		local gb = Tensor.zeros(param._shape, param._dtype, false)
		copyTensorData(param, gb)
		table.insert(gbest, gb)
	end
	for i, param in ipairs(params) do
		param._storage = gbest[i]._storage
	end
	return {
		step = function()
			for p = 1, swarmSize do
				local posForParticle = positions[p]
				for i, param in ipairs(params) do
					local pos = posForParticle[i]
					swapTensorStorage(param, pos)
				end
				local loss = evalFn(p)
				assert(type(loss) == "number", "SwarmPSO: evalFn must return a number (loss)")
				for i, param in ipairs(params) do
					local pos = posForParticle[i]
					swapTensorStorage(param, pos)
				end
				if loss < pbestScore[p] then
					pbestScore[p] = loss
					local pbForParticle = pbest[p]
					for i = 1, #params do
						local pos = posForParticle[i]
						local pb = pbForParticle[i]
						copyTensorData(pos, pb)
					end
				end
				if loss < gbestScore then
					gbestScore = loss
					for i = 1, #params do
						local pos = posForParticle[i]
						local gb = gbest[i]
						copyTensorData(pos, gb)
					end
				end
			end
			for p = 1, swarmSize do
				local posForParticle = positions[p]
				local velForParticle = velocities[p]
				local pbForParticle = pbest[p]
				for i = 1, #params do
					local pos = posForParticle[i]
					local vel = velForParticle[i]
					local pb = pbForParticle[i]
					local gb = gbest[i]
					task.desynchronize()
					for j = 1, #pos._storage do
						local x = pos._storage[j]
						local vj = vel._storage[j]
						local pbestVal = pb._storage[j]
						local gbestVal = gb._storage[j]
						local r1 = math.random()
						local r2 = math.random()

						vj = inertia * vj + cognitive * r1 * (pbestVal - x) + social * r2 * (gbestVal - x)
						vj *= lr

						local newX = x + vj
						if mutationRate > 0 and math.random() < mutationRate then
							newX += (math.random() * 2 - 1) * mutationStrength
						end
						vel._storage[j] = vj
						pos._storage[j] = newX
					end
					task.synchronize()
				end
			end
			for i, param in ipairs(params) do
				param._storage = gbest[i]._storage
				if param._grad then
					task.desynchronize()
					for j = 1, #param._grad._storage do
						param._grad._storage[j] = 0
					end
					task.synchronize()
				end
			end
		end,
		zeroGrad = function()
			for _, param in ipairs(params) do
				if param._grad then
					task.desynchronize()
					for i = 1, #param._grad._storage do
						param._grad._storage[i] = 0
					end
					task.synchronize()
				end
			end
		end,
		getState = function()
			return nil
		end,
		setState = function(_self, _state)
			return
		end,
	}
end


--!native
--!optimize 2
local Types = require("../../Types")
local Tensor = require("../../Tensor")
local Linear = require("../../nn/Linear")
local LSTM = require("../../nn/LSTM")
local Softmax = require("../../nn/Softmax")
local Util = require("../../Util")
local Activations = require("../../nn/Activations")
local assert = Util.Assert

local function cosineSimilarity(u: Types.Tensor, v: Types.Tensor)
	-- u, v: {D, B}
	local dot = Tensor.zeros({ 1, u._shape[2] }, u._dtype, u._requiresGrad or v._requiresGrad)
	local B = u._shape[2]
	local D = u._shape[1]
	task.desynchronize()
	local u_store = u._storage
	local v_store = v._storage
	local out_store = dot._storage
	for b = 1, B do
		local sum_uv, sum_uu, sum_vv = 0, 0, 0
		for i = 1, D do
			local idx = (i - 1) * B + b
			local val_u = u_store[idx]
			local val_v = v_store[idx]
			sum_uv += val_u * val_v
			sum_uu += val_u * val_u
			sum_vv += val_v * val_v
		end
		local norm_u = math.sqrt(sum_uu) + 1e-8
		local norm_v = math.sqrt(sum_vv) + 1e-8
		out_store[b] = sum_uv / (norm_u * norm_v)
	end
	task.synchronize()
	return dot
end

return function(cfg: {
	inputDim: number,
	actionDim: number,
	hiddenDim: number,
	goalDim: number?,
	horizon: number?,
	gammaW: number?,
	gammaM: number?,
	lr: number?,
	optimizerFactory: (params: { number }) -> Types.Optimizer,
})
	assert(cfg.inputDim, "Feudal: inputDim required")
	assert(cfg.actionDim, "Feudal: actionDim required")
	assert(cfg.hiddenDim, "Feudal: hiddenDim required")
	local inputDim = cfg.inputDim
	local actionDim = cfg.actionDim
	local hiddenDim = cfg.hiddenDim
	local goalDim = cfg.goalDim or hiddenDim
	local horizon = cfg.horizon or 10
	local _gammaW = cfg.gammaW or 0.99
	local _gammaM = cfg.gammaM or 0.99
	local _lr = cfg.lr or 1e-3
	local optimizerFactory = cfg.optimizerFactory
	assert(optimizerFactory, "Feudal: optimizerFactory required")
	local perceptor = Linear(inputDim, hiddenDim)
	local mgrLinear1 = Linear(hiddenDim, hiddenDim)
	local mgrLSTM = LSTM(hiddenDim, hiddenDim)
	local mgrGoalHead = Linear(hiddenDim, goalDim)
	local mgrValueHead = Linear(hiddenDim, 1)
	local wkrLinear1 = Linear(hiddenDim + goalDim, hiddenDim)
	local wkrLSTM = LSTM(hiddenDim, hiddenDim)
	local wkrActionHead = Linear(hiddenDim, actionDim)
	local wkrValueHead = Linear(hiddenDim, 1)
	local params = {}
	local function collectParams(mod)
		for _, p in ipairs(mod:parameters()) do
			table.insert(params, p)
		end
	end
	collectParams(perceptor)
	collectParams(mgrLinear1)
	collectParams(mgrLSTM)
	collectParams(mgrGoalHead)
	collectParams(mgrValueHead)
	collectParams(wkrLinear1)
	collectParams(wkrLSTM)
	collectParams(wkrActionHead)
	collectParams(wkrValueHead)
	local optimizer = optimizerFactory(params)
	local mgrH, mgrC
	local wkrH, wkrC
	local currentGoal
	local lastMgrUpdateStep = 0
	local stepCount = 0
	local startState
	local workerBuffer = {} -- {s, g, a, r_ext, r_int, ns, done}
	local managerBuffer = {} -- {s, g, r_sum, ns, done}
	local Agent = {}
	function Agent:reset(batchSize: number)
		batchSize = batchSize or 1
		mgrH = Tensor.zeros({ hiddenDim, batchSize }, "f64", true)
		mgrC = Tensor.zeros({ hiddenDim, batchSize }, "f64", true)
		wkrH = Tensor.zeros({ hiddenDim, batchSize }, "f64", true)
		wkrC = Tensor.zeros({ hiddenDim, batchSize }, "f64", true)
		currentGoal = Tensor.zeros({ goalDim, batchSize }, "f64", true)
		startState = nil
		stepCount = 0
		lastMgrUpdateStep = 0
		table.clear(workerBuffer)
		table.clear(managerBuffer)
	end
	function Agent:act(state: Types.Tensor)
		-- state: {inputDim, B}
		local B = state._shape[2]
		local z_raw = perceptor:forward(state)
		local z = Activations.ReLU(z_raw)
		if stepCount - lastMgrUpdateStep >= horizon or stepCount == 0 then
			local mgr_z = mgrLinear1:forward(z)
			mgr_z = Activations.Tanh(mgr_z)
			local h_new, c_new = mgrLSTM:forward(mgr_z, mgrH, mgrC)
			mgrH, mgrC = h_new, c_new
			local g = mgrGoalHead:forward(h_new)
			local g_norm = Tensor.zeros(g._shape, g._dtype, true)
			task.desynchronize()
			local g_store = g._storage
			local gn_store = g_norm._storage
			for b = 1, B do
				local sum_sq = 0
				for i = 1, goalDim do
					local val = g_store[(i - 1) * B + b]
					sum_sq += val * val
				end
				local scale = 1 / (math.sqrt(sum_sq) + 1e-8)
				for i = 1, goalDim do
					local idx = (i - 1) * B + b
					gn_store[idx] = g_store[idx] * scale
				end
			end
			task.synchronize()
			currentGoal = g_norm
			lastMgrUpdateStep = stepCount
			startState = state:detach()
		end
		-- z: {hiddenDim, B}, goal: {goalDim, B}
		local combinedInput = Tensor.zeros({ hiddenDim + goalDim, B }, state._dtype, true)
		task.desynchronize()
		local z_store = z._storage
		local g_store = currentGoal._storage
		local c_store = combinedInput._storage
		for b = 1, B do
			for i = 1, hiddenDim do
				c_store[(i - 1) * B + b] = z_store[(i - 1) * B + b]
			end
			for i = 1, goalDim do
				c_store[(hiddenDim + i - 1) * B + b] = g_store[(i - 1) * B + b]
			end
		end
		task.synchronize()
		local wz = wkrLinear1:forward(combinedInput)
		local wh_new, wc_new = wkrLSTM:forward(wz, wkrH, wkrC)
		wkrH, wkrC = wh_new, wc_new
		local logits = wkrActionHead:forward(wh_new)
		local P = Softmax.forward(logits)
		local action = 1
		task.desynchronize()
		local r = math.random()
		local cum = 0.0
		for i = 1, actionDim do
			cum += P._storage[i]
			if r <= cum then
				action = i
				break
			end
		end
		task.synchronize()
		stepCount += 1
		return action, currentGoal
	end

	function Agent:observe(transition: {
		state: Types.Tensor,
		action: number,
		reward: number,
		nextState: Types.Tensor,
		done: boolean,
	})
		local s = transition.state
		local a = transition.action
		local r = transition.reward
		local ns = transition.nextState
		local d = transition.done
		local z = Activations.ReLU(perceptor:forward(s))
		local z_next = Activations.ReLU(perceptor:forward(ns))
		local dz = Tensor.zeros(z._shape, z._dtype, false)
		task.desynchronize()
		local zs, zns, dzs = z._storage, z_next._storage, dz._storage
		for i = 1, #zs do
			dzs[i] = zns[i] - zs[i]
		end
		task.synchronize()
		local cosim = cosineSimilarity(dz, currentGoal)
		local r_int = cosim._storage[1]
		table.insert(workerBuffer, {
			s = s,
			g = currentGoal,
			a = a,
			r_ext = r,
			r_int = r_int,
			ns = ns,
			d = d,
		})
		if d or (stepCount - lastMgrUpdateStep == 0) then
			local r_mgr, workerBuffer_len = 0, #workerBuffer
			for i = workerBuffer_len - (stepCount - lastMgrUpdateStep), workerBuffer_len do
				if workerBuffer[i] then
					r_mgr += workerBuffer[i].r_ext
				end
			end
			table.insert(managerBuffer, {
				s = startState,
				g = currentGoal,
				r = r_mgr,
				ns = ns,
				d = d,
			})
		end
	end
	function Agent:trainStep()
		optimizer:zeroGrad()
		local _loss = Tensor.zeros({ 1 }, "f64", true)
		table.clear(workerBuffer)
		table.clear(managerBuffer)
		return { loss = 0 }
	end
	return Agent
end


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

local function step(self, s: Types.Tensor): number
	self.opt:zeroGrad()
	local g = self.target:forward(s) :: Types.Tensor -- {E,1}
	local f = self.predictor:forward(s) :: Types.Tensor -- {E,1}
	local loss = 0
	local grad = Tensor.zeros(f._shape)
	task.desynchronize()
	for i = 1, #f._storage do
		local d = f._storage[i] - g._storage[i]
		loss += d * d
		grad._storage[i] = 2 * d
	end
	task.synchronize()
	Tape.backwardFrom(f, grad)
	self.opt:step()
	local r = loss
	if self.norm then
		self.norm:update(r)
		r = self.norm:normalize(r)
	end
	return r
end

local function new(
	target: Types.Module<Types.Tensor, Types.Tensor>,
	predictor: Types.Module<Types.Tensor, Types.Tensor>,
	optimizer: Types.Optimizer,
	norm: {
		update: (self: { update: (x: number) -> (), normalize: (x: number) -> number }, x: number) -> number,
		normalize: (self: { update: (x: number) -> number, normalize: (x: number) -> number }, x: number) -> number,
	}?
)
	return {
		target = target,
		predictor = predictor,
		opt = optimizer,
		norm = norm,
		step = step,
	}
end

return { new = new }


--[[
Gradien (by @Eternity_Devs)
A strictly-typed, optimized parallel-first ML & DL library for Roblox.
Version 1.4.0-rc5
]]
--!strict
local Gradien = {}

Gradien.Autograd = require("@self/autograd/Tape")
Gradien.Init = require("@self/Initializer")
Gradien.Metrics = require("@self/Metrics")
Gradien.GradClip = require("@self/GradClip")
Gradien.Trainer = require("@self/Trainer")
Gradien.State = require("@self/State")
Gradien.Types = require("@self/Types")
Gradien.Util = require("@self/Util")
Gradien.Tensor = require("@self/Tensor")
Gradien.Debug = require("@self/Debug")
Gradien.Tokenizers = require("@self/tokenizers")

Gradien.Experimental = {
	Models = {
		QIMHNN = require("@self/experimental/models/QIMHNN"),
	},
	NN = {
		MambaBlock = require("@self/experimental/nn/Mamba"),
		KAN = require("@self/experimental/nn/KAN"),
	},
	Optim = {
		Metaheuristic = require("@self/experimental/optim/Metaheuristic"),
		SwarmPSO = require("@self/experimental/optim/SwarmPSO"),
	},
	RL = {
		Feudal = require("@self/experimental/rl/Feudal"),
	},
}

Gradien.Models = {
	MLP = require("@self/models/MLP"),
	ResMLP = require("@self/models/ResMLP"),
	ConvNet = require("@self/models/ConvNet"),
	TransformerEncoder = require("@self/models/TransformerEncoder"),
	SequenceClassifier = require("@self/models/SequenceClassifier"),
	AutoEncoder = require("@self/models/AutoEncoder"),
	Utils = require("@self/models/Utils"),
}
Gradien.Ops = {
	Math = require("@self/ops/Math"),
	BLAS = require("@self/ops/BLAS"),
	Reduce = require("@self/ops/Reduce"),
	Conv2d = require("@self/ops/Conv2d"),
	Pool = require("@self/ops/Pool"),
	Softmax = require("@self/ops/Softmax"),
	AvgPool2d = require("@self/ops/AvgPool2d"),
	Scan = require("@self/ops/Scan"),
	ConvTranspose2d = require("@self/ops/ConvTranspose2d"),
}
Gradien.NN = {
	Linear = require("@self/nn/Linear"),
	NoisyLinear = require("@self/nn/NoisyLinear"),
	Activations = require("@self/nn/Activations"),
	LayerNorm = require("@self/nn/LayerNorm"),
	BatchNorm1d = require("@self/nn/BatchNorm1d"),
	BatchNorm2d = require("@self/nn/BatchNorm2d"),
	Embedding = require("@self/nn/Embedding"),
	Attention = require("@self/nn/Attention"),
	RMSNorm = require("@self/nn/RMSNorm"),
	RNN = require("@self/nn/RNN"),
	GRU = require("@self/nn/GRU"),
	LSTM = require("@self/nn/LSTM"),
	Sequential = require("@self/nn/Sequential"),
	Dropout = require("@self/nn/Dropout"),
	Fused = require("@self/nn/Fused"),
	Softmax = require("@self/nn/Softmax"),
	Losses = require("@self/nn/Losses"),
	Functional = require("@self/nn/Functional"),
	GatedMLP = require("@self/nn/GatedMLP"),
	DropPath = require("@self/nn/DropPath"),
	Residual = require("@self/nn/Residual"),
	SwiGLU = require("@self/nn/SwiGLU"),
	Conv2d = require("@self/nn/Conv2d"),
	MaxPool2d = require("@self/nn/MaxPool2d"),
	AvgPool2d = require("@self/nn/AvgPool2d"),
	ConvTranspose2d = require("@self/nn/ConvTranspose2d"),
	GroupNorm = require("@self/nn/GroupNorm"),
	Flatten = require("@self/nn/Flatten"),
}
Gradien.Optim = {
	SGD = require("@self/optim/SGD"),
	Adam = require("@self/optim/Adam"),
	AdamW = require("@self/optim/AdamW"),
	EMA = require("@self/optim/EMA"),
	RMSProp = require("@self/optim/RMSProp"),
	Adagrad = require("@self/optim/Adagrad"),
	Adafactor = require("@self/optim/Adafactor"),
	Accumulated = require("@self/optim/Accumulated"),
	Schedulers = require("@self/optim/Schedulers"),
	Lion = require("@self/optim/Lion"),
	Lookahead = require("@self/optim/Lookahead"),
	Sophia = require("@self/optim/Sophia"),
	Muon = require("@self/optim/Muon"),
}
Gradien.RL = {
	Replay = require("@self/rl/Replay"),
	UniformReplay = require("@self/rl/UniformReplay"),
	PrioritizedReplay = require("@self/rl/PrioritizedReplay"),
	DQL = require("@self/rl/DQL"),
	DoubleDQN = require("@self/rl/DoubleDQN"),
	C51DQN = require("@self/rl/C51DQN"),
	BDQ = require("@self/rl/BDQ"),
	Dueling = require("@self/rl/Dueling"),
	NStep = require("@self/rl/NStep"),
	LambdaBuffer = require("@self/rl/LambdaBuffer"),
	A2C = require("@self/rl/A2C"),
	PPO = require("@self/rl/PPO"),
	MultiAgent = require("@self/rl/MultiAgent"),
}
Gradien.Classic = {
	SoftmaxRegression = require("@self/classic/SoftmaxRegression"),
	LinearRegression = require("@self/classic/LinearRegression"),
	Ridge = require("@self/classic/Ridge"),
	LogisticRegression = require("@self/classic/LogisticRegression"),
	KMeans = require("@self/classic/KMeans"),
	KNN = require("@self/classic/KNN"),
	GaussianNB = require("@self/classic/GaussianNB"),
	MultinomialNB = require("@self/classic/MultinomialNB"),
}
Gradien.Preprocess = {
	StandardScaler = require("@self/preprocess/StandardScaler"),
	MinMaxScaler = require("@self/preprocess/MinMaxScaler"),
	OneHot = require("@self/preprocess/OneHot"),
	PCA = require("@self/preprocess/PCA"),
	SinusoidalPE = require("@self/preprocess/SinusoidalPE"),
}
Gradien.Data = {
	TensorDataset = require("@self/data/TensorDataset"),
	DataLoader = require("@self/data/DataLoader"),
	Split = require("@self/data/Split"),
	KFold = require("@self/data/KFold"),
}
Gradien.Extra = {
	RND = require("@self/extra/RND"),
}
Gradien.Tools = {
	Gradcheck = require("@self/tools/Gradcheck"),
}

return Gradien


--!native
--!optimize 2
local Types = require("../Types")
local Utils = require("./Utils")
local Util = require("../Util")
local Linear = require("../nn/Linear")
local LayerNorm = require("../nn/LayerNorm")
local Dropout = require("../nn/Dropout")
local Sequential = require("../nn/Sequential")

local assert = Util.Assert

export type AutoEncoderConfig = {
	inputDim: number,
	latentDim: number,
	hiddenDims: { number }?,
	activation: Utils.ActivationName?,
	dropout: number?,
	layerNorm: boolean?,
	finalActivation: Utils.ActivationName?,
	weightInit: ((number, number) -> number)?,
}

local AutoEncoder = {}
AutoEncoder.__index = AutoEncoder

local function buildStack(
	inDim: number,
	hiddenDims: { number },
	outDim: number,
	activation: Utils.ActivationName?,
	dropProb: number,
	useLayerNorm: boolean?,
	init: ((number, number) -> number)?
): (Types.Module<Types.Tensor, Types.Tensor>, { any })
	local layers = {}
	local trainables: { any } = {}
	local function insert(m)
		table.insert(layers, m)
		table.insert(trainables, m)
	end
	local current = inDim
	for _, hidden in ipairs(hiddenDims) do
		insert(Linear(current, hidden, init))
		if useLayerNorm then
			insert(LayerNorm(hidden))
		end
		insert(Utils.makeActivationModule(activation))
		if dropProb > 0 then
			insert(Dropout(dropProb))
		end
		current = hidden
	end
	insert(Linear(current, outDim, init))
	return Sequential(layers), trainables
end

function AutoEncoder.new(cfg: AutoEncoderConfig): Types.Module<Types.Tensor, Types.Tensor>
	assert(cfg.inputDim and cfg.latentDim, "AutoEncoder: inputDim and latentDim required")
	local hidden = if cfg.hiddenDims and #cfg.hiddenDims > 0
		then table.clone(cfg.hiddenDims)
		else {
			math.max(cfg.inputDim, cfg.latentDim * 2),
		}
	local revHidden = table.clone(hidden)
	for i = 1, math.floor(#revHidden / 2) do
		local j = #revHidden - i + 1
		revHidden[i], revHidden[j] = revHidden[j], revHidden[i]
	end
	local drop = math.clamp(cfg.dropout or 0, 0, 0.95)

	local encoder, encTrain =
		buildStack(cfg.inputDim, hidden, cfg.latentDim, cfg.activation, drop, cfg.layerNorm, cfg.weightInit)
	local decoder, decTrain =
		buildStack(cfg.latentDim, revHidden, cfg.inputDim, cfg.activation, drop, cfg.layerNorm, cfg.weightInit)
	local finalAct = if cfg.finalActivation then Utils.makeActivationModule(cfg.finalActivation) else nil

	local trainables: { any } = {}
	for _, m in ipairs(encTrain) do
		table.insert(trainables, m)
	end
	for _, m in ipairs(decTrain) do
		table.insert(trainables, m)
	end
	if finalAct then
		table.insert(trainables, finalAct)
	end

	local self = setmetatable({
		_encoder = encoder,
		_decoder = decoder,
		_finalAct = finalAct,
		_trainables = trainables,
	}, AutoEncoder)

	return self
end

function AutoEncoder:encode(x: Types.Tensor): Types.Tensor
	return self._encoder:forward(x)
end

function AutoEncoder:decode(z: Types.Tensor): Types.Tensor
	local out = self._decoder:forward(z)
	if self._finalAct then
		out = self._finalAct:forward(out)
	end
	return out
end

function AutoEncoder:forward(x: Types.Tensor): Types.Tensor
	local z = self:encode(x)
	return self:decode(z)
end

function AutoEncoder:parameters(): { Types.Tensor }
	local params = {}
	Utils.collectParams(params, self._encoder)
	Utils.collectParams(params, self._decoder)
	return params
end

function AutoEncoder:train(mode: boolean?)
	Utils.trainAll(self._trainables, mode)
end

return setmetatable({
	new = function(cfg: AutoEncoderConfig): Types.Module<Types.Tensor, Types.Tensor>
		return AutoEncoder.new(cfg)
	end,
}, {
	__call = function(_, cfg: AutoEncoderConfig): Types.Module<Types.Tensor, Types.Tensor>
		return AutoEncoder.new(cfg)
	end,
})


--!native
--!optimize 2
local Types = require("../Types")
local Utils = require("./Utils")
local Util = require("../Util")
local Conv2d = require("../nn/Conv2d")
local BatchNorm2d = require("../nn/BatchNorm2d")
local MaxPool2d = require("../nn/MaxPool2d")
local AvgPool2d = require("../nn/AvgPool2d")
local Dropout = require("../nn/Dropout")
local Flatten = require("../nn/Flatten")
local Linear = require("../nn/Linear")
local Sequential = require("../nn/Sequential")

local assert = Util.Assert

export type PoolConfig = {
	kind: "max" | "avg",
	size: number?,
	stride: number?,
}

export type ConvBlockConfig = {
	outChannels: number,
	kernelSize: number?,
	activation: Utils.ActivationName?,
	useBatchNorm: boolean?,
	pool: PoolConfig?,
	dropout: number?,
}

export type ConvNetConfig = {
	inChannels: number,
	inputHeight: number,
	inputWidth: number,
	numClasses: number,
	blocks: { ConvBlockConfig }?,
	headDims: { number }?,
	activation: Utils.ActivationName?,
	dropout: number?,
	headDropout: number?,
	finalActivation: Utils.ActivationName?,
	weightInit: ((number, number) -> number)?,
}

local function defaultBlocks(): { ConvBlockConfig }
	return {
		{ outChannels = 32, kernelSize = 3, pool = { kind = "max", size = 2 } },
		{ outChannels = 64, kernelSize = 3, pool = { kind = "max", size = 2 } },
		{ outChannels = 128, kernelSize = 3, pool = { kind = "max", size = 2 } },
	}
end

local function applyConvSize(size: number, kernel: number): number
	local nextSize = size - kernel + 1
	assert(nextSize > 0, "ConvNet: spatial dimension collapsed (conv)")
	return nextSize
end

local function applyPoolSize(size: number, cfg: PoolConfig): number
	local k = cfg.size or 2
	local stride = cfg.stride or k
	local nextSize = math.floor((size - k) / stride) + 1
	assert(nextSize > 0, "ConvNet: spatial dimension collapsed (pool)")
	return nextSize
end

return function(config: ConvNetConfig): Types.Module<Types.Tensor, Types.Tensor>
	assert(config.inChannels and config.inputHeight and config.inputWidth, "ConvNet: input spec required")
	assert(config.numClasses, "ConvNet: numClasses required")
	local blocks = if config.blocks and #config.blocks > 0 then config.blocks else defaultBlocks()

	local H, W = config.inputHeight, config.inputWidth
	local C = config.inChannels
	for _, block: ConvBlockConfig in ipairs(blocks) do
		local kernel = block.kernelSize or 3
		H = applyConvSize(H, kernel)
		W = applyConvSize(W, kernel)
		if block.pool then
			H = applyPoolSize(H, block.pool)
			W = applyPoolSize(W, block.pool)
		end
		C = block.outChannels
	end
	local flattenDim = C * H * W
	assert(flattenDim > 0, "ConvNet: invalid flatten dimension")

	local convLayers = {}
	local trainables: { any } = {}
	local inChannels = config.inChannels
	for _, block: ConvBlockConfig in ipairs(blocks) do
		local kernel = block.kernelSize or 3
		local conv = Conv2d(inChannels, block.outChannels, kernel, kernel)
		table.insert(convLayers, conv)
		table.insert(trainables, conv)

		if block.useBatchNorm ~= false then
			local bn = BatchNorm2d(block.outChannels)
			table.insert(convLayers, bn)
			table.insert(trainables, bn)
		end

		local blockActivation: Utils.ActivationName?
		if block.activation ~= nil then
			blockActivation = block.activation
		elseif config.activation ~= nil then
			blockActivation = config.activation
		else
			blockActivation = nil
		end
		table.insert(convLayers, Utils.makeActivationModule(blockActivation))

		if block.pool then
			local poolCfg = block.pool
			local size = poolCfg.size or 2
			local stride = poolCfg.stride or size
			local poolModule = if poolCfg.kind == "avg"
				then AvgPool2d(size, size, stride)
				else MaxPool2d(size, size, stride)
			table.insert(convLayers, poolModule)
			table.insert(trainables, poolModule)
		end

		local dropProb = block.dropout or config.dropout or 0
		if dropProb > 0 then
			local drop = Dropout(math.clamp(dropProb, 0, 0.95))
			table.insert(convLayers, drop)
			table.insert(trainables, drop)
		end

		inChannels = block.outChannels
	end

	local featureExtractor = Sequential(convLayers)
	local flatten = Flatten()
	table.insert(trainables, flatten)

	local headLayers = {}
	local headDims = config.headDims or { 256 }
	local headDrop = math.clamp(config.headDropout or config.dropout or 0, 0, 0.95)
	local lastDim = flattenDim
	for _, hidden in ipairs(headDims) do
		assert(hidden > 0, "ConvNet: head hidden dims must be positive")
		local lin = Linear(lastDim, hidden, config.weightInit)
		table.insert(headLayers, lin)
		table.insert(trainables, lin)
		table.insert(headLayers, Utils.makeActivationModule(config.activation))
		if headDrop > 0 then
			local drop = Dropout(headDrop)
			table.insert(headLayers, drop)
			table.insert(trainables, drop)
		end
		lastDim = hidden
	end
	local classifier = Linear(lastDim, config.numClasses, config.weightInit)
	table.insert(headLayers, classifier)
	table.insert(trainables, classifier)
	if config.finalActivation then
		table.insert(headLayers, Utils.makeActivationModule(config.finalActivation))
	end

	local head = Sequential(headLayers)

	return {
		forward = function(_, x: Types.Tensor): Types.Tensor
			local h = featureExtractor:forward(x)
			h = flatten:forward(h)
			return head:forward(h)
		end,
		parameters = function(_): { Types.Tensor }
			local params = {}
			Utils.collectParams(params, featureExtractor)
			Utils.collectParams(params, head)
			return params
		end,
		train = function(_, mode: boolean?)
			Utils.trainAll(trainables, mode)
		end,
	}
end


--!native
--!optimize 2
local Types = require("../Types")
local Utils = require("./Utils")
local Util = require("../Util")
local Linear = require("../nn/Linear")
local LayerNorm = require("../nn/LayerNorm")
local Dropout = require("../nn/Dropout")
local Sequential = require("../nn/Sequential")

local assert = Util.Assert

export type MLPConfig = {
	inputDim: number,
	outputDim: number,
	hiddenDims: { number }?,
	activation: Utils.ActivationName?,
	finalActivation: Utils.ActivationName?,
	dropout: number?,
	layerNorm: boolean?,
	weightInit: ((number, number) -> number)?,
}

return function(config: MLPConfig): Types.Module<Types.Tensor, Types.Tensor>
	assert(config.inputDim, "MLP: inputDim is required")
	assert(config.outputDim, "MLP: outputDim is required")

	local hidden = if config.hiddenDims and #config.hiddenDims > 0
		then table.clone(config.hiddenDims)
		else {
			math.max(config.inputDim, config.outputDim),
		}
	local dropProb = math.clamp(config.dropout or 0, 0, 0.95)

	local layers = {}
	local trainables: { any } = {}
	local inDim = config.inputDim

	for _, hiddenDim in ipairs(hidden) do
		assert(hiddenDim > 0, "MLP: hidden dimensions must be positive")
		local linear = Linear(inDim, hiddenDim, config.weightInit)
		table.insert(layers, linear)
		table.insert(trainables, linear)

		if config.layerNorm then
			local norm = LayerNorm(hiddenDim)
			table.insert(layers, norm)
			table.insert(trainables, norm)
		end

		table.insert(layers, Utils.makeActivationModule(config.activation))

		if dropProb > 0 then
			local drop = Dropout(dropProb)
			table.insert(layers, drop)
			table.insert(trainables, drop)
		end

		inDim = hiddenDim
	end

	local head = Linear(inDim, config.outputDim, config.weightInit)
	table.insert(layers, head)
	table.insert(trainables, head)

	if config.finalActivation then
		table.insert(layers, Utils.makeActivationModule(config.finalActivation))
	end

	local seq = Sequential(layers)

	return {
		forward = function(_, x: Types.Tensor): Types.Tensor
			return seq:forward(x)
		end,
		parameters = function(_): { Types.Tensor }
			return seq:parameters()
		end,
		train = function(_, mode: boolean?)
			Utils.trainAll(trainables, mode)
		end,
	}
end


--!native
--!optimize 2
local Types = require("../Types")
local Utils = require("./Utils")
local Util = require("../Util")
local Linear = require("../nn/Linear")
local LayerNorm = require("../nn/LayerNorm")
local Dropout = require("../nn/Dropout")
local Sequential = require("../nn/Sequential")

local assert = Util.Assert

export type ResMLPConfig = {
	inputDim: number,
	outputDim: number,
	embedDim: number?,
	hiddenDim: number?,
	numBlocks: number?,
	activation: Utils.ActivationName?,
	dropout: number?,
	layerNorm: boolean?,
	finalLayerNorm: boolean?,
	weightInit: ((number, number) -> number)?,
}

type Block = {
	norm: Types.Module<Types.Tensor, Types.Tensor>?,
	ff: Types.Module<Types.Tensor, Types.Tensor>,
	trainables: { any },
}

local function buildBlock(
	embedDim: number,
	hiddenDim: number,
	activation: Utils.ActivationName?,
	dropProb: number,
	useLayerNorm: boolean?,
	init: ((number, number) -> number)?
): Block
	local norm = if useLayerNorm then LayerNorm(embedDim) else nil
	local layers = {}
	local trainables: { any } = {}

	local function insertModule(m)
		table.insert(layers, m)
		table.insert(trainables, m)
	end

	insertModule(Linear(embedDim, hiddenDim, init))
	insertModule(Utils.makeActivationModule(activation))
	if dropProb > 0 then
		insertModule(Dropout(dropProb))
	end
	insertModule(Linear(hiddenDim, embedDim, init))
	if dropProb > 0 then
		insertModule(Dropout(dropProb))
	end

	local ff = Sequential(layers)
	if norm then
		table.insert(trainables, norm)
	end

	return {
		norm = norm,
		ff = ff,
		trainables = trainables,
	}
end

local function buildModel(config: ResMLPConfig)
	assert(config.inputDim and config.outputDim, "ResMLP: inputDim and outputDim required")
	local embedDim = config.embedDim or config.inputDim
	local hiddenDim = config.hiddenDim or (embedDim * 4)
	local blocks = math.max(1, config.numBlocks or 4)
	local dropout = math.clamp(config.dropout or 0, 0, 0.95)

	local inputProj = if config.inputDim ~= embedDim then Linear(config.inputDim, embedDim, config.weightInit) else nil
	local outputHead = Linear(embedDim, config.outputDim, config.weightInit)
	local finalNorm = if config.finalLayerNorm then LayerNorm(embedDim) else nil

	local stack = table.create(blocks)
	for i = 1, blocks do
		stack[i] = buildBlock(embedDim, hiddenDim, config.activation, dropout, config.layerNorm, config.weightInit)
	end

	local trainables: { any } = {}
	if inputProj then
		table.insert(trainables, inputProj)
	end
	for _, block in ipairs(stack) do
		for _, m in ipairs(block.trainables) do
			table.insert(trainables, m)
		end
	end
	if finalNorm then
		table.insert(trainables, finalNorm)
	end
	table.insert(trainables, outputHead)

	local function forward(_, x: Types.Tensor): Types.Tensor
		local h = x
		if inputProj then
			h = inputProj:forward(h)
		end
		for _, block in ipairs(stack) do
			local y = h
			if block.norm then
				y = block.norm:forward(y)
			end
			y = block.ff:forward(y)
			h = Utils.addTensors(h, y)
		end
		if finalNorm then
			h = finalNorm:forward(h)
		end
		return outputHead:forward(h)
	end

	local function parameters(_): { Types.Tensor }
		local params = {}
		if inputProj then
			Utils.collectParams(params, inputProj)
		end
		for _, block in ipairs(stack) do
			if block.norm then
				Utils.collectParams(params, block.norm)
			end
			Utils.collectParams(params, block.ff)
		end
		if finalNorm then
			Utils.collectParams(params, finalNorm)
		end
		Utils.collectParams(params, outputHead)
		return params
	end

	local function train(_, mode: boolean?)
		Utils.trainAll(trainables, mode)
	end

	local module = {
		forward = forward,
		parameters = parameters,
		train = train,
	}
	return module
end

return setmetatable({
	new = function(config: ResMLPConfig)
		return buildModel(config)
	end,
}, {
	__call = function(_, config: ResMLPConfig)
		return buildModel(config)
	end,
})