package io.synor.compute;

import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;

import java.util.*;

/**
 * Multi-dimensional tensor for compute operations.
 *
 * <pre>{@code
 * // Create a 2D tensor
 * Tensor matrix = Tensor.of(new double[][]{{1, 2, 3}, {4, 5, 6}});
 *
 * // Create a random tensor
 * Tensor random = Tensor.rand(512, 512);
 *
 * // Get shape
 * int[] shape = matrix.getShape(); // [2, 3]
 * }</pre>
 */
public class Tensor {
    @JsonProperty("shape")
    private final int[] shape;

    @JsonProperty("data")
    private final double[] data;

    @JsonProperty("dtype")
    private final Precision dtype;

    /**
     * Create a tensor with given shape and data.
     */
    public Tensor(int[] shape, double[] data, Precision dtype) {
        this.shape = shape.clone();
        this.data = data.clone();
        this.dtype = dtype;

        int expectedSize = 1;
        for (int dim : shape) {
            expectedSize *= dim;
        }
        if (data.length != expectedSize) {
            throw new IllegalArgumentException(
                "Data size " + data.length + " does not match shape " + Arrays.toString(shape));
        }
    }

    /**
     * Create a tensor with FP32 precision.
     */
    public Tensor(int[] shape, double[] data) {
        this(shape, data, Precision.FP32);
    }

    public int[] getShape() {
        return shape.clone();
    }

    public double[] getData() {
        return data.clone();
    }

    public Precision getDtype() {
        return dtype;
    }

    @JsonIgnore
    public int getSize() {
        int size = 1;
        for (int dim : shape) {
            size *= dim;
        }
        return size;
    }

    @JsonIgnore
    public int getNdim() {
        return shape.length;
    }

    /**
     * Create a tensor from a 1D array.
     */
    public static Tensor of(double[] data) {
        return new Tensor(new int[]{data.length}, data);
    }

    /**
     * Create a tensor from a 2D array.
     */
    public static Tensor of(double[][] data) {
        int rows = data.length;
        int cols = data[0].length;
        double[] flat = new double[rows * cols];
        for (int i = 0; i < rows; i++) {
            System.arraycopy(data[i], 0, flat, i * cols, cols);
        }
        return new Tensor(new int[]{rows, cols}, flat);
    }

    /**
     * Create a tensor filled with zeros.
     */
    public static Tensor zeros(int... shape) {
        int size = 1;
        for (int dim : shape) {
            size *= dim;
        }
        return new Tensor(shape, new double[size]);
    }

    /**
     * Create a tensor filled with ones.
     */
    public static Tensor ones(int... shape) {
        int size = 1;
        for (int dim : shape) {
            size *= dim;
        }
        double[] data = new double[size];
        Arrays.fill(data, 1.0);
        return new Tensor(shape, data);
    }

    /**
     * Create a tensor with uniform random values [0, 1).
     */
    public static Tensor rand(int... shape) {
        Random random = new Random();
        int size = 1;
        for (int dim : shape) {
            size *= dim;
        }
        double[] data = new double[size];
        for (int i = 0; i < size; i++) {
            data[i] = random.nextDouble();
        }
        return new Tensor(shape, data);
    }

    /**
     * Create a tensor with standard normal random values.
     */
    public static Tensor randn(int... shape) {
        Random random = new Random();
        int size = 1;
        for (int dim : shape) {
            size *= dim;
        }
        double[] data = new double[size];
        for (int i = 0; i < size; i++) {
            data[i] = random.nextGaussian();
        }
        return new Tensor(shape, data);
    }

    /**
     * Create an identity matrix.
     */
    public static Tensor eye(int n) {
        double[] data = new double[n * n];
        for (int i = 0; i < n; i++) {
            data[i * n + i] = 1.0;
        }
        return new Tensor(new int[]{n, n}, data);
    }

    /**
     * Create a range tensor [start, end) with step.
     */
    public static Tensor arange(double start, double end, double step) {
        int size = (int) Math.ceil((end - start) / step);
        double[] data = new double[size];
        for (int i = 0; i < size; i++) {
            data[i] = start + i * step;
        }
        return new Tensor(new int[]{size}, data);
    }

    /**
     * Create a range tensor [start, end) with step 1.
     */
    public static Tensor arange(double start, double end) {
        return arange(start, end, 1.0);
    }

