import Foundation

/// Multi-dimensional tensor for compute operations.
///
/// ```swift
/// // Create a 2D tensor
/// let matrix = Tensor(shape: [2, 3], data: [1, 2, 3, 4, 5, 6])
///
/// // Create random tensor
/// let random = Tensor.rand([512, 512])
///
/// // Operations
/// let mean = random.mean()
/// let transposed = matrix.transpose()
/// ```
public struct Tensor: Codable, Equatable {
    public let shape: [Int]
    public let data: [Double]
    public let dtype: Precision

    /// Total number of elements
    public var size: Int {
        shape.reduce(1, *)
    }

    /// Number of dimensions
    public var ndim: Int {
        shape.count
    }

    public init(shape: [Int], data: [Double], dtype: Precision = .fp32) {
        let expectedSize = shape.reduce(1, *)
        precondition(data.count == expectedSize,
                    "Data size \(data.count) does not match shape \(shape)")
        self.shape = shape
        self.data = data
        self.dtype = dtype
    }

    /// Get element at indices
    public subscript(indices: Int...) -> Double {
        precondition(indices.count == shape.count, "Index dimensions must match tensor dimensions")
        var idx = 0
        var stride = 1
        for i in (0..<shape.count).reversed() {
            idx += indices[i] * stride
            stride *= shape[i]
        }
        return data[idx]
    }

    /// Reshape tensor to new shape
    public func reshape(_ newShape: [Int]) -> Tensor {
        let newSize = newShape.reduce(1, *)
        precondition(newSize == size, "Cannot reshape tensor of size \(size) to shape \(newShape)")
        return Tensor(shape: newShape, data: data, dtype: dtype)
    }

    /// Transpose 2D tensor
    public func transpose() -> Tensor {
        precondition(ndim == 2, "Transpose only supported for 2D tensors")
        let rows = shape[0]
        let cols = shape[1]
        var transposed = [Double](repeating: 0, count: data.count)
        for i in 0..<rows {
            for j in 0..<cols {
                transposed[j * rows + i] = data[i * cols + j]
            }
        }
        return Tensor(shape: [cols, rows], data: transposed, dtype: dtype)
    }

    /// Compute mean of all elements
    public func mean() -> Double {
        data.reduce(0, +) / Double(data.count)
    }

    /// Compute sum of all elements
    public func sum() -> Double {
        data.reduce(0, +)
    }

    /// Compute standard deviation
    public func std() -> Double {
        let meanVal = mean()
        let variance = data.map { ($0 - meanVal) * ($0 - meanVal) }.reduce(0, +) / Double(data.count)
        return sqrt(variance)
    }

    /// Find maximum value
    public func max() -> Double {
        data.max() ?? .nan
    }

    /// Find minimum value
    public func min() -> Double {
        data.min() ?? .nan
    }

    /// Apply ReLU activation
    public func relu() -> Tensor {
        Tensor(shape: shape, data: data.map { Swift.max(0, $0) }, dtype: dtype)
    }

    /// Apply sigmoid activation
    public func sigmoid() -> Tensor {
        Tensor(shape: shape, data: data.map { 1 / (1 + exp(-$0)) }, dtype: dtype)
    }

    /// Apply softmax activation
    public func softmax() -> Tensor {
        let maxVal = max()
        let expValues = data.map { exp($0 - maxVal) }
        let sum = expValues.reduce(0, +)
        return Tensor(shape: shape, data: expValues.map { $0 / sum }, dtype: dtype)
    }

    /// Convert to nested array (1D or 2D)
    public func toNestedArray() -> Any {
        switch ndim {
        case 1:
            return data
        case 2:
            let rows = shape[0]
            let cols = shape[1]
            return (0..<rows).map { i in
                (0..<cols).map { j in data[i * cols + j] }
            }
        default:
            fatalError("toNestedArray only supports 1D and 2D tensors")
        }
    }

    // MARK: - Factory Methods

    /// Create tensor from 1D array
    public static func of(_ data: [Double]) -> Tensor {
        Tensor(shape: [data.count], data: data)
    }

    /// Create tensor from 2D array
    public static func of(_ data: [[Double]]) -> Tensor {
        let rows = data.count
        let cols = data[0].count
        let flat = data.flatMap { $0 }
        return Tensor(shape: [rows, cols], data: flat)
    }

    /// Create tensor filled with zeros
    public static func zeros(_ shape: [Int]) -> Tensor {
        let size = shape.reduce(1, *)
        return Tensor(shape: shape, data: [Double](repeating: 0, count: size))
    }

    /// Create tensor filled with ones
    public static func ones(_ shape: [Int]) -> Tensor {
        let size = shape.reduce(1, *)
        return Tensor(shape: shape, data: [Double](repeating: 1, count: size))
    }

    /// Create tensor with uniform random values [0, 1)
    public static func rand(_ shape: [Int]) -> Tensor {
        let size = shape.reduce(1, *)
        return Tensor(shape: shape, data: (0..<size).map { _ in Double.random(in: 0..<1) })
    }

    /// Create tensor with standard normal random values
    public static func randn(_ shape: [Int]) -> Tensor {
        let size = shape.reduce(1, *)
        return Tensor(shape: shape, data: (0..<size).map { _ in
            // Box-Muller transform
            let u1 = Double.random(in: 0..<1)
            let u2 = Double.random(in: 0..<1)
            return sqrt(-2 * log(u1)) * cos(2 * .pi * u2)
        })
    }

    /// Create identity matrix
    public static func eye(_ n: Int) -> Tensor {
        var data = [Double](repeating: 0, count: n * n)
        for i in 0..<n {
            data[i * n + i] = 1
        }
        return Tensor(shape: [n, n], data: data)
    }

    /// Create range tensor
    public static func arange(_ start: Double, _ end: Double, _ step: Double = 1) -> Tensor {
        let size = Int(ceil((end - start) / step))
        return Tensor(shape: [size], data: (0..<size).map { start + Double($0) * step })
    }

    /// Create linearly spaced tensor
    public static func linspace(_ start: Double, _ end: Double, _ num: Int) -> Tensor {
        let step = (end - start) / Double(num - 1)
        return Tensor(shape: [num], data: (0..<num).map { start + Double($0) * step })
    }
}

extension Tensor: CustomStringConvertible {
    public var description: String {
        "Tensor(shape=\(shape), dtype=\(dtype.rawValue))"
    }
}