using System;
using System.Buffers;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Numerics;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;

namespace HyperMath.src.Core
{
    /// <summary>
    /// Represents a multi-dimensional immutable tensor with generic type T.
    /// </summary>
    /// <typeparam name="T">The type of elements in the tensor. Supported types: float, double.</typeparam>
    public class Tensor<T> : IDisposable where T : struct, INumber<T>, IExponentialFunctions<T>, IComparable, IConvertible, IEquatable<T>, IFormattable
    {
        /// <summary>
        /// Manages the shared data array with reference counting.
        /// </summary>
        private class SharedData
        {
            public readonly T[] Data;
            private int refCount;

            public SharedData(int size)
            {
                Data = ArrayPool<T>.Shared.Rent(size);
                Array.Clear(Data, 0, size);
                refCount = 1;
            }

            public void AddRef()
            {
                Interlocked.Increment(ref refCount);
            }

            public void Release()
            {
                if (Interlocked.Decrement(ref refCount) == 0)
                {
                    ArrayPool<T>.Shared.Return(Data, clearArray: true);
                }
            }
        }

        private readonly SharedData sharedData;
        private readonly int[] shape;
        private readonly int[] strides;
        private readonly int offset;
        private readonly bool isView;
        private bool disposed = false;

        /// <summary>
        /// Gets the shape of the tensor.
        /// </summary>
        public int[] Shape => (int[])shape.Clone();

        /// <summary>
        /// Gets the total number of elements in the tensor.
        /// </summary>
        public int TotalSize { get; }

        /// <summary>
        /// Threshold for enabling parallelization.
        /// </summary>
        public static int ParallelizationThreshold { get; set; } = 1000;

        /// <summary>
        /// Initializes a new tensor with the specified dimensions.
        /// </summary>
        /// <param name="dimensions">Sizes for each dimension.</param>
        public Tensor(params int[] dimensions)
        {
            if (dimensions == null || dimensions.Length == 0)
                throw new ArgumentException("Tensor must have at least one dimension.", nameof(dimensions));

            foreach (var dim in dimensions)
            {
                if (dim <= 0)
                    throw new ArgumentException("Each dimension size must be positive.", nameof(dimensions));
            }

            shape = (int[])dimensions.Clone();
            strides = CalculateStrides(shape);
            TotalSize = shape.Aggregate(1, (a, b) => a * b);
            sharedData = new SharedData(TotalSize);
            offset = 0;
            isView = false;
        }

        private Tensor(SharedData sharedData, int[] shape, int[] strides, int offset)
        {
            this.sharedData = sharedData;
            this.shape = shape;
            this.strides = strides;
            this.offset = offset;
            TotalSize = shape.Aggregate(1, (a, b) => a * b);
            isView = true;
            sharedData.AddRef();
        }

        /// <summary>
        /// Calculates the strides for row-major order.
        /// </summary>
        /// <param name="shape">Shape of the tensor.</param>
        /// <returns>Strides array.</returns>
        private static int[] CalculateStrides(int[] shape)
        {
            int[] strides = new int[shape.Length];
            strides[shape.Length - 1] = 1;
            for (int i = shape.Length - 2; i >= 0; i--)
            {
                strides[i] = strides[i + 1] * shape[i + 1];
            }
            return strides;
        }

        /// <summary>
        /// Creates a tensor view (slice) from the current tensor.
        /// </summary>
        /// <param name="ranges">Ranges for each dimension.</param>
        /// <returns>A new tensor view.</returns>
        public Tensor<T> Slice(params Range[] ranges)
        {
            EnsureNotDisposed();

            if (ranges == null)
                throw new ArgumentNullException(nameof(ranges));

            if (ranges.Length != shape.Length)
                throw new ArgumentException("Number of ranges must match number of tensor dimensions.", nameof(ranges));

            int[] newShape = new int[shape.Length];
            int newOffset = offset;

            for (int i = 0; i < shape.Length; i++)
            {
                var range = ranges[i];
                int start = range.Start.IsFromEnd ? shape[i] - range.Start.Value : range.Start.Value;
                int end = range.End.IsFromEnd ? shape[i] - range.End.Value : range.End.Value;

                if (start < 0 || end > shape[i] || start > end)
                    throw new ArgumentOutOfRangeException(nameof(ranges), $"Invalid range for dimension {i}.");

                newShape[i] = end - start;
                newOffset += start * strides[i];
            }

            int[] newStrides = (int[])strides.Clone();

            return new Tensor<T>(sharedData, newShape, newStrides, newOffset);
        }

        /// <summary>
        /// Provides advanced slicing with steps.
        /// </summary>
        /// <param name="ranges">Ranges with steps for each dimension.</param>
        /// <returns>A new tensor view with steps.</returns>
        public Tensor<T> SliceWithSteps(params (int start, int end, int step)[] ranges)
        {
            EnsureNotDisposed();

            if (ranges == null)
                throw new ArgumentNullException(nameof(ranges));

            if (ranges.Length != shape.Length)
                throw new ArgumentException("Number of ranges must match number of tensor dimensions.", nameof(ranges));

            int[] newShape = new int[shape.Length];
            int newOffset = offset;
            int[] newStrides = new int[shape.Length];

            for (int i = 0; i < shape.Length; i++)
            {
                var (start, end, step) = ranges[i];
                if (step <= 0)
                    throw new ArgumentException("Step must be positive.", nameof(ranges));

                if (start < 0 || start >= shape[i])
                    throw new ArgumentOutOfRangeException(nameof(ranges), $"Start index out of bounds for dimension {i}.");

                if (end < start || end > shape[i])
                    throw new ArgumentOutOfRangeException(nameof(ranges), $"End index out of bounds for dimension {i}.");

                newShape[i] = (int)Math.Ceiling((double)(end - start) / step);
                newOffset += start * strides[i];
                newStrides[i] = strides[i] * step;
            }

            Array.Copy(strides, newStrides, strides.Length);

            return new Tensor<T>(sharedData, newShape, newStrides, newOffset);
        }

        /// <summary>
        /// Gets the element at the specified indices.
        /// </summary>
        /// <param name="indices">Multi-dimensional indices.</param>
        /// <returns>Value at the specified indices.</returns>
        public T GetElement(params int[] indices)
        {
            EnsureNotDisposed();
            int flatIndex = CalculateFlatIndex(indices);
            return sharedData.Data[flatIndex];
        }

        /// <summary>
        /// Returns a new tensor with the specified value set at the given indices.
        /// </summary>
        /// <param name="value">Value to set.</param>
        /// <param name="indices">Multi-dimensional indices.</param>
        /// <returns>New tensor with the value set.</returns>
        public Tensor<T> SetElement(T value, params int[] indices)
        {
            EnsureNotDisposed();
            int flatIndex = CalculateFlatIndex(indices);
            Tensor<T> newTensor = this.Clone();
            newTensor.sharedData.Data[flatIndex] = value;
            return newTensor;
        }

        /// <summary>
        /// Calculates the flat index from multi-dimensional indices.
        /// </summary>
        /// <param name="indices">Multi-dimensional indices.</param>
        /// <returns>Flat index.</returns>
        private int CalculateFlatIndex(int[] indices)
        {
            if (indices == null)
                throw new ArgumentNullException(nameof(indices));

            if (indices.Length != shape.Length)
                throw new ArgumentException("Number of indices must match number of tensor dimensions.", nameof(indices));

            int flatIndex = offset;
            for (int i = 0; i < shape.Length; i++)
            {
                if (indices[i] < 0 || indices[i] >= shape[i])
                    throw new ArgumentOutOfRangeException(nameof(indices), $"Index {indices[i]} out of bounds for dimension {i}.");

                flatIndex += indices[i] * strides[i];
            }

            if (flatIndex < 0 || flatIndex >= sharedData.Data.Length)
                throw new IndexOutOfRangeException("Flat index out of range.");

            return flatIndex;
        }

