package io.synor.compute

import kotlinx.serialization.Serializable
import kotlin.math.exp
import kotlin.math.sqrt
import kotlin.random.Random

/**
 * Multi-dimensional tensor for compute operations.
 *
 * ```kotlin
 * // Create a 2D tensor
 * val matrix = Tensor.of(arrayOf(
 *     doubleArrayOf(1.0, 2.0, 3.0),
 *     doubleArrayOf(4.0, 5.0, 6.0)
 * ))
 *
 * // Create random tensor
 * val random = Tensor.rand(512, 512)
 *
 * // Operations
 * val mean = random.mean()
 * val transposed = matrix.transpose()
 * ```
 */
@Serializable
data class Tensor(
    val shape: IntArray,
    val data: DoubleArray,
    val dtype: Precision = Precision.FP32
) {
    init {
        val expectedSize = shape.fold(1) { acc, dim -> acc * dim }
        require(data.size == expectedSize) {
            "Data size ${data.size} does not match shape ${shape.contentToString()}"
        }
    }

    /** Total number of elements */
    val size: Int get() = data.size

    /** Number of dimensions */
    val ndim: Int get() = shape.size

    /** Get element at indices */
    operator fun get(vararg indices: Int): Double {
        require(indices.size == shape.size) { "Index dimensions must match tensor dimensions" }
        var idx = 0
        var stride = 1
        for (i in shape.indices.reversed()) {
            idx += indices[i] * stride
            stride *= shape[i]
        }
        return data[idx]
    }

    /** Reshape tensor to new shape */
    fun reshape(vararg newShape: Int): Tensor {
        val newSize = newShape.fold(1) { acc, dim -> acc * dim }
        require(newSize == size) {
            "Cannot reshape tensor of size $size to shape ${newShape.contentToString()}"
        }
        return Tensor(newShape, data.copyOf(), dtype)
    }

    /** Transpose 2D tensor */
    fun transpose(): Tensor {
        require(ndim == 2) { "Transpose only supported for 2D tensors" }
        val rows = shape[0]
        val cols = shape[1]
        val transposed = DoubleArray(data.size)
        for (i in 0 until rows) {
            for (j in 0 until cols) {
                transposed[j * rows + i] = data[i * cols + j]
            }
        }
        return Tensor(intArrayOf(cols, rows), transposed, dtype)
    }

    /** Compute mean of all elements */
    fun mean(): Double = data.average()

    /** Compute sum of all elements */
    fun sum(): Double = data.sum()

    /** Compute standard deviation */
    fun std(): Double {
        val mean = mean()
        val variance = data.map { (it - mean) * (it - mean) }.average()
        return sqrt(variance)
    }

    /** Find maximum value */
    fun max(): Double = data.max()

    /** Find minimum value */
    fun min(): Double = data.min()

    /** Apply ReLU activation */
    fun relu(): Tensor = Tensor(shape.copyOf(), data.map { maxOf(0.0, it) }.toDoubleArray(), dtype)

    /** Apply sigmoid activation */
    fun sigmoid(): Tensor = Tensor(
        shape.copyOf(),
        data.map { 1.0 / (1.0 + exp(-it)) }.toDoubleArray(),
        dtype
    )

    /** Apply softmax activation */
    fun softmax(): Tensor {
        val maxVal = max()
        val expValues = data.map { exp(it - maxVal) }
        val sum = expValues.sum()
        return Tensor(shape.copyOf(), expValues.map { it / sum }.toDoubleArray(), dtype)
    }

    /** Convert to nested list */
    fun toNestedList(): Any {
        return when (ndim) {
            1 -> data.toList()
            2 -> {
                val rows = shape[0]
                val cols = shape[1]
                (0 until rows).map { i ->
                    (0 until cols).map { j -> data[i * cols + j] }
                }
            }
            else -> throw UnsupportedOperationException("toNestedList only supports 1D and 2D tensors")
        }
    }

    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        if (other !is Tensor) return false
        return shape.contentEquals(other.shape) &&
                data.contentEquals(other.data) &&
                dtype == other.dtype
    }

    override fun hashCode(): Int {
        var result = shape.contentHashCode()
        result = 31 * result + data.contentHashCode()
        result = 31 * result + dtype.hashCode()
        return result
    }

    override fun toString(): String = "Tensor(shape=${shape.contentToString()}, dtype=${dtype.value})"

    companion object {
        /** Create tensor from 1D array */
        fun of(data: DoubleArray): Tensor = Tensor(intArrayOf(data.size), data.copyOf())

        /** Create tensor from 2D array */
        fun of(data: Array<DoubleArray>): Tensor {
            val rows = data.size
            val cols = data[0].size
            val flat = DoubleArray(rows * cols)
            for (i in 0 until rows) {
                System.arraycopy(data[i], 0, flat, i * cols, cols)
            }
            return Tensor(intArrayOf(rows, cols), flat)
        }

        /** Create tensor filled with zeros */
        fun zeros(vararg shape: Int): Tensor =
            Tensor(shape, DoubleArray(shape.fold(1) { acc, d -> acc * d }))

        /** Create tensor filled with ones */
        fun ones(vararg shape: Int): Tensor {
            val size = shape.fold(1) { acc, d -> acc * d }
            return Tensor(shape, DoubleArray(size) { 1.0 })
        }

        /** Create tensor with uniform random values [0, 1) */
        fun rand(vararg shape: Int): Tensor {
            val size = shape.fold(1) { acc, d -> acc * d }
            return Tensor(shape, DoubleArray(size) { Random.nextDouble() })
        }

        /** Create tensor with standard normal random values */
        fun randn(vararg shape: Int): Tensor {
            val size = shape.fold(1) { acc, d -> acc * d }
            return Tensor(shape, DoubleArray(size) { Random.nextGaussian() })
        }

        /** Create identity matrix */
        fun eye(n: Int): Tensor {
            val data = DoubleArray(n * n)
            for (i in 0 until n) {
                data[i * n + i] = 1.0
            }
            return Tensor(intArrayOf(n, n), data)
        }

        /** Create range tensor */
        fun arange(start: Double, end: Double, step: Double = 1.0): Tensor {
            val size = ((end - start) / step).toInt()
            return Tensor(intArrayOf(size), DoubleArray(size) { start + it * step })
        }

        /** Create linearly spaced tensor */
        fun linspace(start: Double, end: Double, num: Int): Tensor {
            val step = (end - start) / (num - 1)
            return Tensor(intArrayOf(num), DoubleArray(num) { start + it * step })
        }
    }
}

private fun Random.nextGaussian(): Double {
    // Box-Muller transform
    val u1 = nextDouble()
    val u2 = nextDouble()
    return sqrt(-2 * kotlin.math.ln(u1)) * kotlin.math.cos(2 * Math.PI * u2)
}
