import numpy as np
import roseau.load_flow as rlf

# Define optional quantities
un = rlf.Q_(400.0, "V")  # Nominal voltage
u_min = rlf.Q_(90, "%")  # Minimum voltage level
u_max = rlf.Q_(110, "%")  # Maximum
ampacity = rlf.Q_(500, "A")  # Ampacity

# Create two buses
source_bus = rlf.Bus(
    id="sb", phases="abcn", nominal_voltage=un, min_voltage_level=u_min, max_voltage_level=u_max
)
load_bus = rlf.Bus(
    id="lb", phases="abcn", nominal_voltage=un, min_voltage_level=u_min, max_voltage_level=u_max
)

# Define the reference of potentials to be the neutral of the source bus
ground = rlf.Ground(id="gnd")
pref = rlf.PotentialRef(id="pref", element=ground)  # Fix the potential of the ground at 0 V
ground.connect(source_bus, phase="n")

# Create a LV source at the first bus: 230V phase-to-neutral
vs = rlf.VoltageSource(id="vs", bus=source_bus, phases="abcn", voltages=rlf.Q_(230, "V"))

# Add a load at the second bus: 10kW per phase
load = rlf.PowerLoad(id="load", bus=load_bus, phases="abcn", powers=rlf.Q_(10 + 0j, "kVA"))

# Add a LV line between the source bus and the load bus
z = rlf.Q_((0.1 + 0.0j), "ohm/km")  # R = 0.1 Ohm, X = 0
lp = rlf.LineParameters("lp", z_line=z * np.eye(4, dtype=complex), ampacities=ampacity)
line = rlf.Line(id="line", bus1=source_bus, bus2=load_bus, parameters=lp, length=rlf.Q_(2.0, "km"))

# Create the electrical network
en = rlf.ElectricalNetwork.from_element(source_bus)
en