namespace SynorCompute;

/// <summary>
/// Multi-dimensional tensor for compute operations.
/// </summary>
/// <example>
/// <code>
/// // Create a 2D tensor
/// var matrix = new Tensor(new[] { 2, 3 }, new[] { 1.0, 2.0, 3.0, 4.0, 5.0, 6.0 });
///
/// // Create random tensor
/// var random = Tensor.Rand(512, 512);
///
/// // Operations
/// var mean = random.Mean();
/// var transposed = matrix.Transpose();
/// </code>
/// </example>
public class Tensor : IEquatable<Tensor>
{
    public int[] Shape { get; }
    public double[] Data { get; }
    public Precision Dtype { get; }

    public int Size => Data.Length;
    public int Ndim => Shape.Length;

    public Tensor(int[] shape, double[] data, Precision dtype = Precision.Fp32)
    {
        int expectedSize = shape.Aggregate(1, (a, b) => a * b);
        if (data.Length != expectedSize)
        {
            throw new ArgumentException($"Data size {data.Length} does not match shape [{string.Join(", ", shape)}]");
        }

        Shape = (int[])shape.Clone();
        Data = (double[])data.Clone();
        Dtype = dtype;
    }

    /// <summary>
    /// Get element at indices.
    /// </summary>
    public double this[params int[] indices]
    {
        get
        {
            if (indices.Length != Shape.Length)
            {
                throw new ArgumentException("Index dimensions must match tensor dimensions");
            }

            int idx = 0;
            int stride = 1;
            for (int i = Shape.Length - 1; i >= 0; i--)
            {
                idx += indices[i] * stride;
                stride *= Shape[i];
            }
            return Data[idx];
        }
    }

    // Factory Methods

    public static Tensor Of(double[] data) => new([data.Length], data);

    public static Tensor Of(double[,] data)
    {
        int rows = data.GetLength(0);
        int cols = data.GetLength(1);
        var flat = new double[rows * cols];
        for (int i = 0; i < rows; i++)
        {
            for (int j = 0; j < cols; j++)
            {
                flat[i * cols + j] = data[i, j];
            }
        }
        return new Tensor([rows, cols], flat);
    }

    public static Tensor Zeros(params int[] shape)
    {
        int size = shape.Aggregate(1, (a, b) => a * b);
        return new Tensor(shape, new double[size]);
    }

    public static Tensor Ones(params int[] shape)
    {
        int size = shape.Aggregate(1, (a, b) => a * b);
        var data = new double[size];
        Array.Fill(data, 1.0);
        return new Tensor(shape, data);
    }

    public static Tensor Rand(params int[] shape)
    {
        int size = shape.Aggregate(1, (a, b) => a * b);
        var random = new Random();
        var data = new double[size];
        for (int i = 0; i < size; i++)
        {
            data[i] = random.NextDouble();
        }
        return new Tensor(shape, data);
    }

    public static Tensor Randn(params int[] shape)
    {
        int size = shape.Aggregate(1, (a, b) => a * b);
        var random = new Random();
        var data = new double[size];
        for (int i = 0; i < size; i++)
        {
            // Box-Muller transform
            double u1 = random.NextDouble();
            double u2 = random.NextDouble();
            data[i] = Math.Sqrt(-2 * Math.Log(u1)) * Math.Cos(2 * Math.PI * u2);
        }
        return new Tensor(shape, data);
    }

    public static Tensor Eye(int n)
    {
        var data = new double[n * n];
        for (int i = 0; i < n; i++)
        {
            data[i * n + i] = 1.0;
        }
        return new Tensor([n, n], data);
    }

    public static Tensor Arange(double start, double end, double step = 1.0)
    {
        int size = (int)Math.Ceiling((end - start) / step);
        var data = new double[size];
        for (int i = 0; i < size; i++)
        {
            data[i] = start + i * step;
        }
        return new Tensor([size], data);
    }

    public static Tensor Linspace(double start, double end, int num)
    {
        var data = new double[num];
        double step = (end - start) / (num - 1);
        for (int i = 0; i < num; i++)
        {
            data[i] = start + i * step;
        }
        return new Tensor([num], data);
    }

    // Operations

    public Tensor Reshape(params int[] newShape)
    {
        int newSize = newShape.Aggregate(1, (a, b) => a * b);
        if (newSize != Size)
        {
            throw new ArgumentException($"Cannot reshape tensor of size {Size} to shape [{string.Join(", ", newShape)}]");
        }
        return new Tensor(newShape, Data, Dtype);
    }

    public Tensor Transpose()
    {
        if (Ndim != 2)
        {
            throw new InvalidOperationException("Transpose only supported for 2D tensors");
        }

        int rows = Shape[0];
        int cols = Shape[1];
        var transposed = new double[Data.Length];

        for (int i = 0; i < rows; i++)
        {
            for (int j = 0; j < cols; j++)
            {
                transposed[j * rows + i] = Data[i * cols + j];
            }
        }

        return new Tensor([cols, rows], transposed, Dtype);
    }

    // Reductions

    public double Mean() => Data.Average();
    public double Sum() => Data.Sum();

    public double Std()
    {
        double mean = Mean();
        double sumSq = Data.Sum(x => (x - mean) * (x - mean));
        return Math.Sqrt(sumSq / Data.Length);
    }

    public double Max() => Data.Max();
    public double Min() => Data.Min();

    // Activations

    public Tensor Relu() => new(Shape, Data.Select(x => Math.Max(0, x)).ToArray(), Dtype);

    public Tensor Sigmoid() => new(Shape, Data.Select(x => 1.0 / (1.0 + Math.Exp(-x))).ToArray(), Dtype);

    public Tensor Softmax()
    {
        double maxVal = Max();
        var expValues = Data.Select(x => Math.Exp(x - maxVal)).ToArray();
        double sum = expValues.Sum();
        return new Tensor(Shape, expValues.Select(x => x / sum).ToArray(), Dtype);
    }

    // Conversion

    public object ToNestedList()
    {
        return Ndim switch
        {
            1 => Data.ToList(),
            2 => Enumerable.Range(0, Shape[0])
                .Select(i => Enumerable.Range(0, Shape[1])
                    .Select(j => Data[i * Shape[1] + j])
                    .ToList())
                .ToList(),
            _ => throw new InvalidOperationException("ToNestedList only supports 1D and 2D tensors")
        };
    }

    // Equality

    public bool Equals(Tensor? other)
    {
        if (other is null) return false;
        if (ReferenceEquals(this, other)) return true;
        return Shape.SequenceEqual(other.Shape) &&
               Data.SequenceEqual(other.Data) &&
               Dtype == other.Dtype;
    }

    public override bool Equals(object? obj) => Equals(obj as Tensor);

    public override int GetHashCode() => HashCode.Combine(
        Shape.Aggregate(0, HashCode.Combine),
        Data.Aggregate(0, (acc, val) => HashCode.Combine(acc, val.GetHashCode())),
        Dtype
    );

    public override string ToString() => $"Tensor(shape=[{string.Join(", ", Shape)}], dtype={Dtype})";
}