        /// <summary>
        /// Clones the tensor, creating a new tensor with the same data.
        /// </summary>
        /// <returns>A new tensor cloned from the current tensor.</returns>
        public Tensor<T> Clone()
        {
            EnsureNotDisposed();

            Tensor<T> newTensor = new Tensor<T>(shape);
            Array.Copy(sharedData.Data, offset, newTensor.sharedData.Data, 0, TotalSize);
            return newTensor;
        }

        /// <summary>
        /// Adds two tensors element-wise.
        /// </summary>
        /// <param name="a">First tensor.</param>
        /// <param name="b">Second tensor.</param>
        /// <returns>Element-wise sum of tensors.</returns>
        public static Tensor<T> operator +(Tensor<T> a, Tensor<T> b)
        {
            return a.ElementWiseAdd(b);
        }

        /// <summary>
        /// Subtracts two tensors element-wise.
        /// </summary>
        /// <param name="a">First tensor.</param>
        /// <param name="b">Second tensor.</param>
        /// <returns>Element-wise difference of tensors.</returns>
        public static Tensor<T> operator -(Tensor<T> a, Tensor<T> b)
        {
            return a.ElementWiseSubtract(b);
        }

        /// <summary>
        /// Multiplies two tensors element-wise.
        /// </summary>
        /// <param name="a">First tensor.</param>
        /// <param name="b">Second tensor.</param>
        /// <returns>Element-wise product of tensors.</returns>
        public static Tensor<T> operator *(Tensor<T> a, Tensor<T> b)
        {
            return a.ElementWiseMultiply(b);
        }

        /// <summary>
        /// Divides two tensors element-wise.
        /// </summary>
        /// <param name="a">First tensor.</param>
        /// <param name="b">Second tensor.</param>
        /// <returns>Element-wise division of tensors.</returns>
        public static Tensor<T> operator /(Tensor<T> a, Tensor<T> b)
        {
            return a.ElementWiseDivide(b);
        }

        /// <summary>
        /// Adds a scalar to each element of the tensor.
        /// </summary>
        /// <param name="tensor">The tensor.</param>
        /// <param name="scalar">The scalar value.</param>
        /// <returns>A new tensor with the scalar added.</returns>
        public static Tensor<T> operator +(Tensor<T> tensor, T scalar)
        {
            return tensor.AddScalar(scalar);
        }

        /// <summary>
        /// Subtracts a scalar from each element of the tensor.
        /// </summary>
        /// <param name="tensor">The tensor.</param>
        /// <param name="scalar">The scalar value.</param>
        /// <returns>A new tensor with the scalar subtracted.</returns>
        public static Tensor<T> operator -(Tensor<T> tensor, T scalar)
        {
            return tensor.SubtractScalar(scalar);
        }

        /// <summary>
        /// Multiplies each element of the tensor by a scalar.
        /// </summary>
        /// <param name="tensor">The tensor.</param>
        /// <param name="scalar">The scalar value.</param>
        /// <returns>A new tensor with the scalar multiplied.</returns>
        public static Tensor<T> operator *(Tensor<T> tensor, T scalar)
        {
            return tensor.MultiplyScalar(scalar);
        }

        /// <summary>
        /// Divides each element of the tensor by a scalar.
        /// </summary>
        /// <param name="tensor">The tensor.</param>
        /// <param name="scalar">The scalar value.</param>
        /// <returns>A new tensor with the scalar divided.</returns>
        public static Tensor<T> operator /(Tensor<T> tensor, T scalar)
        {
            return tensor.DivideScalar(scalar);
        }

        /// <summary>
        /// Performs element-wise addition with another tensor.
        /// </summary>
        /// <param name="other">The other tensor.</param>
        /// <returns>Element-wise sum.</returns>
        public Tensor<T> ElementWiseAdd(Tensor<T> other)
        {
            EnsureNotDisposed();
            other?.EnsureNotDisposed();

            Tensor<T> result = ShapeMatches(other)
                ? new Tensor<T>(shape)
                : BroadcastOperation(other, (a, b) => a + b, out _, out _);

            if (ShapeMatches(other))
            {
                Parallel.For(0, TotalSize, i =>
                {
                    result.sharedData.Data[i] = sharedData.Data[offset + i] + other.sharedData.Data[other.offset + i];
                });
            }

            return result;
        }

        /// <summary>
        /// Performs element-wise subtraction with another tensor.
        /// </summary>
        /// <param name="other">The other tensor.</param>
        /// <returns>Element-wise difference.</returns>
        public Tensor<T> ElementWiseSubtract(Tensor<T> other)
        {
            EnsureNotDisposed();
            other?.EnsureNotDisposed();

            Tensor<T> result = ShapeMatches(other)
                ? new Tensor<T>(shape)
                : BroadcastOperation(other, (a, b) => a - b, out _, out _);

            if (ShapeMatches(other))
            {
                Parallel.For(0, TotalSize, i =>
                {
                    result.sharedData.Data[i] = sharedData.Data[offset + i] - other.sharedData.Data[other.offset + i];
                });
            }

            return result;
        }

        /// <summary>
        /// Performs element-wise multiplication with another tensor.
        /// </summary>
        /// <param name="other">The other tensor.</param>
        /// <returns>Element-wise product.</returns>
        public Tensor<T> ElementWiseMultiply(Tensor<T> other)
        {
            EnsureNotDisposed();
            other?.EnsureNotDisposed();

            Tensor<T> result = ShapeMatches(other)
                ? new Tensor<T>(shape)
                : BroadcastOperation(other, (a, b) => a * b, out _, out _);

            if (ShapeMatches(other))
            {
                Parallel.For(0, TotalSize, i =>
                {
                    result.sharedData.Data[i] = sharedData.Data[offset + i] * other.sharedData.Data[other.offset + i];
                });
            }

            return result;
        }

        /// <summary>
        /// Performs element-wise division with another tensor.
        /// </summary>
        /// <param name="other">The other tensor.</param>
        /// <returns>Element-wise division.</returns>
        public Tensor<T> ElementWiseDivide(Tensor<T> other)
        {
            EnsureNotDisposed();
            other?.EnsureNotDisposed();

            Tensor<T> result = ShapeMatches(other)
                ? new Tensor<T>(shape)
                : BroadcastOperation(other, (a, b) => a / b, out _, out _);

            if (ShapeMatches(other))
            {
                Parallel.For(0, TotalSize, i =>
                {
                    result.sharedData.Data[i] = sharedData.Data[offset + i] / other.sharedData.Data[other.offset + i];
                });
            }

            return result;
        }

        /// <summary>
        /// Adds a scalar to each element of the tensor.
        /// </summary>
        /// <param name="scalar">Scalar to add.</param>
        /// <returns>A new tensor with the scalar added.</returns>
        public Tensor<T> AddScalar(T scalar)
        {
            EnsureNotDisposed();

            Tensor<T> result = new Tensor<T>(shape);

            if (TotalSize >= ParallelizationThreshold)
            {
                if (IsSupportedSIMD())
                {
                    ConcurrentOperationSIMD(sharedData.Data, offset, result.sharedData.Data, 0, scalar, (a, s) => a + s);
                }
                else
                {
                    Parallel.For(0, TotalSize, i =>
                    {
                        result.sharedData.Data[i] = sharedData.Data[offset + i] + scalar;
                    });
                }
            }
            else
            {
                for (int i = 0; i < TotalSize; i++)
                {
                    result.sharedData.Data[i] = sharedData.Data[offset + i] + scalar;
                }
            }

            return result;
        }

