package holeg.power_flow;

import Jama.LUDecomposition;
import Jama.Matrix;

import java.util.Arrays;

public class NewtonRaphsonSolver implements Solver {
    @Override
    public SolverResult solve(PowerFlowProblem problem, SolverSettings settings) {
        if (problem == null)
            throw new IllegalArgumentException("problem is null");
        if (settings == null)
            throw new IllegalArgumentException("settings is null");

        try {
            long timeStart = System.nanoTime();

            // Check data first
            if (!problem.checkForValidData())
                return SolverResult.error(SolverError.InvalidData);

            // Create performance monitor
            SolverPerformance performance = new SolverPerformance();

            // Build admittance matrix
            AdmittanceMatrix admittanceMatrix = AdmittanceMatrix.build(problem.buses, problem.lines);
            Matrix G_bus = admittanceMatrix.getG();
            Matrix B_bus = admittanceMatrix.getB();

            // Copy data over to do math with it
            int busCount = problem.buses.length;
            double[] V = new double[busCount];
            double[] delta = new double[busCount];
            double[] Pset = new double[busCount];
            double[] Qset = new double[busCount];
            for (int i = 0; i < busCount; i++) {
                V[i] = problem.buses[i].voltage;
                delta[i] = problem.buses[i].delta;
                Pset[i] = problem.buses[i].Pg - problem.buses[i].Pl; // generation - load
                Qset[i] = problem.buses[i].Qg - problem.buses[i].Ql;
            }

            // Find PQ/PV nodes
            int[] pq = Util.indicesOf(problem.buses, BusType.PQ);
            int[] pv = Util.indicesOf(problem.buses, BusType.PV);

            // Main solver iteration
            double tolerance = Double.MAX_VALUE;

            // Jacobian caching
            Matrix inverseJ = Matrix.identity(1, 1);
            double factor = 1;
            int jacobianRecalculation = 0;

            for (int iter = 1; iter <= settings.maxIterations; iter++) {
                // Calculate current P/Q values
                double[] P = new double[busCount];
                double[] Q = new double[busCount];

                for (int i = 0; i < busCount; i++) {
                    for (int j = 0; j < busCount; j++) {
                        P[i] += V[i] * V[j] * (G_bus.get(i, j) * Math.cos(delta[i] - delta[j]) + B_bus.get(i, j) * Math.sin(delta[i] - delta[j]));
                        Q[i] += V[i] * V[j] * (G_bus.get(i, j) * Math.sin(delta[i] - delta[j]) - B_bus.get(i, j) * Math.cos(delta[i] - delta[j]));
                    }
                }

                // Reactive power limit
                if (iter > 2 && iter < 20) {
                    for (int i = 0; i < pv.length; i++) {
                        Bus bus = problem.buses[i];
                        double q = Q[i] + bus.Ql;
                        if (q < bus.Qmin)
                            V[i] += 0.01;
                        else if (q > bus.Qmax)
                            V[i] -= 0.01;
                    }
                }

                // Calculate error given by delta between set values and current values
                double[] dP = Util.subtract(Pset, P);
                double[] dQ = Util.subtract(Qset, Q);

                // Create error vector
                double[] error = new double[busCount - 1 + pq.length];
                for (int i = 0; i < busCount - 1; i++)
                    error[i] = dP[i + 1];
                for (int i = 0; i < pq.length; i++)
                    error[i + busCount - 1] = dQ[pq[i]];

                // Update performance
                performance.iterations = iter;
                performance.error = tolerance;

                // Calculate mismatch between values excepted and current values
                tolerance = Arrays.stream(error).map(Math::abs).max().orElse(0);
                if (tolerance < settings.minError)
                    break;

                if (jacobianRecalculation <= 0) {
                    // Jacobian matrix J1
                    Matrix J1 = new Matrix(busCount - 1, busCount - 1);
                    for (int i = 0; i < busCount - 1; i++) {
                        int m = i + 1;
                        for (int k = 0; k < busCount - 1; k++) {
                            int n = k + 1;
                            if (n == m) {
                                for (int l = 0; l < busCount - 1; l++)
                                    J1.set(i, k, J1.get(i, k) + V[m] * V[l] * (-G_bus.get(m, l) * Math.sin(delta[m] - delta[l]) + B_bus.get(m, l) * Math.cos(delta[m] - delta[l])));
                                J1.set(i, k, J1.get(i, k) - Math.pow(V[m], 2) * B_bus.get(m, m));
                            } else
                                J1.set(i, k, V[m] * V[n] * (G_bus.get(m, n) * Math.sin(delta[m] - delta[n]) - B_bus.get(m, n) * Math.cos(delta[m] - delta[n])));
                        }
                    }

                    // Jacobian matrix J2
                    Matrix J2 = new Matrix(busCount - 1, pq.length);
                    for (int i = 0; i < busCount - 1; i++) {
                        int m = i + 1;
                        for (int k = 0; k < pq.length; k++) {
                            int n = pq[k];
                            if (n == m) {
                                for (int l = 0; l < busCount; l++)
                                    J2.set(i, k, J2.get(i, k) + V[l] * (G_bus.get(m, l) * Math.cos(delta[m] - delta[l]) + B_bus.get(m, l) * Math.sin(delta[m] - delta[l])));
                                J2.set(i, k, J2.get(i, k) + V[m] * G_bus.get(m, m));
                            } else
                                J2.set(i, k, V[m] * (G_bus.get(m, n) * Math.cos(delta[m] - delta[n]) + B_bus.get(m, n) * Math.sin(delta[m] - delta[n])));
                        }
                    }

                    // Jacobian matrix J3
                    Matrix J3 = new Matrix(pq.length, busCount - 1);
                    for (int i = 0; i < pq.length; i++) {
                        int m = pq[i];
                        for (int k = 0; k < busCount - 1; k++) {
                            int n = k + 1;
                            if (n == m) {
                                for (int l = 0; l < busCount; l++)
                                    J3.set(i, k, J3.get(i, k) + V[m] * V[l] * (G_bus.get(m, l) * Math.cos(delta[m] - delta[l]) + B_bus.get(m, l) * Math.sin(delta[m] - delta[l])));
                                J3.set(i, k, J3.get(i, k) - Math.pow(V[m], 2) * G_bus.get(m, m));
                            } else
                                J3.set(i, k, V[m] * V[n] * (-G_bus.get(m, n) * Math.cos(delta[m] - delta[n]) - B_bus.get(m, n) * Math.sin(delta[m] - delta[n])));
                        }
                    }

                    // Jacobian matrix J4
                    Matrix J4 = new Matrix(pq.length, pq.length);
                    for (int i = 0; i < pq.length; i++) {
                        int m = pq[i];
                        for (int k = 0; k < pq.length; k++) {
                            int n = pq[k];
                            if (n == m) {
                                for (int l = 0; l < busCount; l++)
                                    J4.set(i, k, J4.get(i, k) + V[l] * (G_bus.get(m, l) * Math.sin(delta[m] - delta[l]) - B_bus.get(m, l) * Math.cos(delta[m] - delta[l])));
                                J4.set(i, k, J4.get(i, k) - V[m] * B_bus.get(m, m));
                            } else
                                J4.set(i, k, V[m] * (G_bus.get(m, n) * Math.sin(delta[m] - delta[n]) - B_bus.get(m, n) * Math.cos(delta[m] - delta[n])));
                        }
                    }

                    // Construct matrix like this:
                    // J1 | J2
                    // -------
                    // J3 | J4
                    Matrix J12 = Util.concatColumns(J1, J2);
                    Matrix J34 = Util.concatColumns(J3, J4);
                    Matrix J = Util.concatRows(J12, J34);

                    // Inverse of matrix
                    try {
                        LUDecomposition lu = new LUDecomposition(J);
                        inverseJ = lu.solve(Matrix.identity(J.getRowDimension(), J.getColumnDimension()));
                    }
                    catch (Exception e) {
                        for (int row = 0; row < J.getRowDimension(); row++) {
                            for (int col = 0; col < J.getColumnDimension(); col++)
                                System.out.printf("%f ", J.get(col, row));
                            System.out.println(";");
                        }
                        throw e;
                    }

                    // Try to find good factor to slow down convergence
                    factor = 0.1;
                    try {
                        double[] eigenvalues = inverseJ.eig().getRealEigenvalues();
                        double maxEigenvalue = Arrays.stream(eigenvalues).map(Math::abs).max().orElse(0);

                        // Will only convergence when eigenvalues of matrix are < 1
                        if (maxEigenvalue > 1)
                            factor = 1 / maxEigenvalue;
                    } catch (Exception unused) {
                        // ignored
                    }

                    jacobianRecalculation = settings.jacobianRecalculationInterval;
                }
                else
                    jacobianRecalculation--;

                // Newton step
                Matrix x = inverseJ.times(Util.fromArray(error));

                // Update current state (delta, voltage)
                for (int i = 0; i < busCount - 1; i++)
                    delta[i + 1] += x.get(i, 0) * factor;
                for (int i = 0; i < pq.length; i++)
                    V[pq[i]] += x.get(i + busCount - 1, 0) * factor;
            }

            // Tolerance is to high
            if (tolerance >= settings.maxError)
                return SolverResult.error(SolverError.ConvergenceError);

            // Complex voltages
            ComplexNumber[] voltages = new ComplexNumber[V.length];
            for (int i = 0; i < V.length; i++)
                voltages[i] = ComplexNumber.fromPolar(V[i], delta[i]);

            // Calculate flows
            FlowData flowData = FlowCalculation.calculate(admittanceMatrix, problem.buses, problem.lines, voltages);

            // Save finished time
            performance.timeInMilliseconds = (System.nanoTime() - timeStart) / 1e6;

            // Return success
            return SolverResult.success(performance, voltages, flowData);
        }
        catch(Exception e) {
            e.printStackTrace();
        }
        return SolverResult.error(SolverError.Internal);
    }
}