import numpy as np
from numba import njit


@njit(cache=True)
def solve_bank_raw(diagonal, off_diagonal, bridge_diagonal, port_weights, coupling,
               pick_coefficient, pick_weights, rhs, bridge_rhs,
               pivots, reduced, free_solution, boundary_response, pick_response, output, bridge_output):
    strings, nodes = diagonal.shape
    bridge_load = 0.
    bridge_drive = 0.
    for string in range(strings):
        pivots[string, 0] = diagonal[string, 0]
        reduced[string, 0] = rhs[string, 0]
        if abs(pick_coefficient[string]) > 0:
            pick_response[string, 0] = pick_weights[string, 0]
        for node in range(1, nodes):
            multiplier = off_diagonal[string, node - 1] / pivots[string, node - 1]
            pivots[string, node] = diagonal[string, node] - multiplier * off_diagonal[string, node - 1]
            reduced[string, node] = rhs[string, node] - multiplier * reduced[string, node - 1]
            if abs(pick_coefficient[string]) > 0:
                pick_response[string, node] = pick_weights[string, node] - multiplier * pick_response[string, node - 1]
        if np.min(np.abs(pivots[string])) == 0 or not np.isfinite(pivots[string]).all():
            raise ValueError('Invalid complex string-bank solve pivot')
        free_solution[string, -1] = reduced[string, -1] / pivots[string, -1]
        boundary_response[string, -1] = 1 / pivots[string, -1]
        if abs(pick_coefficient[string]) > 0:
            pick_response[string, -1] /= pivots[string, -1]
        for node in range(nodes - 2, -1, -1):
            free_solution[string, node] = (reduced[string, node] - off_diagonal[string, node] * free_solution[string, node + 1]) / pivots[string, node]
            boundary_response[string, node] = -off_diagonal[string, node] * boundary_response[string, node + 1] / pivots[string, node]
            if abs(pick_coefficient[string]) > 0:
                pick_response[string, node] = (pick_response[string, node] - off_diagonal[string, node] * pick_response[string, node + 1]) / pivots[string, node]
        if abs(pick_coefficient[string]) > 0:
            compliance = 0.
            free_port = 0.
            boundary_port = 0.
            for node in range(nodes):
                compliance += pick_weights[string, node] * pick_response[string, node]
                free_port += pick_weights[string, node] * free_solution[string, node]
                boundary_port += pick_weights[string, node] * boundary_response[string, node]
            scale = pick_coefficient[string] / (1 + pick_coefficient[string] * compliance)
            for node in range(nodes):
                free_solution[string, node] -= scale * pick_response[string, node] * free_port
                boundary_response[string, node] -= scale * pick_response[string, node] * boundary_port
        bridge_load += coupling[string] - coupling[string] ** 2 * boundary_response[string, -1]
        bridge_drive += coupling[string] * free_solution[string, -1]
    compliance = 0.
    free_port = 0.
    for mode in range(len(port_weights)):
        compliance += port_weights[mode] ** 2 / bridge_diagonal[mode]
        free_port += port_weights[mode] * bridge_rhs[mode] / bridge_diagonal[mode]
    denominator = 1 + bridge_load * compliance
    if abs(denominator) == 0 or not np.isfinite(denominator):
        raise ValueError('Invalid complex shared bridge solve denominator')
    bridge_port = (free_port + compliance * bridge_drive) / denominator
    for mode in range(len(port_weights)):
        bridge_output[mode] = (bridge_rhs[mode] + port_weights[mode] * (bridge_drive - bridge_load * bridge_port)) / bridge_diagonal[mode]
    for string in range(strings):
        for node in range(nodes):
            output[string, node] = free_solution[string, node] + coupling[string] * boundary_response[string, node] * bridge_port


@njit(cache=True)
def solve_bank(diagonal, off_diagonal, bridge_diagonal, port_weights, coupling,
               pick_coefficient, pick_weights, rhs, bridge_rhs,
               pivots, reduced, free_solution, boundary_response, pick_response, output, bridge_output):
    solve_bank_raw(diagonal, off_diagonal, bridge_diagonal, port_weights, coupling,
                   pick_coefficient, pick_weights, rhs, bridge_rhs, pivots, reduced,
                   free_solution, boundary_response, pick_response, output, bridge_output)
    ratio = 0.
    for string in range(len(diagonal)):
        ratio = max(ratio, abs(pick_coefficient[string]) * np.sum(pick_weights[string] ** 2) / np.min(np.abs(diagonal[string])))
    if ratio < 1e4:
        return
    residual = np.empty_like(rhs)
    bridge_residual = np.empty_like(bridge_rhs)
    correction = np.empty_like(output)
    bridge_correction = np.empty_like(bridge_output)
    for iteration in range(4):
        bridge_port = sum(port_weights * bridge_output)
        bridge_drive = 0.
        for string in range(len(diagonal)):
            pick_port = sum(pick_weights[string] * output[string])
            for node in range(diagonal.shape[1]):
                value = diagonal[string, node] * output[string, node]
                if node:
                    value += off_diagonal[string, node - 1] * output[string, node - 1]
                if node + 1 < diagonal.shape[1]:
                    value += off_diagonal[string, node] * output[string, node + 1]
                value += pick_coefficient[string] * pick_weights[string, node] * pick_port
                if node == diagonal.shape[1] - 1:
                    value -= coupling[string] * bridge_port
                residual[string, node] = rhs[string, node] - value
            bridge_drive += coupling[string] * output[string, -1]
        for mode in range(len(bridge_output)):
            value = bridge_diagonal[mode] * bridge_output[mode]
            value += port_weights[mode] * (np.sum(coupling) * bridge_port - bridge_drive)
            bridge_residual[mode] = bridge_rhs[mode] - value
        solve_bank_raw(diagonal, off_diagonal, bridge_diagonal, port_weights, coupling,
                       pick_coefficient, pick_weights, residual, bridge_residual, pivots, reduced,
                       free_solution, boundary_response, pick_response, correction, bridge_correction)
        output += correction
        bridge_output += bridge_correction