graph LR
    CPU[CPU<br>Accelerate / LAPACK] <-->|Zero-Copy| RAM[(Apple Silicon<br>Unified Memory)]
    GPU[GPU<br>Metal Performance Shaders] <-->|Zero-Copy| RAM

dependencies: [
  .package(url: "https://github.com/acemoglu/SwiftMetalNumerics.git", branch: "main")
]

import SwiftMetalNumerics

if SwiftMetalNumericsCapabilities.isMetalGPUAvailable {
    // GPU paths are preferred; library still falls back where implemented
}

// From explicit elements (row-major [Float])
let A = try NumericMatrix(rows: 2, columns: 3, elements: [1, 2, 3, 4, 5, 6])

// Zero-filled
let Z = try NumericMatrix(rows: 4, columns: 4)

// Identity and diagonal helpers
let I = try NumericMatrix.identity(size: 5)
let D = try NumericMatrix.diagonal(values: [1, 2, 3])

// Read-only / mutable access to scalars without extra copies of the backing store
let sum = A.withFloatBuffer { buf in
    buf.reduce(0, +)
}

var B = A
try B.withMutableFloatBuffer { ptr in
    if let base = ptr.baseAddress { base[0] = 99 }
}

let M = try NumericMatrix(rows: 2, columns: 2, elements: [1, 2, 3, 4])
let N = try NumericMatrix(rows: 2, columns: 2, elements: [5, 6, 7, 8])

let product = try await M.multiplied(by: N)
// same as: try await M * N

let Mt = try await M.transposed()
let H = try await M.hadamard(N)
// or: try await M .* N

// Scalar scaling: matrix * scalar / scalar * matrix (not matrix×matrix)
let scaled = try await M * 0.5
let scaled2 = try await 2.0 * M

// CPU-only GEMM (e.g. benchmarks vs GPU):
let cpu = try await M.multipliedUsingAccelerateCPU(by: N)

let S = try await M.adding(N)
let D = try await M.subtracting(N)
// Operators: try await M + N, try await M - N

let E = try await M.scaled(by: 0.25)
let rsqrt = try await M.sqrt()
let e = try await M.exp()
let l = try await M.log()
let s = try await M.sin()

let mx = try await M.max()
let mn = try await M.min()
let mu = try await M.mean()
let t = try await I.trace()   // square matrices only

// Solve A x = b for square A
let A = try NumericMatrix(rows: 3, columns: 3, elements: [
    2, 0, 0,
    0, 2, 0,
    0, 0, 2
])
let x = try await A.solve(b: [1, 2, 3])

let inv = try await A.inversed()
let det = try await A.determinant()
let (U, S, Vt) = try await A.svd()
let L = try await A.cholesky()   // symmetric positive-definite A only

let height = 32
let width = 32
let image = try NumericMatrix(
    rows: height,
    columns: width,
    elements: Array(repeating: 0.1, count: height * width)
)

let kernel = try NumericMatrix(rows: 3, columns: 3, elements: [
    0, 0.25, 0,
    0.25, 0, 0.25,
    0, 0.25, 0
])

let out = try await image.convolved(with: kernel)

// Real input → complex spectrum
let n = 512
let real = (0..<n).map { Float($0) }
var signal = Signal1D(realSamples: real)
let spectrum = try await signal.forwardFFT()

// Inverse spectrum → time-domain signal
let reconstructed = try await spectrum.inverseFFT()

// General complex input
let complexSamples: [ComplexFloat] = (0..<n).map {
    ComplexFloat(real: Float($0), imaginary: 0)
}
signal = Signal1D(samples: complexSamples)
let spec2 = try await signal.forwardFFT()

// CPU-only FFT (benchmarks / forcing vDSP):
let cpuSpectrum = try await Signal1D.forwardFFTReal1DUsingAccelerate(real)

let sr = 48_000
let seconds: Float = 1.0
let count = Int(Float(sr) * seconds)
let samples = (0..<count).map { i in
    // 440 Hz sine
    sin(2 * .pi * 440 * Float(i) / Float(sr))
}

// Magnitude spectrogram: shape [numFrames, frameSize]
let mag = try await STFT.compute(
    signal: samples,
    frameSize: 1024,
    hopSize: 256,
    window: .hann,
    outputKind: .magnitude
)

// Complex interleaved: shape [numFrames, frameSize, 2] where [...,0]=real [...,1]=imag
let complex = try await STFT.compute(
    signal: samples,
    frameSize: 1024,
    hopSize: 256,
    window: .hann,
    outputKind: .complexInterleaved
)

// Rank-1
let v = try NumericTensor(contiguousRowMajorShape: [8], elements: (0..<8).map { Float($0) })

// Rank-4 NCHW: [batch, channels, height, width]
let n = 4, c = 3, h = 32, w = 32
let x = try NumericTensor(
    contiguousRowMajorShape: [n, c, h, w],
    elements: Array(repeating: 0.01, count: n * c * h * w)
)

