from numpy.typing import NDArray
import numpy as np
from collections import defaultdict
from typing import Sequence

# All function needed to create the geometry

def create_seeds(S, n, m):
    """
    Create seed points for streamlines in cylindrical coordinates.

    Parameters:
    S: array-like of shape (N, 2) with [r, z] coordinates
    n: number of radial divisions (0 to r)
    m: number of angular divisions around z-axis

    Returns:
    numpy array of shape (len(S) * n * m, 3) with [x, y, z] coordinates
    """
    seeds = []

    for r, z in S:
        # Create radial points from 0 to r
        radial_points = np.linspace(0, r, n)

        # Create angular divisions around z-axis
        angles = np.linspace(0, 2 * np.pi, m, endpoint=False)

        for rad in radial_points:
            for angle in angles:
                x = rad * np.cos(angle)
                y = rad * np.sin(angle)
                seeds.append([x.item(), y.item(), z])

    return seeds

def get_circuit_currents(I, C):
    result = {}
    for current, circuit in zip(I, C):
        if circuit not in result:
            result[circuit] = current
    if result is not None and 0 in result:
        del result[0]
    return result

def distribute_inversely(x: float, weights: Sequence[float]) -> np.ndarray:
    weights = np.array(weights, dtype=float)
    inv_weights = 1.0 / weights
    inv_sum = inv_weights.sum()
    distribution = x * (inv_weights / inv_sum)
    return distribution

def discretize_loop(
    r: float,
    z: float,
    correction: float,
    subdivisions: int,
) -> np.ndarray:
    """
    Transform a 2D axisymmetric point into a 3D discretized current loop.

    Converts a single point (r, z) in the cylindrical r-z plane into a 3D
    circular loop subdivided into straight-line segments.  This discretization
    is required for numerical inductance calculations using the Neumann
    formula, which integrates over current paths.

    Each segment is represented by its centre position (x, y, z) and tangent
    vector (dx, dy, dz) weighted by the correction factor.

    Args:
        r:            Radial position of the loop centre (m).
        z:            Axial position of the loop centre (m).
        correction:   Scaling factor applied to each segment (typically
                      turns/serial or turns/(serial*parallel)).
        subdivisions: Number of segments to divide the loop into.

    Returns:
        np.ndarray: 7 x subdivisions array where each column is
            [x, y, z, dx, dy, dz, correction].
            x, y, z are segment centre coordinates; dx, dy, dz are segment
            tangent vector components; correction is the weighting factor.

    Note:
        Uses endpoint=False to avoid double-counting at 0 and 2*pi.
        Tangent vectors are computed analytically for circular geometry.
    """
    angle = np.linspace(0, 2 * np.pi, subdivisions, endpoint=False)
    # Tangent components: derivative of the position vector along the circle.
    dx = -r * np.sin(angle) * 2 * np.pi / subdivisions
    dy =  r * np.cos(angle) * 2 * np.pi / subdivisions
    dz =  np.zeros(subdivisions)
    return np.column_stack((
        r * np.cos(angle),
        r * np.sin(angle),
        np.repeat(z, subdivisions),
        dx,
        dy,
        dz,
        np.repeat(correction, subdivisions),
    )).T


def get_solenoid_scaler(
    turns: int,
    serial: int,
    parallel: int,
    no_parallel: bool,
) -> float:
    """
    Calculate the inductance scaling factor for a single NI coil element.

    Returns the effective number of turns per discretised element, accounting
    for series and parallel subdivisions.  This factor scales the mutual
    inductance between individual discretised elements so that it represents
    the actual magnet geometry.

    The scaling depends on the no_parallel flag:
    - no_parallel=False (default): parallel subdivisions model current paths
      in parallel (same inductance), so only serial contributes.
      scale = turns / serial
    - no_parallel=True: parallel subdivisions represent physical sub-turns
      that contribute additively to inductance.
      scale = turns / (serial * parallel)

    Args:
        turns:       Total number of turns for the coil element.
        serial:      Number of serial subdivisions.
        parallel:    Number of parallel subdivisions.
        no_parallel: When True the parallel parameter subdivides turns
                     (contributing to inductance) rather than modelling
                     parallel current paths.

    Returns:
        float: Scaling factor (effective turns per element).
    """
    if no_parallel:
        return turns / (serial * parallel)
    return turns / serial