    /**
     * Create a linearly spaced tensor.
     */
    public static Tensor linspace(double start, double end, int num) {
        double[] 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(new int[]{num}, data);
    }

    /**
     * Get element at index.
     */
    public double get(int... indices) {
        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];
    }

    /**
     * Reshape tensor to new shape.
     */
    public Tensor reshape(int... newShape) {
        int newSize = 1;
        for (int dim : newShape) {
            newSize *= dim;
        }
        if (newSize != getSize()) {
            throw new IllegalArgumentException(
                "Cannot reshape tensor of size " + getSize() + " to shape " + Arrays.toString(newShape));
        }
        return new Tensor(newShape, data, dtype);
    }

    /**
     * Transpose 2D tensor.
     */
    public Tensor transpose() {
        if (shape.length != 2) {
            throw new IllegalStateException("Transpose only supported for 2D tensors");
        }
        int rows = shape[0];
        int cols = shape[1];
        double[] 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(new int[]{cols, rows}, transposed, dtype);
    }

    /**
     * Compute mean of all elements.
     */
    public double mean() {
        double sum = 0;
        for (double v : data) {
            sum += v;
        }
        return sum / data.length;
    }

    /**
     * Compute sum of all elements.
     */
    public double sum() {
        double sum = 0;
        for (double v : data) {
            sum += v;
        }
        return sum;
    }

    /**
     * Compute standard deviation.
     */
    public double std() {
        double mean = mean();
        double sumSq = 0;
        for (double v : data) {
            sumSq += (v - mean) * (v - mean);
        }
        return Math.sqrt(sumSq / data.length);
    }

    /**
     * Find maximum value.
     */
    public double max() {
        double max = Double.NEGATIVE_INFINITY;
        for (double v : data) {
            if (v > max) max = v;
        }
        return max;
    }

    /**
     * Find minimum value.
     */
    public double min() {
        double min = Double.POSITIVE_INFINITY;
        for (double v : data) {
            if (v < min) min = v;
        }
        return min;
    }

    /**
     * Apply ReLU activation.
     */
    public Tensor relu() {
        double[] result = new double[data.length];
        for (int i = 0; i < data.length; i++) {
            result[i] = Math.max(0, data[i]);
        }
        return new Tensor(shape, result, dtype);
    }

    /**
     * Apply sigmoid activation.
     */
    public Tensor sigmoid() {
        double[] result = new double[data.length];
        for (int i = 0; i < data.length; i++) {
            result[i] = 1.0 / (1.0 + Math.exp(-data[i]));
        }
        return new Tensor(shape, result, dtype);
    }

    /**
     * Apply softmax activation.
     */
    public Tensor softmax() {
        double max = max();
        double[] exp = new double[data.length];
        double sum = 0;
        for (int i = 0; i < data.length; i++) {
            exp[i] = Math.exp(data[i] - max);
            sum += exp[i];
        }
        for (int i = 0; i < data.length; i++) {
            exp[i] /= sum;
        }
        return new Tensor(shape, exp, dtype);
    }

    /**
     * Convert to nested list representation.
     */
    public Object toNestedList() {
        if (shape.length == 1) {
            List<Double> list = new ArrayList<>();
            for (double v : data) {
                list.add(v);
            }
            return list;
        } else if (shape.length == 2) {
            List<List<Double>> list = new ArrayList<>();
            int rows = shape[0];
            int cols = shape[1];
            for (int i = 0; i < rows; i++) {
                List<Double> row = new ArrayList<>();
                for (int j = 0; j < cols; j++) {
                    row.add(data[i * cols + j]);
                }
                list.add(row);
            }
            return list;
        }
        throw new UnsupportedOperationException("toNestedList only supports 1D and 2D tensors");
    }

    @Override
    public String toString() {
        return "Tensor(shape=" + Arrays.toString(shape) + ", dtype=" + dtype.getValue() + ")";
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        Tensor tensor = (Tensor) o;
        return Arrays.equals(shape, tensor.shape) &&
               Arrays.equals(data, tensor.data) &&
               dtype == tensor.dtype;
    }

    @Override
    public int hashCode() {
        int result = Objects.hash(dtype);
        result = 31 * result + Arrays.hashCode(shape);
        result = 31 * result + Arrays.hashCode(data);
        return result;
    }
}