        /// <summary>
        /// Subtracts a scalar from each element of the tensor.
        /// </summary>
        /// <param name="scalar">Scalar to subtract.</param>
        /// <returns>A new tensor with the scalar subtracted.</returns>
        public Tensor<T> SubtractScalar(T scalar)
        {
            EnsureNotDisposed();

            Tensor<T> result = new Tensor<T>(shape);

            if (TotalSize >= ParallelizationThreshold)
            {
                if (IsSupportedSIMD())
                {
                    ConcurrentOperationSIMD(sharedData.Data, offset, result.sharedData.Data, 0, scalar, (a, s) => a - s);
                }
                else
                {
                    Parallel.For(0, TotalSize, i =>
                    {
                        result.sharedData.Data[i] = sharedData.Data[offset + i] - scalar;
                    });
                }
            }
            else
            {
                for (int i = 0; i < TotalSize; i++)
                {
                    result.sharedData.Data[i] = sharedData.Data[offset + i] - scalar;
                }
            }

            return result;
        }

        /// <summary>
        /// Multiplies each element of the tensor by a scalar.
        /// </summary>
        /// <param name="scalar">Scalar to multiply.</param>
        /// <returns>A new tensor with the scalar multiplied.</returns>
        public Tensor<T> MultiplyScalar(T scalar)
        {
            EnsureNotDisposed();

            Tensor<T> result = new Tensor<T>(shape);

            if (TotalSize >= ParallelizationThreshold)
            {
                if (IsSupportedSIMD())
                {
                    ConcurrentOperationSIMD(sharedData.Data, offset, result.sharedData.Data, 0, scalar, (a, s) => a * s);
                }
                else
                {
                    Parallel.For(0, TotalSize, i =>
                    {
                        result.sharedData.Data[i] = sharedData.Data[offset + i] * scalar;
                    });
                }
            }
            else
            {
                for (int i = 0; i < TotalSize; i++)
                {
                    result.sharedData.Data[i] = sharedData.Data[offset + i] * scalar;
                }
            }

            return result;
        }

        /// <summary>
        /// Divides each element of the tensor by a scalar.
        /// </summary>
        /// <param name="scalar">Scalar to divide.</param>
        /// <returns>A new tensor with the scalar divided.</returns>
        public Tensor<T> DivideScalar(T scalar)
        {
            EnsureNotDisposed();

            Tensor<T> result = new Tensor<T>(shape);

            if (TotalSize >= ParallelizationThreshold)
            {
                if (IsSupportedSIMD())
                {
                    ConcurrentOperationSIMD(sharedData.Data, offset, result.sharedData.Data, 0, scalar, (a, s) => a / s);
                }
                else
                {
                    Parallel.For(0, TotalSize, i =>
                    {
                        result.sharedData.Data[i] = sharedData.Data[offset + i] / scalar;
                    });
                }
            }
            else
            {
                for (int i = 0; i < TotalSize; i++)
                {
                    result.sharedData.Data[i] = sharedData.Data[offset + i] / scalar;
                }
            }

            return result;
        }

        /// <summary>
        /// Performs element-wise addition with broadcasting support.
        /// </summary>
        /// <param name="other">The other tensor to add.</param>
        /// <param name="operation">The binary operation to perform.</param>
        /// <param name="resultShape">Output parameter for Shape of result tensor.</param>
        /// <param name="resultStrides">Output parameter for Strides of result tensor.</param>
        /// <returns>New tensor after the operation.</returns>
        private Tensor<T> BroadcastOperation(Tensor<T> other, Func<T, T, T> operation, out int[] resultShape, out int[] resultStrides)
        {
            // Determine the shape using broadcasting rules
            resultShape = DetermineBroadcastShape(this.shape, other.shape);
            resultStrides = CalculateStrides(resultShape);

            Tensor<T> result = new Tensor<T>(resultShape);

            // Precompute broadcasted strides
            var aStrides = ComputeBroadcastStrides(this.shape, resultShape, this.strides);
            var bStrides = ComputeBroadcastStrides(other.shape, resultShape, other.strides);

            Parallel.For(0, result.TotalSize, i =>
            {
                int[] indices = UnflattenIndex(i, resultShape);
                int aIndex = ComputeFlatIndex(indices, aStrides);
                int bIndex = ComputeFlatIndex(indices, bStrides);
                result.sharedData.Data[i] = operation(this.sharedData.Data[aIndex], other.sharedData.Data[bIndex]);
            });

            return result;
        }

        /// <summary>
        /// Determines if two tensors have the same shape.
        /// </summary>
        /// <param name="other">Other tensor.</param>
        /// <returns>True if shapes match; otherwise, false.</returns>
        private bool ShapeMatches(Tensor<T> other)
        {
            if (other == null)
                throw new ArgumentNullException(nameof(other));

            if (shape.Length != other.shape.Length)
                return false;

            for (int i = 0; i < shape.Length; i++)
            {
                if (shape[i] != other.shape[i])
                    return false;
            }

            return true;
        }

        /// <summary>
        /// Determines the broadcasted shape from two shapes.
        /// </summary>
        /// <param name="aShape">Shape of the first tensor.</param>
        /// <param name="bShape">Shape of the second tensor.</param>
        /// <returns>Broadcasted shape.</returns>
        private static int[] DetermineBroadcastShape(int[] aShape, int[] bShape)
        {
            int maxLength = Math.Max(aShape.Length, bShape.Length);
            int[] broadcastShape = new int[maxLength];

            for (int i = 0; i < maxLength; i++)
            {
                int aDim = i < aShape.Length ? aShape[aShape.Length - 1 - i] : 1;
                int bDim = i < bShape.Length ? bShape[bShape.Length - 1 - i] : 1;

                if (aDim != bDim && aDim != 1 && bDim != 1)
                    throw new InvalidOperationException($"Shapes {string.Join(",", aShape)} and {string.Join(",", bShape)} are not broadcastable.");

                broadcastShape[maxLength - 1 - i] = Math.Max(aDim, bDim);
            }

            return broadcastShape;
        }

        /// <summary>
        /// Computes broadcasted strides for a tensor based on the result shape.
        /// </summary>
        /// <param name="aShape">Original shape of the tensor.</param>
        /// <param name="resultShape">Broadcasted shape.</param>
        /// <param name="aStrides">Original strides of the tensor.</param>
        /// <returns>Broadcasted strides.</returns>
        private static int[] ComputeBroadcastStrides(int[] aShape, int[] resultShape, int[] aStrides)
        {
            int aOffset = aShape.Length - resultShape.Length;
            int[] aBroadcastStrides = new int[resultShape.Length];

            for (int i = 0; i < resultShape.Length; i++)
            {
                if (i + aOffset < 0 || i + aOffset >= aShape.Length)
                {
                    aBroadcastStrides[i] = 0;
                }
                else if (aShape[i + aOffset] == 1)
                {
                    aBroadcastStrides[i] = 0;
                }
                else
                {
                    aBroadcastStrides[i] = aStrides[i + aOffset];
                }
            }

            return aBroadcastStrides;
        }