def create_dictionaries(turns, serial):
    multiplicity = {}
    for i in range(len(turns)):
        multiplicity[i] = turns[i]/serial[i]
    return multiplicity

def create_circuit_information(ABCD, n_solenoids, node_index_array_per_coil, circuit, I_initial):
    # Creates a list of circuits, nodes added to a specific circuit are connected in series.
    circuits = np.zeros((ABCD.shape[0],))
    signs = np.zeros((ABCD.shape[0],))
    for i in range(n_solenoids):
        circuits[node_index_array_per_coil[i]] += circuit[i]
        signs[node_index_array_per_coil[i]] += np.sign(I_initial[i])
    circuits = [int(i.item()) for i in circuits]
    signs = [int(i.item()) for i in signs]

    circuits_list = {}
    unique_circuits = list(set(circuits))
    for i in unique_circuits:
        circuits_list[i] = [j for j in range(len(circuits)) if circuits[j] == i]
    # Circuit 0 marks unpowered elements; remove it when present.
    # It may be absent when all coils are insulated (no NI elements at all).
    if 0 in circuits_list:
        del circuits_list[0]
    unique_circuits = list(set(circuits_list))

    circuit_coil_list = defaultdict(list)
    for key, value in zip(circuit, range(n_solenoids)):
        circuit_coil_list[key].append(value)
    if 0 in circuit_coil_list.keys():
        circuit_coil_list.pop(0)

    return unique_circuits, circuits_list, circuits, circuit_coil_list, signs

def get_circuit_inductance(circuit, inductance):
    # Getting the inductance per circuit:
    inductance_dict = {}
    for i in set(circuit):
        # Getting the signs of the currents
        if i is not 0:
            c = np.where(np.array(circuit)==i, 1, 0)
            inductance_dict[i] = np.sum(inductance * np.outer(c, c)).item()
    return inductance_dict

def create_circuit_array(ABCD, n_solenoids, node_index_array_per_coil, circuit):
    # Creates a list of circuits, nodes added to a specific circuit are connected in series.
    circuits = np.zeros((ABCD.shape[0],))
    for i in range(n_solenoids):
        circuits[node_index_array_per_coil[i]] += circuit[i]
    circuits = [int(i.item()) for i in circuits]
    return circuits

class incuctance_class:
    # This is the inductance class with all operations related to the inductance.

    def __init__(self, I_initial):
        # Initializes the class.
        self.inductance_matrix = []
        self.inductance_sign_matrix = []
        self.condensed_circuit_matrix = []
        self.condensed_coil_matrix = []
        self.sign = []

        # Getting the sign array and the corrected mutual inductance matrix
        self.calc_sign(I_initial)

    def calc_sign(self, I_initial):
        self.sign = np.sign(I_initial)
        self.inductance_sign_matrix = self.inductance_matrix * np.outer(self.sign, self.sign)

class coil_class:
    # This is a coil class with all information about the coils.

    def __init__(self):
        # Initializes the class.
        self.multiplicity: NDArray[np.float64]      = np.array([], dtype=np.float64)

class circuit_class:
    # This is a circuit class with all information of the circuit.

    def __init__(self):
        # Initializes the class.
        self.multiplicity: NDArray[np.float64]      = np.array([], dtype=np.float64)
        self.Areas: NDArray[np.float64]             = np.array([], dtype=np.float64)
        self.Resistance_Total: float                = 0.0
        self.Resistance: NDArray[np.float64]        = np.array([], dtype=np.float64)
        self.Length: NDArray[np.float64]            = np.array([], dtype=np.float64)
        self.Index_connections: NDArray[np.int]     = np.array([], dtype=np.int)
        self.Index_resistors: NDArray[np.int]       = np.array([], dtype=np.int)
        self.Inductance: float                      = 0.0
        self.Inductance_list: NDArray[np.float64]   = np.array([], dtype=np.float64)
        self.Tau: float                             = 0.0
        self.Tau_list: NDArray[np.float64]          = np.array([], dtype=np.float64)
        self.Mass: float                            = 0.0
        self.Coils: list                            = []
