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 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]
    del circuits_list[0] # Removing circuit zero '0', since that one is used for non-powered circular elements.
    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                            = []