        /// <summary>
        /// Performs a matrix multiplication with another tensor.
        /// </summary>
        /// <param name="other">The other tensor to multiply with.</param>
        /// <returns>The result of matrix multiplication.</returns>
        public Tensor<T> MatMul(Tensor<T> other)
        {
            EnsureNotDisposed();
            other?.EnsureNotDisposed();

            if (this.shape.Length != 2 || other.shape.Length != 2)
                throw new ArgumentException("MatMul is only supported for 2D tensors.");

            if (this.shape[1] != other.shape[0])
                throw new ArgumentException($"Inner tensor dimensions must match for MatMul. Found {this.shape[1]} and {other.shape[0]}.");

            int m = this.shape[0];
            int n = this.shape[1];
            int p = other.shape[1];

            Tensor<T> result = new Tensor<T>(m, p);

            // Optimized matrix multiplication using parallelism and SIMD
            Parallel.For(0, m, i =>
            {
                for (int j = 0; j < p; j++)
                {
                    T sum = T.Zero;
                    int aIndex = i * this.strides[0];
                    int bIndex = j;

                    if (IsSupportedSIMD())
                    {
                        // SIMD optimized inner loop
                        int k = 0;
                        int simdWidth = Vector<T>.Count;

                        Vector<T> vecSum = Vector<T>.Zero;
                        for (; k <= n - simdWidth; k += simdWidth)
                        {
                            Vector<T> vecA = new Vector<T>(sharedData.Data, aIndex + k);
                            Vector<T> vecB = new Vector<T>(other.sharedData.Data, other.offset + k * other.strides[0] + j);
                            vecSum += vecA * vecB;
                        }

                        for (; k < n; k++)
                        {
                            sum += sharedData.Data[aIndex + k] * other.sharedData.Data[other.offset + k * other.strides[0] + j];
                        }

                        for (int s = 0; s < Vector<T>.Count; s++)
                        {
                            sum += vecSum[s];
                        }
                    }
                    else
                    {
                        // Scalar inner loop
                        for (int k = 0; k < n; k++)
                        {
                            sum += sharedData.Data[aIndex + k] * other.sharedData.Data[other.offset + k * other.strides[0] + j];
                        }
                    }

                    result.sharedData.Data[i * result.strides[0] + j] = sum;
                }
            });

            return result;
        }

        /// <summary>
        /// Performs batched matrix multiplication with another tensor.
        /// Assumes batch dimensions are the leading dimensions.
        /// </summary>
        /// <param name="other">The other tensor to perform batched multiplication with.</param>
        /// <returns>The result of batched matrix multiplication.</returns>
        public Tensor<T> BatchMatrixMultiply(Tensor<T> other)
        {
            EnsureNotDisposed();
            other?.EnsureNotDisposed();

            if (this.shape.Length < 2 || other.shape.Length < 2)
                throw new ArgumentException("BatchedMatMul requires tensors with at least 2 dimensions.");

            int aBatch = this.shape.Length > 2 ? this.shape.Take(this.shape.Length - 2).Aggregate(1, (a, b) => a * b) : 1;
            int aRows = this.shape[this.shape.Length - 2];
            int aCols = this.shape[this.shape.Length - 1];

            int bBatch = other.shape.Length > 2 ? other.shape.Take(other.shape.Length - 2).Aggregate(1, (a, b) => a * b) : 1;
            int bRows = other.shape[other.shape.Length - 2];
            int bCols = other.shape[other.shape.Length - 1];

            if (aCols != bRows)
                throw new ArgumentException($"Inner tensor dimensions must match for MatMul. Found {aCols} and {bRows}.");

            if (aBatch != bBatch)
                throw new ArgumentException($"Batch size must match for BatchedMatMul. Found {aBatch} and {bBatch}.");

            Tensor<T> result = new Tensor<T>(aBatch, aRows, bCols);

            Parallel.For(0, aBatch, batch =>
            {
                for (int i = 0; i < aRows; i++)
                {
                    for (int j = 0; j < bCols; j++)
                    {
                        T sum = T.Zero;
                        int aIndex = batch * this.strides[0] + i * this.strides[this.shape.Length - 2];
                        int bIndex = batch * other.strides[0] + j;

                        if (IsSupportedSIMD())
                        {
                            int k = 0;
                            int n = aCols;
                            int simdWidth = Vector<T>.Count;

                            Vector<T> vecSum = Vector<T>.Zero;
                            for (; k <= n - simdWidth; k += simdWidth)
                            {
                                Vector<T> vecA = new Vector<T>(sharedData.Data, aIndex + k);
                                Vector<T> vecB = new Vector<T>(other.sharedData.Data, other.offset + k * other.strides[other.shape.Length - 1] + j);
                                vecSum += vecA * vecB;
                            }

                            for (; k < n; k++)
                            {
                                sum += sharedData.Data[aIndex + k] * other.sharedData.Data[other.offset + k * other.strides[other.shape.Length - 1] + j];
                            }

                            for (int s = 0; s < Vector<T>.Count; s++)
                            {
                                sum += vecSum[s];
                            }
                        }
                        else
                        {
                            for (int k = 0; k < aCols; k++)
                            {
                                sum += sharedData.Data[aIndex + k] * other.sharedData.Data[other.offset + k * other.strides[other.shape.Length - 1] + j];
                            }
                        }

                        result.sharedData.Data[batch * result.strides[0] + i * result.strides[result.shape.Length - 2] + j] = sum;
                    }
                }
            });

            return result;
        }

        /// <summary>
        /// Performs a generalized tensor dot product with another tensor over specified axes.
        /// </summary>
        /// <param name="other">The other tensor to perform the dot product with.</param>
        /// <param name="thisAxes">The axes of the current tensor to contract over.</param>
        /// <param name="otherAxes">The axes of the other tensor to contract over.</param>
        /// <returns>A new tensor representing the dot product.</returns>
        public Tensor<T> DotProduct(Tensor<T> other, int[]? thisAxes = null, int[]? otherAxes = null)
        {
            EnsureNotDisposed();
            other?.EnsureNotDisposed();

            // Default axes: last axis of this tensor and first axis of other tensor
            if (thisAxes == null || otherAxes == null)
            {
                thisAxes = thisAxes ?? new[] { shape.Length - 1 };
                otherAxes = otherAxes ?? new[] { 0 };
            }

            if (thisAxes.Length != otherAxes.Length)
                throw new ArgumentException("The number of axes to contract must be the same for both tensors.");

            // Validate the axes
            for (int i = 0; i < thisAxes.Length; i++)
            {
                if (thisAxes[i] < 0 || thisAxes[i] >= shape.Length)
                    throw new ArgumentOutOfRangeException($"thisAxes[{i}] = {thisAxes[i]} is out of bounds for tensor with {shape.Length} dimensions.");
                if (otherAxes[i] < 0 || otherAxes[i] >= other.shape.Length)
                    throw new ArgumentOutOfRangeException($"otherAxes[{i}] = {otherAxes[i]} is out of bounds for other tensor with {other.shape.Length} dimensions.");

                if (shape[thisAxes[i]] != other.shape[otherAxes[i]])
                    throw new ArgumentException($"Dimension mismatch on axis {i}: {shape[thisAxes[i]]} != {other.shape[otherAxes[i]]}");
            }

            // Determine free and contraction axes
            var thisFreeAxes = Enumerable.Range(0, shape.Length).Where(d => !thisAxes.Contains(d)).ToArray();
            var otherFreeAxes = Enumerable.Range(0, other.shape.Length).Where(d => !otherAxes.Contains(d)).ToArray();

            // Calculate the shape of the result tensor
            var resultShape = thisFreeAxes.Select(d => shape[d])
                                .Concat(otherFreeAxes.Select(d => other.shape[d]))
                                .ToArray();

            Tensor<T> result = new Tensor<T>(resultShape);

            // Precompute strides
            var thisBroadcastStrides = ComputeBroadcastStrides(shape, resultShape, strides, thisFreeAxes, thisAxes);
            var otherBroadcastStrides = ComputeBroadcastStrides(other.shape, resultShape, other.strides, otherFreeAxes, otherAxes);

            // Calculate the length of the contraction
            int contractionLength = thisAxes.Aggregate(1, (a, b) => a * shape[b]);

            // Perform the dot product
            Parallel.For(0, result.TotalSize, i =>
            {
                int[] resultIndices = UnflattenIndex(i, resultShape);
                int thisIndex = 0;
                int otherIndex = 0;

                // Calculate base indices for this tensor
                for (int a = 0; a < thisFreeAxes.Length; a++)
                {
                    thisIndex += resultIndices[a] * thisBroadcastStrides[a];
                }

                // Calculate base indices for other tensor
                for (int b = 0; b < otherFreeAxes.Length; b++)
                {
                    otherIndex += resultIndices[thisFreeAxes.Length + b] * otherBroadcastStrides[b];
                }

                // Accumulate over the contraction axes
                T sum = T.Zero;
                for (int k = 0; k < contractionLength; k++)
                {
                    int a_k = thisIndex + k * thisBroadcastStrides[thisFreeAxes.Length + thisAxes.Length - 1];
                    int b_k = otherIndex + k * otherBroadcastStrides[otherFreeAxes.Length + otherAxes.Length - 1];
                    sum += sharedData.Data[a_k] * other.sharedData.Data[b_k];
                }

                result.sharedData.Data[i] = sum;
            });

            return result;
        }

