% Newton Raphson Load flow solution
clear all; close all; clc;

impmat;

iteration = 200; % max iterations
tolerance = 1e-6;

n_bus = size(powerdata, 1);

error = zeros(n_bus, 1);

V = powerdata(:, 3);
delta = powerdata(:, 4);


P = (powerdata(:, 5) - powerdata(:, 7)) / base_power;
Q = (powerdata(:, 6) - powerdata(:, 8)) / base_power;

% converting to r<theta
theta = angle(Y);
Y = abs(Y);

P_n = P;
Q_n = Q;

n_p = sum(powerdata(:, 2) ~= 1);
n_q = sum(powerdata(:, 2) == 3);



% Init jacobians
J1 = zeros(n_p, n_p);
J2 = zeros(n_p, n_q);
J3 = zeros(n_q, n_p);
J4 = zeros(n_q, n_q);

for iter = 1:iteration

    % Calculate delta P and Q

    for idx = 1:n_bus
        if powerdata(idx, 2) ~= 1
            % PV or PQ bus; update P
            P_n(idx) = V(idx) * sum(Y(idx, :) .* (V .* cos(theta(idx, :)' - repmat(delta(idx), n_bus, 1) + delta))');
        end
        if powerdata(idx, 2) == 3
            % PQ bus; update P and Q
            Q_n(idx) = V(idx) * sum(Y(idx, :) .* (V .* sin(theta(idx, :)' - repmat(delta(idx), n_bus, 1) + delta))');
        end
    end

    delta_p_q = [P - P_n; Q - Q_n];

    error = max(abs(delta_p_q));

    if error < tolerance
        break
    end

    % generate jacobian

    % J1; P vs delta
    for idx = 1:n_p
        for jdx = 1:n_p
            if idx == jdx
                J1(idx, idx) = (V(idx + 1) * sum(Y(idx + 1, :) .* (V .* sin(theta(idx + 1, :)' - repmat(delta(idx + 1), n_bus, 1) + delta))')) - ...
                    (V(idx + 1) * V(idx + 1) * Y(idx + 1, idx + 1) * sin(theta(idx + 1, idx + 1)));
            else
                J1(idx, jdx) = -V(idx + 1) * V(jdx + 1) * Y(idx + 1, jdx + 1) * sin(theta(idx + 1, jdx + 1) - delta(idx + 1) + delta(jdx + 1));
            end
        end
    end

    % J4; Q vs V
    for idx = 1:n_q
        for jdx = 1:n_q
            if idx == jdx
                J4(idx, idx) = (-2 * V(idx + 3) * Y(idx + 3, idx + 3) * sin(theta(idx + 3, idx + 3))) - ...
                    (sum(V' .* (Y(idx + 3, :) .* sin(theta(idx + 3, :)' - repmat(delta(idx + 3), n_bus, 1) + delta)')) - (V(idx + 3) * Y(idx + 3, idx + 3) * sin(theta(idx + 3, idx + 3))));
            else
                J4(idx, jdx) = -V(idx + 3) * Y(idx + 3, jdx + 3) * sin(theta(idx + 3, jdx + 3) - delta(idx + 3) + delta(jdx + 3));
            end
        end
    end

    % J2; P vs V
    for idx = 1:n_p
        for jdx = 1:n_q
            if idx == jdx
                J2(idx, idx) = 2 * V(idx + 1) * Y(idx + 1, idx + 1) * cos(theta(idx + 1, idx + 1)) + ...
                    sum(Y(idx, :) .* (V .* cos(theta(idx, :)' - repmat(delta(idx), n_bus, 1) + delta))');
            else
                J2(idx, jdx) = V(idx + 1) * Y(idx + 1, jdx + 1) * cos(theta(idx + 1, jdx + 1) - delta(idx + 1) + delta(jdx + 1));
            end
        end
    end

    % J3; Q vs delta
    for idx = 1:n_q
        for jdx = 1:n_p
            if idx == jdx
                J3(idx, idx) = (V(idx) * sum(Y(idx, :) .* (V .* cos(theta(idx, :)' - repmat(delta(idx), n_bus, 1) + delta))')) - ...
                    (V(idx) * V(idx) * Y(idx, idx) * cos(theta(idx, idx) - delta(idx) + delta(jdx)));
            else
                J3(idx, jdx) = - V(idx) * V(jdx) * Y(idx, jdx) * cos(theta(idx, jdx) - delta(idx) + delta(jdx));
            end
        end
    end

    J = [J1 J2; J3 J4];
    indices = [2, 3, 4, 5, 9, 10];
    delta_delta_v = J \ delta_p_q(indices);

    V(n_p:end, 1) = V(n_p:end, 1) + delta_delta_v(n_p+1:end, 1);
    delta(2:end, 1) = delta(2:end, 1) + delta_delta_v(1:n_p, 1);
end