"""
Newton-Raphson AC Power Flow Solver

Full Newton-Raphson solution of the steady-state AC power-flow problem
in polar coordinates. Supports slack / PV / PQ bus types.

For bus i:  P_i = V_i * sum_k V_k (G_ik cos(th_i-th_k) + B_ik sin(th_i-th_k))
            Q_i = V_i * sum_k V_k (G_ik sin(th_i-th_k) - B_ik cos(th_i-th_k))

Mismatch vector  dF = [dP (for PV+PQ buses); dQ (for PQ buses)]
State vector     x  = [theta (for PV+PQ buses); V (for PQ buses)]
Update           x_(n+1) = x_n + J^-1 dF
"""

import numpy as np


class PowerFlowResult:
    def __init__(self, converged, iterations, V, theta, mismatch_history,
                 bus_types, P_spec, Q_spec):
        self.converged = converged
        self.iterations = iterations
        self.V = V
        self.theta = theta
        self.mismatch_history = mismatch_history
        self.bus_types = bus_types
        self.P_spec = P_spec
        self.Q_spec = Q_spec


def _calc_power_injections(V, theta, Ybus):
    n = len(V)
    G = Ybus.real
    B = Ybus.imag
    P = np.zeros(n)
    Q = np.zeros(n)
    for i in range(n):
        for k in range(n):
            ang = theta[i] - theta[k]
            P[i] += V[i] * V[k] * (G[i, k] * np.cos(ang) + B[i, k] * np.sin(ang))
            Q[i] += V[i] * V[k] * (G[i, k] * np.sin(ang) - B[i, k] * np.cos(ang))
    return P, Q


def _build_jacobian(V, theta, Ybus, pv_pq_idx, pq_idx):
    """Analytical Jacobian."""
    n = len(V)
    G = Ybus.real
    B = Ybus.imag
    P, Q = _calc_power_injections(V, theta, Ybus)

    npv_pq = len(pv_pq_idx)
    npq = len(pq_idx)

    J11 = np.zeros((npv_pq, npv_pq))   # dP/dtheta
    J12 = np.zeros((npv_pq, npq))      # dP/dV
    J21 = np.zeros((npq, npv_pq))      # dQ/dtheta
    J22 = np.zeros((npq, npq))         # dQ/dV

    for a, i in enumerate(pv_pq_idx):
        for b, k in enumerate(pv_pq_idx):
            if i == k:
                J11[a, b] = -Q[i] - B[i, i] * V[i] ** 2
            else:
                ang = theta[i] - theta[k]
                J11[a, b] = V[i] * V[k] * (G[i, k] * np.sin(ang) - B[i, k] * np.cos(ang))

    for a, i in enumerate(pv_pq_idx):
        for b, k in enumerate(pq_idx):
            if i == k:
                J12[a, b] = P[i] / V[i] + G[i, i] * V[i]
            else:
                ang = theta[i] - theta[k]
                J12[a, b] = V[i] * (G[i, k] * np.cos(ang) + B[i, k] * np.sin(ang))

    for a, i in enumerate(pq_idx):
        for b, k in enumerate(pv_pq_idx):
            if i == k:
                J21[a, b] = P[i] - G[i, i] * V[i] ** 2
            else:
                ang = theta[i] - theta[k]
                J21[a, b] = -V[i] * V[k] * (G[i, k] * np.cos(ang) + B[i, k] * np.sin(ang))

    for a, i in enumerate(pq_idx):
        for b, k in enumerate(pq_idx):
            if i == k:
                J22[a, b] = Q[i] / V[i] - B[i, i] * V[i]
            else:
                ang = theta[i] - theta[k]
                J22[a, b] = V[i] * (G[i, k] * np.sin(ang) - B[i, k] * np.cos(ang))

    J = np.block([[J11, J12], [J21, J22]])
    return J, P, Q


def run_newton_raphson(Ybus, bus_types, P_spec, Q_spec, V0=None, theta0=None,
                        tol=1e-6, max_iter=25):
    """
    Parameters
    Ybus : (n,n) complex ndarray
    bus_types : list[str] length n, entries in {'slack','PV','PQ'}
    P_spec, Q_spec : (n,) ndarray of specified injections (pu, generation - load)
    tol : convergence tolerance on max abs mismatch (pu)
    max_iter : iteration cap

    Returns
    PowerFlowResult
    """
    n = len(bus_types)
    V = np.ones(n) if V0 is None else V0.copy()
    theta = np.zeros(n) if theta0 is None else theta0.copy()

    slack_idx = [i for i, t in enumerate(bus_types) if t == "slack"]
    pv_idx = [i for i, t in enumerate(bus_types) if t == "PV"]
    pq_idx = [i for i, t in enumerate(bus_types) if t == "PQ"]
    pv_pq_idx = sorted(pv_idx + pq_idx)

    mismatch_history = []
    converged = False
    it = 0

    for it in range(1, max_iter + 1):
        P, Q = _calc_power_injections(V, theta, Ybus)
        dP = P_spec[pv_pq_idx] - P[pv_pq_idx]
        dQ = Q_spec[pq_idx] - Q[pq_idx]
        mismatch = np.concatenate([dP, dQ])
        max_mis = np.max(np.abs(mismatch)) if len(mismatch) else 0.0
        mismatch_history.append(max_mis)

        if max_mis < tol:
            converged = True
            break

        J, _, _ = _build_jacobian(V, theta, Ybus, pv_pq_idx, pq_idx)
        dx = np.linalg.solve(J, mismatch)

        dtheta = dx[:len(pv_pq_idx)]
        dV = dx[len(pv_pq_idx):]

        for a, i in enumerate(pv_pq_idx):
            theta[i] += dtheta[a]
        for a, i in enumerate(pq_idx):
            V[i] += dV[a]

    return PowerFlowResult(converged, it, V, theta, mismatch_history,
                            bus_types, P_spec, Q_spec)


def branch_flows(V, theta, branches, n_bus):
    """
    Compute complex power flow S_ft and S_tf at each
    branch terminal, in pu, given solved bus voltages.
    Returns a list of dicts with from/to loading info.
    """
    Vc = V * np.exp(1j * theta)
    results = []
    for (f, t, data, is_xfmr) in branches:
        R, X, B = data["R"], data["X"], data.get("B", 0.0)
        z = complex(R, X)
        y = 1.0 / z
        y_sh = complex(0.0, B / 2.0)
        fi, ti = f - 1, t - 1
        Vf, Vt = Vc[fi], Vc[ti]
        I_ft = (Vf - Vt) * y + Vf * y_sh
        I_tf = (Vt - Vf) * y + Vt * y_sh
        S_ft = Vf * np.conj(I_ft)
        S_tf = Vt * np.conj(I_tf)
        loss = S_ft + S_tf
        results.append({
            "from": f, "to": t, "is_xfmr": is_xfmr,
            "S_ft_MVA": S_ft * 100.0, "S_tf_MVA": S_tf * 100.0,
            "loading_MVA": max(abs(S_ft), abs(S_tf)) * 100.0,
            "loss_MW": loss.real * 100.0,
            "rate_mva": data.get("rate_mva", None),
        })
    return results