        /// <summary>
        /// Computes broadcasted strides for tensor operations.
        /// </summary>
        /// <param name="originalShape">Original shape of the tensor.</param>
        /// <param name="resultShape">Shape after broadcasting.</param>
        /// <param name="originalStrides">Original strides of the tensor.</param>
        /// <param name="freeAxes">Free axes of the tensor.</param>
        /// <param name="contractAxes">Contraction axes of the tensor.</param>
        /// <returns>Array of broadcasted strides.</returns>
        private static int[] ComputeBroadcastStrides(int[] originalShape, int[] resultShape, int[] originalStrides, int[] freeAxes, int[] contractAxes)
        {
            var broadcastStrides = new int[resultShape.Length];
            int originalRank = originalShape.Length;

            for (int i = 0; i < resultShape.Length; i++)
            {
                int resultDim = resultShape[i];
                if (i >= originalShape.Length)
                {
                    broadcastStrides[i] = 0;
                }
                else
                {
                    if (originalShape[i] == resultDim)
                    {
                        broadcastStrides[i] = originalStrides[i];
                    }
                    else if (originalShape[i] == 1)
                    {
                        broadcastStrides[i] = 0;
                    }
                    else
                    {
                        throw new InvalidOperationException("Shapes are not broadcastable.");
                    }
                }
            }

            return broadcastStrides;
        }

        /// <summary>
        /// Determines if SIMD is supported for the current type.
        /// </summary>
        /// <returns>True if SIMD is supported; otherwise, false.</returns>
        private bool IsSupportedSIMD()
        {
            return typeof(T) == typeof(float) || typeof(T) == typeof(double);
        }

        /// <summary>
        /// Performs a concurrent operation using SIMD where possible.
        /// </summary>
        /// <param name="source">Source data array.</param>
        /// <param name="sourceOffset">Offset in the source array.</param>
        /// <param name="destination">Destination data array.</param>
        /// <param name="destinationOffset">Offset in the destination array.</param>
        /// <param name="scalar">Scalar value for the operation.</param>
        /// <param name="operation">Binary operation to perform.</param>
        private void ConcurrentOperationSIMD(T[] source, int sourceOffset, T[] destination, int destinationOffset, T scalar, Func<T, T, T> operation)
        {
            if (typeof(T) == typeof(float))
            {
                Vector<float> vecScalar = new Vector<float>((float)(object)scalar);
                int vectorSize = Vector<float>.Count;
                int i = 0;

                // Process in chunks of vectorSize
                for (; i <= TotalSize - vectorSize; i += vectorSize)
                {
                    var vec = new Vector<float>(source, sourceOffset + i);
                    Vector<float> resultVec = vec + vecScalar; // Operations must match the Func
                    resultVec.CopyTo(destination, destinationOffset + i);
                }

                // Process remaining elements
                for (; i < TotalSize; i++)
                {
                    destination[destinationOffset + i] = operation((T)(object)source[sourceOffset + i], scalar);
                }
            }
            else if (typeof(T) == typeof(double))
            {
                Vector<double> vecScalar = new Vector<double>((double)(object)scalar);
                int vectorSize = Vector<double>.Count;
                int i = 0;

                // Process in chunks of vectorSize
                for (; i <= TotalSize - vectorSize; i += vectorSize)
                {
                    var vec = new Vector<double>(source, sourceOffset + i);
                    Vector<double> resultVec = vec + vecScalar; // Operations must match the Func
                    resultVec.CopyTo(destination, destinationOffset + i);
                }

                // Process remaining elements
                for (; i < TotalSize; i++)
                {
                    destination[destinationOffset + i] = operation((T)(object)source[sourceOffset + i], scalar);
                }
            }
            else
            {
                throw new InvalidOperationException("SIMD operations are only supported for float and double types.");
            }
        }

        /// <summary>
        /// Performs a generalized element-wise operation with broadcasting support.
        /// </summary>
        /// <param name="other">The other tensor to operate with.</param>
        /// <param name="operation">The binary operation to perform.</param>
        /// <returns>A new tensor after the operation.</returns>
        private Tensor<T> GeneralizedElementWiseOperation(Tensor<T> other, Func<T, T, T> operation)
        {
            EnsureNotDisposed();
            other?.EnsureNotDisposed();

            Tensor<T> result = ShapeMatches(other)
                ? new Tensor<T>(shape)
                : BroadcastOperation(other, operation, out int[] resultShape, out int[] resultStrides);

            if (ShapeMatches(other))
            {
                Parallel.For(0, TotalSize, i =>
                {
                    result.sharedData.Data[i] = operation(sharedData.Data[offset + i], other.sharedData.Data[other.offset + i]);
                });
            }

            return result;
        }

        /// <summary>
        /// Performs element-wise power operation with another tensor.
        /// </summary>
        /// <param name="other">The exponent tensor.</param>
        /// <returns>A new tensor after the power operation.</returns>
        public Tensor<T> ElementWisePower(Tensor<T> other)
        {
            return GeneralizedElementWiseOperation(other, (a, b) => INumber<T>.Pow(a, b));
        }

        /// <summary>
        /// Implements broadcasting and performs the specified binary operation.
        /// </summary>
        /// <param name="other">The other tensor.</param>
        /// <param name="operation">The binary operation.</param>
        /// <param name="resultShape">Output parameter for the result shape.</param>
        /// <param name="resultStrides">Output parameter for the result strides.</param>
        /// <returns>New tensor after the operation.</returns>
        private Tensor<T> BroadcastOperation(Tensor<T> other, Func<T, T, T> operation, out int[] resultShape, out int[] resultStrides)
        {
            resultShape = DetermineBroadcastShape(this.shape, other.shape);
            resultStrides = CalculateStrides(resultShape);

            Tensor<T> result = new Tensor<T>(resultShape);

            // Precompute broadcasted strides
            var aStridesBroadcasted = ComputeBroadcastStrides(this.shape, resultShape, this.strides);
            var bStridesBroadcasted = ComputeBroadcastStrides(other.shape, resultShape, other.strides);

            Parallel.For(0, result.TotalSize, i =>
            {
                int[] indices = UnflattenIndex(i, resultShape);
                int aIndex = ComputeFlatIndex(indices, aStridesBroadcasted);
                int bIndex = ComputeFlatIndex(indices, bStridesBroadcasted);
                result.sharedData.Data[i] = operation(this.sharedData.Data[aIndex], other.sharedData.Data[bIndex]);
            });

            return result;
        }

