"""This script generates a version 1 JSON file.

Run with roseau-load-flow==0.9.1, NOT WITH THE CURRENT VERSION::

    uv run --no-project --script roseau/load_flow/io/tests/data/generators/v1.py
"""

# Script metadata generated by `uv`
# /// script
# requires-python = ">=3.10,<3.12"
# dependencies = [
#     "roseau-load-flow==0.9.1",
# ]
# ///
import cmath
from pathlib import Path

import numpy as np
from shapely import LineString, Point

import roseau.load_flow as rlf
from roseau.load_flow.io.dict import NETWORK_JSON_VERSION

if NETWORK_JSON_VERSION != 1:
    raise ValueError(
        f"This script is meant to generate version 1 JSON files, current version: {NETWORK_JSON_VERSION}"
        f"\nExecute with `uv run --no-project --script roseau/load_flow/io/tests/data/generators/v1.py`"
    )

OUTPUT = Path(__file__).parent.parent / f"network_json_v{NETWORK_JSON_VERSION}.json"

bus1 = rlf.Bus(id=1, phases="abc", geometry=Point(0.0, 0.0), min_voltage=19e3, max_voltage=21e3)
bus2 = rlf.Bus(id=2, phases="abc", geometry=Point(0.0, 1.0), min_voltage=19e3, max_voltage=21e3)
bus3 = rlf.Bus(id=3, phases="abcn", geometry=Point(0.0, 1.0), min_voltage=207, max_voltage=253)
bus4 = rlf.Bus(id=4, phases="abcn", geometry=Point(0.0, 2.0), min_voltage=207, max_voltage=253)
bus5 = rlf.Bus(id=5, phases="abcn", geometry=Point(0.0, 2.0), min_voltage=207, max_voltage=253)
bus6 = rlf.Bus(id=6, phases="bc", geometry=Point(0.0, 3.0), min_voltage=207, max_voltage=253)
bus7 = rlf.Bus(id=7, phases="abn", geometry=Point(0.0, 3.0), min_voltage=207, max_voltage=253)

ground = rlf.Ground(id="gnd")
ground.connect(bus3, phase="n")
ground.connect(bus5, phase="n")
potential_ref = rlf.PotentialRef(id="pref", element=ground)

vs = rlf.VoltageSource(id=1, bus=bus1, voltages=20e3 * np.exp([0, -2j * np.pi / 3, 2j * np.pi / 3]), phases="abc")
fp = rlf.FlexibleParameter(
    control_p=rlf.Control.p_max_u_consumption(u_min=18e3, u_down=19e3),
    control_q=rlf.Control.constant(),
    projection=rlf.Projection(type="euclidean"),
    s_max=45e3,
)
power = cmath.rect(40e3, np.acos(0.95))
rlf.PowerLoad(id=1, bus=bus2, phases="abc", powers=[power, power, power], flexible_params=[fp, fp, fp])
rlf.PowerLoad(
    id=2,
    bus=bus2,
    phases="abc",
    powers=[
        40459.7989783205 + 20229.89948916025j,
        40459.79897941102 + 20229.89948970551j,
        40459.79897941102 + 20229.89948970551j,
    ],
)


lp_mv = rlf.LineParameters(
    id=1,
    z_line=0.35 * np.eye(3, dtype=complex),
    y_shunt=1e-6j * np.eye(3, dtype=complex),
    line_type=rlf.LineType.UNDERGROUND,
    conductor_type=rlf.ConductorType.AL,
    insulator_type=rlf.InsulatorType.PVC,
    section=rlf.Q_(240, "mm^2"),
    max_current=250,
)
lp_lv = rlf.LineParameters(
    id=2,
    z_line=0.35 * np.eye(4, dtype=complex),
    line_type=rlf.LineType.OVERHEAD,
    conductor_type=rlf.ConductorType.CU,
    insulator_type=None,
    section=rlf.Q_(150, "mm^2"),
    max_current=160,
)
rlf.Line(
    id=1, bus1=bus1, bus2=bus2, parameters=lp_mv, length=10.0, geometry=LineString([(0, 0), (1, 0)]), ground=ground
)
rlf.Line(id=2, bus1=bus3, bus2=bus4, parameters=lp_lv, length=1.0, geometry=LineString([(0, 1), (0, 2)]))

tp = rlf.TransformerParameters.from_catalogue("SE_Minera_AA0Ak_160kVA")
tp.max_power = 200e3
rlf.Transformer(id=3, bus1=bus2, bus2=bus3, parameters=tp, geometry=Point(0.0, 1.0), tap=1.025)
tp_single = rlf.TransformerParameters(
    id="2", type="single", uhv=400, ulv=400, sn=10e3, z2=0.05, ym=0.01 + 0.02j, max_power=10e3
)
rlf.Transformer(id=4, bus1=bus4, bus2=bus5, parameters=tp_single, geometry=Point(0.0, 2.0), phases1="bc", phases2="bc")
tp_center = rlf.TransformerParameters(
    id="3", type="center", uhv=400, ulv=400, sn=10e3, z2=0.05, ym=0.01 + 0.02j, max_power=10e3
)
rlf.Transformer(id=5, bus1=bus4, bus2=bus5, parameters=tp_center, geometry=Point(0.0, 2.0), phases1="ab", phases2="abn")

lp_lv_2_cond = rlf.LineParameters(id=3, z_line=0.35 * np.eye(2, dtype=complex))
lp_lv_3_cond = rlf.LineParameters(id=4, z_line=0.35 * np.eye(3, dtype=complex))
rlf.Line(
    id=6, bus1=bus5, bus2=bus6, phases="bc", parameters=lp_lv_2_cond, length=1.0, geometry=LineString([(0, 2), (0, 3)])
)
rlf.Line(
    id=7, bus1=bus5, bus2=bus7, phases="abn", parameters=lp_lv_3_cond, length=1.0, geometry=LineString([(0, 2), (0, 3)])
)
rlf.PowerLoad(id=3, bus=bus6, phases="bc", powers=[1000 + 200j])
rlf.CurrentLoad(id=4, bus=bus7, phases="abn", currents=[2 + 1j, 0.5 + 0.1j])
rlf.ImpedanceLoad(id=5, bus=bus7, phases="abn", impedances=[100 + 200j, 300 + 4000j])

en = rlf.ElectricalNetwork.from_element(bus1)
en.solve_load_flow()

# Bug in PowerLoad.res_flexible_powers with include_results=True
assert en.loads[1].flexible_params is not None
en.loads[1]._res_flexible_powers = en.loads[1].res_powers.m

en.to_json(OUTPUT, include_results=True)
