package holeg.power_flow;

import Jama.Matrix;

import java.io.PrintStream;
import java.util.Arrays;

public class ComplexMatrix {
    private final int width;
    private final int height;
    private final ComplexNumber[] data;

    public ComplexMatrix(int size) {
        this(size, size);
    }

    public ComplexMatrix(int width, int height) {
        if (width <= 0)
            throw new IllegalArgumentException("width is invalid");
        if (height <= 0)
            throw new IllegalArgumentException("height is invalid");

        this.width = width;
        this.height = height;
        this.data = new ComplexNumber[width * height];
        Arrays.fill(this.data, ComplexNumber.Zero);
    }

    public static ComplexMatrix sameSizeAs(ComplexMatrix other) {
        return new ComplexMatrix(other.getWidth(), other.getHeight());
    }

    public static ComplexMatrix fromArray(ComplexNumber[] numbers) {
        if (numbers == null)
            throw new IllegalArgumentException("number is null");
        if (numbers.length == 0)
            throw new IllegalArgumentException("number is empty");

        ComplexMatrix out = new ComplexMatrix(numbers.length, 1);
        for (int i = 0; i < numbers.length; i++)
            out.uncheckedSet(i, 0, numbers[i]);
        return out;
    }

    public int getWidth() {
        return width;
    }

    public int getHeight() {
        return height;
    }

    private ComplexNumber uncheckedGet(int row, int column) {
        return data[column * width + row];
    }

    public ComplexNumber get(int row, int column) {
        if (row < 0 || row >= width)
            throw new IllegalArgumentException("row is invalid");
        if (column < 0 || column >= height)
            throw new IllegalArgumentException("column is invalid");
        return uncheckedGet(row, column);
    }

    public void set(int row, int column, double realValue) {
        set(row, column, new ComplexNumber(realValue, 0));
    }

    private void uncheckedSet(int row, int column, ComplexNumber value) {
        data[column * width + row] = value;
    }

    public void set(int row, int column, ComplexNumber value) {
        if (row < 0 || row >= width)
            throw new IllegalArgumentException("row is invalid");
        if (column < 0 || column >= height)
            throw new IllegalArgumentException("column is invalid");
        if (value == null)
            throw new IllegalArgumentException("value is null");
        uncheckedSet(row, column, value);
    }

    public Matrix getRealMatrix() {
        Matrix m = new Matrix(getWidth(), getHeight());
        for (int i = 0; i < getWidth(); i++)
            for (int j = 0; j < getHeight(); j++)
                m.set(i, j, uncheckedGet(i, j).real);
        return m;
    }

    public Matrix getImaginaryMatrix() {
        Matrix m = new Matrix(getWidth(), getHeight());
        for (int i = 0; i < getWidth(); i++)
            for (int j = 0; j < getHeight(); j++)
                m.set(i, j, uncheckedGet(i, j).imaginary);
        return m;
    }

    public ComplexMatrix add(ComplexMatrix other) {
        if (getWidth() != other.getWidth())
            throw new IllegalArgumentException("width does not match");
        if (getHeight() != other.getHeight())
            throw new IllegalArgumentException("height does not match");
        ComplexMatrix out = new ComplexMatrix(getWidth(), getHeight());
        for (int i = 0; i < getWidth(); i++)
            for (int j = 0; j < getHeight(); j++)
                out.uncheckedSet(i, j, uncheckedGet(i, j).add(other.uncheckedGet(i, j)));
        return out;
    }

    public ComplexMatrix subtract(ComplexMatrix other) {
        if (getWidth() != other.getWidth())
            throw new IllegalArgumentException("width does not match");
        if (getHeight() != other.getHeight())
            throw new IllegalArgumentException("height does not match");
        ComplexMatrix out = new ComplexMatrix(getWidth(), getHeight());
        for (int i = 0; i < getWidth(); i++)
            for (int j = 0; j < getHeight(); j++)
                out.uncheckedSet(i, j, uncheckedGet(i, j).subtract(other.uncheckedGet(i, j)));
        return out;
    }

    public ComplexMatrix multiply(ComplexNumber[] numbers) {
        return multiply(ComplexMatrix.fromArray(numbers));
    }

    public ComplexMatrix multiply(ComplexMatrix other) {
        if (getHeight() != other.getWidth())
            throw new IllegalArgumentException("height does not match other width");

        ComplexMatrix out = new ComplexMatrix(getWidth(), other.getHeight());
        for (int i = 0; i < getWidth(); i++) {
            for (int j = 0; j < other.getHeight(); j++) {
                ComplexNumber sum = ComplexNumber.Zero;

                for (int k = 0; k < getHeight(); k++) {
                    ComplexNumber a = uncheckedGet(i, k);
                    ComplexNumber b = other.uncheckedGet(k, j);
                    sum = sum.add(a.multiply(b));
                }
                out.uncheckedSet(i, j, sum);
            }
        }
        return out;
    }

    public void printTo(PrintStream stream) {
        if (stream == null)
            throw new IllegalArgumentException("stream is null");

        for (int row = 0; row < getWidth(); row++) {
            for (int col = 0; col < getHeight(); col++)
                stream.printf("%s, ", uncheckedGet(row, col).toString());
            stream.println();
        }
    }
}