        /// <summary>
        /// Computes the flat index from multi-dimensional indices and strides.
        /// </summary>
        /// <param name="indices">Multi-dimensional indices.</param>
        /// <param name="strides">Strides array.</param>
        /// <returns>Flat index.</returns>
        private static int ComputeFlatIndex(int[] indices, int[] strides)
        {
            int flatIndex = 0;
            for (int i = 0; i < indices.Length; i++)
            {
                flatIndex += indices[i] * strides[i];
            }
            return flatIndex;
        }

        /// <summary>
        /// Determines the broadcasted shape from two shapes.
        /// </summary>
        /// <param name="aShape">Shape of tensor a.</param>
        /// <param name="bShape">Shape of tensor b.</param>
        /// <returns>Broadcasted shape.</returns>
        private static int[] DetermineBroadcastShape(int[] aShape, int[] bShape)
        {
            int maxLength = Math.Max(aShape.Length, bShape.Length);
            int[] broadcastShape = new int[maxLength];

            for (int i = 0; i < maxLength; i++)
            {
                int aDim = i < aShape.Length ? aShape[aShape.Length - 1 - i] : 1;
                int bDim = i < bShape.Length ? bShape[bShape.Length - 1 - i] : 1;

                if (aDim != bDim && aDim != 1 && bDim != 1)
                    throw new InvalidOperationException($"Shapes {string.Join(",", aShape)} and {string.Join(",", bShape)} are not broadcastable.");

                broadcastShape[maxLength - 1 - i] = Math.Max(aDim, bDim);
            }

            return broadcastShape;
        }

        /// <summary>
        /// Computes broadcasted strides for a tensor based on the result shape.
        /// </summary>
        /// <param name="originalShape">Original shape of the tensor.</param>
        /// <param name="resultShape">Broadcasted shape.</param>
        /// <param name="originalStrides">Original strides of the tensor.</param>
        /// <returns>Broadcasted strides.</returns>
        private static int[] ComputeBroadcastStrides(int[] originalShape, int[] resultShape, int[] originalStrides)
        {
            int originalLength = originalShape.Length;
            int resultLength = resultShape.Length;
            int[] broadcastStrides = new int[resultLength];

            for (int i = 0; i < resultLength; i++)
            {
                int a = i >= resultLength - originalLength ? originalShape[i - (resultLength - originalLength)] : 1;
                if (a == 1)
                {
                    broadcastStrides[i] = 0;
                }
                else
                {
                    broadcastStrides[i] = originalStrides.Length >= i - (resultLength - originalLength) && i - (resultLength - originalLength) >= 0
                        ? originalStrides[i - (resultLength - originalLength)]
                        : 0;
                }
            }

            return broadcastStrides;
        }

        /// <summary>
        /// Checks if the tensor has a supported type for SIMD operations.
        /// </summary>
        /// <returns>True if supported; otherwise, false.</returns>
        private bool IsSupportedSIMD()
        {
            return typeof(T) == typeof(float) || typeof(T) == typeof(double);
        }

        /// <summary>
        /// Performs a concurrent operation across tensors using SIMD where possible.
        /// </summary>
        /// <param name="operation">The binary operation to perform.</param>
        /// <returns>A new tensor after the operation.</returns>
        private Tensor<T> ConcurrentOperation(Func<Vector<T>, Vector<T>> operation)
        {
            EnsureNotDisposed();

            Tensor<T> result = new Tensor<T>(shape);

            int vectorSize = Vector<T>.Count;
            int simdLength = (TotalSize / vectorSize) * vectorSize;

            Parallel.For(0, simdLength, i =>
            {
                Vector<T> vec = new Vector<T>(sharedData.Data, offset + i);
                Vector<T> resVec = operation(vec);
                resVec.CopyTo(result.sharedData.Data, i);
            });

            for (int i = simdLength; i < TotalSize; i++)
            {
                // Handle remaining elements if any
                result.sharedData.Data[i] = default;
            }

            return result;
        }

        #region Visualization and Serialization

        /// <summary>
        /// Serializes the tensor to a JSON string.
        /// </summary>
        /// <returns>JSON representation of the tensor.</returns>
        public string ToJson()
        {
            EnsureNotDisposed();
            return JsonSerializer.Serialize(new
            {
                Shape = this.shape,
                Data = this.ToArray()
            });
        }

        /// <summary>
        /// Deserializes a JSON string to create a tensor.
        /// </summary>
        /// <param name="json">JSON string representing a tensor.</param>
        /// <returns>Deserialized tensor.</returns>
        public static Tensor<T> FromJson(string json)
        {
            if (string.IsNullOrEmpty(json))
                throw new ArgumentException("JSON string cannot be null or empty.", nameof(json));

            var deserialized = JsonSerializer.Deserialize<SerializedTensor<T>>(json);
            if (deserialized == null)
                throw new InvalidOperationException("Failed to deserialize tensor.");

            Tensor<T> tensor = new Tensor<T>(deserialized.Shape);
            Array.Copy(deserialized.Data, 0, tensor.sharedData.Data, 0, tensor.TotalSize);
            return tensor;
        }

        /// <summary>
        /// Saves the tensor as an image (only for 2D tensors with float or double types).
        /// </summary>
        /// <param name="path">Path to save the image.</param>
        public void SaveAsImage(string path)
        {
            EnsureNotDisposed();

            if (shape.Length != 2)
                throw new InvalidOperationException("SaveAsImage is only supported for 2D tensors.");

            if (typeof(T) != typeof(float) && typeof(T) != typeof(double))
                throw new NotSupportedException("SaveAsImage is only supported for float and double types.");

            int height = shape[0];
            int width = shape[1];
            byte[] imageData = new byte[height * width * 3]; // RGB

            for (int i = 0; i < height; i++)
            {
                for (int j = 0; j < width; j++)
                {
                    T value = GetElement(i, j);
                    byte intensity = (byte)Math.Clamp(Convert.ToInt32((object)value), 0, 255);
                    int index = (i * width + j) * 3;
                    imageData[index] = intensity;     // R
                    imageData[index + 1] = intensity; // G
                    imageData[index + 2] = intensity; // B
                }
            }

            // Save as BMP
            ImageHelper.SaveAsBmp(path, width, height, imageData);
        }

        /// <summary>
        /// Helper class for image operations.
        /// </summary>
        private static class ImageHelper
        {
            public static void SaveAsBmp(string path, int width, int height, byte[] rgbData)
            {
                using var fs = new System.IO.FileStream(path, System.IO.FileMode.Create);
                using var bw = new System.IO.BinaryWriter(fs);

                int padding = (4 - (width * 3) % 4) % 4;
                int stride = width * 3 + padding;
                int fileSize = 14 + 40 + stride * height;

                // BITMAPFILEHEADER
                bw.Write((byte)'B');
                bw.Write((byte)'M');
                bw.Write(fileSize);
                bw.Write((short)0);
                bw.Write((short)0);
                bw.Write(14 + 40);

                // BITMAPINFOHEADER
                bw.Write(40);
                bw.Write(width);
                bw.Write(height);
                bw.Write((short)1);
                bw.Write((short)24);
                bw.Write(0); // Compression
                bw.Write(stride * height);
                bw.Write(0); // XPelsPerMeter
                bw.Write(0); // YPelsPerMeter
                bw.Write(0); // ClrUsed
                bw.Write(0); // ClrImportant

                // Pixel Data (Bottom-up)
                for (int i = height - 1; i >= 0; i--)
                {
                    int rowStart = i * width * 3;
                    for (int j = 0; j < width * 3; j++)
                    {
                        bw.Write(rgbData[rowStart + j]);
                    }

                    // Padding
                    for (int p = 0; p < padding; p++)
                        bw.Write((byte)0);
                }
            }
        }

