import numpy as np
from imas import DBEntry, IDSFactory, convert_ids


ORIGINAL_DD_VERSION = "3.42.0"
NEW_DD_VERSION = "4.1.0"

with DBEntry("test.nc", "w", dd_version=ORIGINAL_DD_VERSION) as db:
    # Create the IDS factor with a fixed version of the DD (4.1.0)
    # because the following code uses that API
    ids_factory = IDSFactory(version=NEW_DD_VERSION)

    limiter_ids = ids_factory.wall()
    description_2d = limiter_ids.description_2d
    description_2d.resize(1)
    description_2d = description_2d[0]
    unit = description_2d.limiter.unit
    unit.resize(1)
    unit = unit[0]
    outline = unit.outline

    limiter_ids.ids_properties.homogeneous_time = 1
    outline.r = np.array([0.0, 0.0, 1.0, 1.0])
    outline.z = np.array([0.0, 1.0, 1.0, 0.0])

    # Convert the IDS back to 3.42.0 before writing because this is the version
    # of the DBEntry.
    db.put(convert_ids(limiter_ids, ORIGINAL_DD_VERSION))


with DBEntry("test.nc", "r", dd_version=ORIGINAL_DD_VERSION) as db:
    print(f"=== Opening database with DD version {ORIGINAL_DD_VERSION} ===")
    wall_ids = db.get("wall")
    limiter = wall_ids.description_2d[0].limiter.unit[0].outline
    print(f"Limiter r = {limiter.r.value}")
    print(f"Limiter z = {limiter.z.value}")


with DBEntry("test.nc", "r", dd_version=ORIGINAL_DD_VERSION) as db:
    print(
        f"=== Opening database with DD version {ORIGINAL_DD_VERSION}, converting to {NEW_DD_VERSION} ==="
    )
    # the error occurs at this line, when trying to convert the 3.42.0 wall IDS to 4.1.0
    wall_ids = convert_ids(db.get("wall"), NEW_DD_VERSION)
    limiter = wall_ids.description_2d[0].limiter.unit[0].outline
    print(f"Limiter r = {limiter.r.value}")
    print(f"Limiter z = {limiter.z.value}")