// Read back scalars (requires contiguous row-major)
let first10 = try x.makeContiguousRowMajor().withFloatBuffer { buf in
    Array(buf.prefix(10))
}
_ = first10

let H = 64, W = 64
let img = try NumericTensor(
    contiguousRowMajorShape: [H, W],
    elements: Array(repeating: 1.0, count: H * W)
)

// Forward: [H, W, 2] (real/imag)
let spec = try await img.forwardFFT2D()

// Inverse: returns real rank-2 tensor [H, W] (imag discarded)
let back = try await spec.inverseFFT2D()
_ = back

import Metal

let device = MTLCreateSystemDefaultDevice()!
guard MPSGraphFloat16Support.isSupported(on: device) else {
    fatalError("FP16 ops not supported on this GPU")
}

// Matrix FP32 → FP16 (validated against default device) → multiply → back to FP32
let X = try NumericMatrix(rows: 2, columns: 2, elements: [1, 2, 3, 4])
let X16 = try await X.convertedToFloat16()
let Y16 = try NumericMatrix16(rows: 2, columns: 2, elements: [1, 0, 0, 1])
let P16 = try await X16.multiplied(by: Y16)
let back32 = try P16.toFloat32Matrix()
_ = back32

// Random Xavier init (FP32)
let fc = try LinearLayer(inFeatures: 10, outFeatures: 4, bias: true, dataType: .float32)
let batch = try NumericMatrix(rows: 8, columns: 10, elements: Array(repeating: 0.1, count: 8 * 10))
let logits = try await fc.forward(batch)

// Fixed weights [inFeatures × outFeatures], bias [1 × outFeatures]
let W = try NumericMatrix(rows: 10, columns: 4, elements: Array(repeating: 0.01, count: 40))
let b = try NumericMatrix(rows: 1, columns: 4, elements: [0, 0, 0, 1])
let fc2 = try LinearLayer(weights: W, bias: b, inFeatures: 10, outFeatures: 4)
let logits2 = try await fc2.forward(batch)
_ = (logits, logits2)

// FP32 NCHW convolution
let conv = try Conv2DLayer(
    inChannels: 3,
    outChannels: 16,
    kernelHeight: 3,
    kernelWidth: 5,
    stride: (1, 2),
    padding: (1, 2),
    bias: true,
    dataType: .float32
)

let x = try NumericTensor(
    contiguousRowMajorShape: [4, 3, 32, 32],
    elements: Array(repeating: 0.02, count: 4 * 3 * 32 * 32)
)

let y = try await conv.forward(x)
_ = y

let maxPool = try MaxPooling2DLayer(kernelSize: 2, stride: 2, padding: 0)
let avgPool = try AveragePooling2DLayer(kernelSize: (3, 3), stride: (2, 2), padding: (1, 1))

let y1 = try await maxPool.forward(x)
let y2 = try await avgPool.forward(x)
_ = (y1, y2)

let relu = ActivationLayer(type: .relu)
let sig = ActivationLayer(type: .sigmoid)
let sm = ActivationLayer(type: .softmax(axis: 1))

let h = try await relu.forward(logits)
let p = try await sm.forward(h)
_ = (sig, p)

let outShape = NeuralSpatialMath.conv2DOutputShapeNCHW(
    batch: 4,
    height: 32,
    width: 32,
    outChannels: 16,
    kernelHeight: 3,
    kernelWidth: 5,
    strideHeight: 1,
    strideWidth: 2,
    paddingTop: 1,
    paddingBottom: 1,
    paddingLeft: 2,
    paddingRight: 2
)
_ = outShape

import CoreML

let W = try NumericMatrix(rows: 3, columns: 3, elements: Array(repeating: 1, count: 9))
let arr = try W.asMLMultiArray()
// Use `arr` while the owning matrix is retained by the deallocator; avoid mutating `W` concurrently.

// Copy **in** from a rank-2 Float32 MLMultiArray
if let roundtrip = NumericMatrix(mlMultiArray: arr) {
    _ = roundtrip
}

let report = try await GEMMBenchmarkRunner.run(
    sampleCount: 10,
    sizes: [256, 512, 1024],
    includeGpu: nil // default: uses SwiftMetalNumericsCapabilities.isMetalGPUAvailable
)
print(report.formattedTable)

// Or:
// try await GEMMBenchmarkRunner.runPrintingToStandardOut()

import Metal

let device = MTLCreateSystemDefaultDevice()!
let (ptr, bytes) = try BufferManager.allocatePageAlignedByteBuffer(minimumByteCount: 4096)
defer { BufferManager.freeAligned(ptr) }

// Fill `ptr` with your payload, then:
let buffer = try BufferManager.makeSharedBufferNoCopy(
    device: device,
    bytes: ptr,
    length: bytes,
    deallocator: { p, _ in
        BufferManager.freeAligned(p)
    }
)
_ = buffer

var data = [Float](repeating: 1, count: 1024)
_ = try data.withUnsafeMutableBufferPointer { buf in
    try BufferManager.withBridgedSharedBuffer(device: device, contiguous: buf) { mtlBuffer in
        mtlBuffer.length
    }
}