        /// <summary>
        /// Internal class for serialization.
        /// </summary>
        private class SerializedTensor<TData> where TData : struct
        {
            public int[] Shape { get; set; }
            public TData[] Data { get; set; }
        }

        #endregion

        #region Serialization

        /// <summary>
        /// Serializes the tensor to a binary format.
        /// </summary>
        /// <param name="path">Path to save the binary file.</param>
        public void SaveToBinary(string path)
        {
            EnsureNotDisposed();

            using var fs = new System.IO.FileStream(path, System.IO.FileMode.Create);
            using var bw = new System.IO.BinaryWriter(fs);

            // Write shape
            bw.Write(shape.Length);
            foreach (var dim in shape)
            {
                bw.Write(dim);
            }

            // Write data
            for (int i = 0; i < TotalSize; i++)
            {
                var value = sharedData.Data[offset + i];
                WriteValue(bw, value);
            }
        }

        /// <summary>
        /// Deserializes a tensor from a binary file.
        /// </summary>
        /// <param name="path">Path to the binary file.</param>
        /// <returns>Deserialized tensor.</returns>
        public static Tensor<T> LoadFromBinary(string path)
        {
            if (string.IsNullOrEmpty(path))
                throw new ArgumentException("Path cannot be null or empty.", nameof(path));

            using var fs = new System.IO.FileStream(path, System.IO.FileMode.Open);
            using var br = new System.IO.BinaryReader(fs);

            // Read shape
            int rank = br.ReadInt32();
            int[] shape = new int[rank];
            for (int i = 0; i < rank; i++)
            {
                shape[i] = br.ReadInt32();
            }

            Tensor<T> tensor = new Tensor<T>(shape);

            // Read data
            for (int i = 0; i < tensor.TotalSize; i++)
            {
                tensor.sharedData.Data[i] = ReadValue(br);
            }

            return tensor;
        }

        /// <summary>
        /// Writes a single value to the binary writer.
        /// </summary>
        /// <param name="bw">BinaryWriter instance.</param>
        /// <param name="value">Value to write.</param>
        private static void WriteValue(System.IO.BinaryWriter bw, T value)
        {
            if (typeof(T) == typeof(float))
            {
                bw.Write((float)(object)value);
            }
            else if (typeof(T) == typeof(double))
            {
                bw.Write((double)(object)value);
            }
            else if (typeof(T) == typeof(decimal))
            {
                bw.Write((decimal)(object)value);
            }
            else
            {
                throw new NotSupportedException($"Type {typeof(T)} is not supported for binary serialization.");
            }
        }

        /// <summary>
        /// Reads a single value from the binary reader.
        /// </summary>
        /// <param name="br">BinaryReader instance.</param>
        /// <returns>Deserialized value.</returns>
        private static T ReadValue(System.IO.BinaryReader br)
        {
            if (typeof(T) == typeof(float))
            {
                return (T)(object)br.ReadSingle();
            }
            else if (typeof(T) == typeof(double))
            {
                return (T)(object)br.ReadDouble();
            }
            else if (typeof(T) == typeof(decimal))
            {
                return (T)(object)br.ReadDecimal();
            }
            else
            {
                throw new NotSupportedException($"Type {typeof(T)} is not supported for binary deserialization.");
            }
        }

        #endregion

        #region Broadcasting and Operations

        /// <summary>
        /// Determines if two tensors are broadcastable and performs the operation.
        /// </summary>
        /// <param name="other">The other tensor.</param>
        /// <param name="operation">The binary operation.</param>
        /// <returns>A new tensor after the operation.</returns>
        public Tensor<T> ElementWiseOperation(Tensor<T> other, Func<T, T, T> operation)
        {
            EnsureNotDisposed();
            other?.EnsureNotDisposed();

            if (shape.Length == other.shape.Length && shape.SequenceEqual(other.shape))
            {
                // Shapes match; perform operation directly
                Tensor<T> result = new Tensor<T>(shape);
                if (TotalSize >= ParallelizationThreshold)
                {
                    Parallel.For(0, TotalSize, i =>
                    {
                        result.sharedData.Data[i] = operation(sharedData.Data[offset + i], other.sharedData.Data[other.offset + i]);
                    });
                }
                else
                {
                    for (int i = 0; i < TotalSize; i++)
                    {
                        result.sharedData.Data[i] = operation(sharedData.Data[offset + i], other.sharedData.Data[other.offset + i]);
                    }
                }

                return result;
            }
            else
            {
                // Perform broadcasting
                return this.BroadcastOperation(other, operation, out _, out _);
            }
        }

        /// <summary>
        /// Broadcasts tensors to a common shape and performs the specified operation.
        /// </summary>
        /// <param name="other">The other tensor.</param>
        /// <param name="operation">The binary operation.</param>
        /// <param name="resultShape">Output shape of the result tensor.</param>
        /// <param name="resultStrides">Output strides of the result tensor.</param>
        /// <returns>A new tensor after broadcasting and operation.</returns>
        private Tensor<T> BroadcastOperation(Tensor<T> other, Func<T, T, T> operation, out int[] resultShape, out int[] resultStrides)
        {
            resultShape = DetermineBroadcastShape(this.shape, other.shape);
            resultStrides = CalculateStrides(resultShape);

            Tensor<T> result = new Tensor<T>(resultShape);

            // Precompute broadcasted strides
            var aStridesBroadcasted = ComputeBroadcastStrides(this.shape, resultShape, this.strides);
            var bStridesBroadcasted = ComputeBroadcastStrides(other.shape, resultShape, other.strides);

            Parallel.For(0, result.TotalSize, i =>
            {
                int[] indices = UnflattenIndex(i, resultShape);
                int aIndex = ComputeFlatIndex(indices, aStridesBroadcasted);
                int bIndex = ComputeFlatIndex(indices, bStridesBroadcasted);
                result.sharedData.Data[i] = operation(this.sharedData.Data[aIndex], other.sharedData.Data[bIndex]);
            });

            return result;
        }

        #endregion

        #region Tensor Manipulation Methods

        /// <summary>
        /// Transposes the tensor along two specified axes.
        /// </summary>
        /// <param name="axis1">First axis.</param>
        /// <param name="axis2">Second axis.</param>
        /// <returns>A new tensor with transposed axes.</returns>
        public Tensor<T> Transpose(int axis1, int axis2)
        {
            EnsureNotDisposed();

            if (axis1 < 0 || axis1 >= shape.Length)
                throw new ArgumentOutOfRangeException(nameof(axis1), "Axis1 is out of bounds.");

            if (axis2 < 0 || axis2 >= shape.Length)
                throw new ArgumentOutOfRangeException(nameof(axis2), "Axis2 is out of bounds.");

            if (axis1 == axis2)
                return this; // No change needed

            int[] newShape = (int[])shape.Clone();
            (newShape[axis1], newShape[axis2]) = (newShape[axis2], newShape[axis1]);

            int[] newStrides = CalculateStrides(newShape);

            Tensor<T> result = new Tensor<T>(sharedData, newShape, newStrides, offset);

            Parallel.For(0, TotalSize, i =>
            {
                int[] indices = UnflattenIndex(i, shape);
                (indices[axis1], indices[axis2]) = (indices[axis2], indices[axis1]);
                int newFlatIndex = result.CalculateFlatIndex(indices);
                result.sharedData.Data[newFlatIndex] = sharedData.Data[offset + i];
            });

            return result;
        }

        /// <summary>
        /// Applies an activation function element-wise.
        /// </summary>
        /// <param name="activation">The activation function to apply.</param>
        /// <returns>A new tensor after applying the activation.</returns>
        public Tensor<T> ApplyActivation(Func<T, T> activation)
        {
            EnsureNotDisposed();

            Tensor<T> result = new Tensor<T>(shape);

            if (TotalSize >= ParallelizationThreshold)
            {
                Parallel.For(0, TotalSize, i =>
                {
                    result.sharedData.Data[i] = activation(sharedData.Data[offset + i]);
                });
            }
            else
            {
                for (int i = 0; i < TotalSize; i++)
                {
                    result.sharedData.Data[i] = activation(sharedData.Data[offset + i]);
                }
            }

            return result;
        }

        /// <summary>
        /// Applies a batch normalization operation.
        /// </summary>
        /// <param name="gamma">Scale parameter.</param>
        /// <param name="beta">Shift parameter.</param>
        /// <param name="mean">Mean for normalization.</param>
        /// <param name="variance">Variance for normalization.</param>
        /// <param name="epsilon">Epsilon to prevent division by zero.</param>
        /// <returns>A new tensor after batch normalization.</returns>
        public Tensor<T> BatchNormalize(T gamma, T beta, Tensor<T> mean, Tensor<T> variance, T epsilon = default)
        {
            EnsureNotDisposed();
            mean?.EnsureNotDisposed();
            variance?.EnsureNotDisposed();

            if (mean.shape.Length != this.shape.Length || variance.shape.Length != this.shape.Length)
                throw new ArgumentException("Mean and variance must have the same shape as the tensor.");

            Tensor<T> result = new Tensor<T>(shape);

            if (TotalSize >= ParallelizationThreshold)
            {
                Parallel.For(0, TotalSize, i =>
                {
                    T std = INumber<T>.Sqrt(variance.sharedData.Data[variance.offset + i] + epsilon);
                    result.sharedData.Data[i] = (sharedData.Data[offset + i] - mean.sharedData.Data[mean.offset + i]) / std * gamma + beta;
                });
            }
            else
            {
                for (int i = 0; i < TotalSize; i++)
                {
                    T std = INumber<T>.Sqrt(variance.sharedData.Data[variance.offset + i] + epsilon);
                    result.sharedData.Data[i] = (sharedData.Data[offset + i] - mean.sharedData.Data[mean.offset + i]) / std * gamma + beta;
                }
            }

            return result;
        }

        #endregion

        #region Automatic Differentiation (Autograd)

        // Placeholder for automatic differentiation implementation
        // Implement reverse-mode automatic differentiation as needed.

        #endregion

        #region Tensor Broadcasting

        /// <summary>
        /// Implements NumPy-like tensor broadcasting.
        /// </summary>
        /// <param name="other">The other tensor.</param>
        /// <param name="operation">The binary operation to perform.</param>
        /// <returns>A new tensor after broadcasting and operation.</returns>
        public Tensor<T> BroadcastOperation(Tensor<T> other, Func<T, T, T> operation)
        {
            EnsureNotDisposed();
            other?.EnsureNotDisposed();

            int[] resultShape = DetermineBroadcastShape(this.shape, other.shape);
            int[] resultStrides = CalculateStrides(resultShape);

            Tensor<T> result = new Tensor<T>(resultShape);

            // Compute strides for both tensors with respect to the result shape
            int[] aBroadcastStrides = ComputeBroadcastStrides(this.shape, resultShape, this.strides);
            int[] bBroadcastStrides = ComputeBroadcastStrides(other.shape, resultShape, other.strides);

            Parallel.For(0, result.TotalSize, i =>
            {
                int[] indices = UnflattenIndex(i, resultShape);
                int aFlatIndex = ComputeFlatIndex(indices, aBroadcastStrides);
                int bFlatIndex = ComputeFlatIndex(indices, bBroadcastStrides);
                result.sharedData.Data[i] = operation(this.sharedData.Data[aFlatIndex], other.sharedData.Data[bFlatIndex]);
            });

            return result;
        }

        #endregion

        #region Helper Methods

        /// <summary>
        /// Performs a chunked concurrent operation with SIMDed operations.
        /// </summary>
        /// <param name="operation">The binary operation to perform.</param>
        /// <returns>A new tensor after the operation.</returns>
        private Tensor<T> ConcurrentOperation(Func<Vector<T>, Vector<T>> operation)
        {
            EnsureNotDisposed();

            Tensor<T> result = new Tensor<T>(shape);

            int vectorSize = Vector<T>.Count;
            int simdLength = (TotalSize / vectorSize) * vectorSize;

            Parallel.For(0, simdLength, i =>
            {
                Vector<T> vec = new Vector<T>(sharedData.Data, offset + i);
                Vector<T> resVec = operation(vec);
                resVec.CopyTo(result.sharedData.Data, i);
            });

            for (int i = simdLength; i < TotalSize; i++)
            {
                result.sharedData.Data[i] = default;
            }

            return result;
        }

        /// <summary>
        /// Converts a flat index to multi-dimensional indices based on the shape.
        /// </summary>
        /// <param name="flatIndex">Flat index.</param>
        /// <param name="shape">Shape of the tensor.</param>
        /// <returns>Array of multi-dimensional indices.</returns>
        private static int[] UnflattenIndex(int flatIndex, int[] shape)
        {
            int[] indices = new int[shape.Length];
            for (int i = shape.Length - 1; i >= 0; i--)
            {
                indices[i] = flatIndex % shape[i];
                flatIndex /= shape[i];
            }
            return indices;
        }


        private void EnsureNotDisposed()
        {
            if (disposed)
                throw new ObjectDisposedException(nameof(Tensor<T>));
        }

        /// <summary>
        /// Dispose pattern implementation.
        /// </summary>
        public void Dispose()
        {
            if (!disposed)
            {
                if (isView)
                {
                    sharedData.Release();
                }
                else
                {
                    sharedData.Release();
                }
                disposed = true;
            }
        }

        #endregion

        #region ToArray and Display

        /// <summary>
        /// Converts the tensor data to a 1D array.
        /// </summary>
        /// <returns>1D array containing tensor data.</returns>
        public T[] ToArray()
        {
            EnsureNotDisposed();

            T[] array = new T[TotalSize];
            Array.Copy(sharedData.Data, offset, array, 0, TotalSize);
            return array;
        }

        /// <summary>
        /// Displays the tensor data (for debugging purposes).
        /// </summary>
        public void Display()
        {
            EnsureNotDisposed();

            if (shape.Length == 1)
            {
                Console.Write("[");
                for (int i = 0; i < shape[0]; i++)
                {
                    Console.Write(sharedData.Data[offset + i]);
                    if (i < shape[0] - 1)
                        Console.Write(", ");
                }
                Console.WriteLine("]");
            }
            else if (shape.Length == 2)
            {
                for (int i = 0; i < shape[0]; i++)
                {
                    Console.Write("[");
                    for (int j = 0; j < shape[1]; j++)
                    {
                        Console.Write(sharedData.Data[offset + i * strides[0] + j]);
                        if (j < shape[1] - 1)
                            Console.Write(", ");
                    }
                    Console.WriteLine("]");
                }
            }
            else
            {
                Console.WriteLine("Tensor data display is only implemented for 1D and 2D tensors.");
            }
        }

        #endregion
    }
}