from contextlib import contextmanager

import magpylib as magpy
from magpylib._src.defaults.defaults_utility import ALLOWED_UNITS_MODES

magpy.units.reset()
magpy.units.mode = "keep"
magpy.units.package = "pint"

Q_ = magpy.units.registry.Quantity

@contextmanager
def catch_raise(prefix="", max_len=200):
    try:
        yield
    except Exception as msg:
        s = str(msg).split("\n")[0][:max_len]
        print(f"{prefix}\033[91m raises: {s}\033[0m")


inputs_types = {
    "array-like": {
        "dimension": [0.001, 0.002, 0.003],
        "polarization": [0.001, 0.002, 0.003],
    },
    "unit-like": {
        "dimension": Q_([1,2,3], "mm"),
        "polarization": Q_([1,2,3], "mT"),
    },
    "mixed": {
        "dimension": Q_([1,2,3], "mm"),
        "polarization": [0.001, 0.002, 0.003],
    },
}
observers_types = {
    "array-like": [.001,.002,.003],
    "unit-like": Q_([1,2,3], "mm"),
    "mixed":[Q_([1,2,3], "mm"), [.001,.002,.003]],
}
for type_name, inputs in inputs_types.items():
    print(f"\n* inputs (\033[92m{type_name}\033[0m): {inputs}")
    for mode in ALLOWED_UNITS_MODES:
        magpy.units.reset()
        magpy.units.mode = mode
        print(f"\n -> units mode: \033[94m{mode!r}\033[0m")
        c = None
        # test inputs
        pref = "      stored: "
        with catch_raise(prefix=pref):
            c = magpy.magnet.Cuboid(**inputs)
            stored = {k: getattr(c, k) for k in inputs}
            print(f"{pref}{stored}")
        # test getB
        if False:#c is not None:
            for type_name, obs in observers_types.items():
                pref = f"      (\033[92m{type_name}\033[0m) getB({obs!r}): "
                with catch_raise(prefix=pref):
                    print(f"{pref}{c.getB(obs)!r}")

import magpylib as magpy

magpy.units.reset()
magpy.units.mode = "keep"
magpy.units.package = "pint"

c = magpy.magnet.Cuboid(dimension=[(1, 1, 1), "cm"], polarization=(0, 0, 1))
sens = magpy.Sensor(position=[(0, 0, 1), "cm"])

B1 = c.getB(sens)
c.polarization = [(0, 0, 1000), "mT"]

B2 = c.getB(sens)
with magpy.units.registry.context("Gaussian"):
    c.magnetization = [(0, 0, 1), "T"]  # instead of A/m, still works

B3 = c.getB(sens)
B1, B2, B3

import magpylib as magpy
import numpy as np
import pyvista as pv

magpy.units.reset()
magpy.units.mode = "keep"
magpy.units.package = "pint"

objects = {
    "Cuboid": magpy.magnet.Cuboid(
        polarization=(0, -0.1, 0),
        dimension=[(1, 1, 1), "cm"],
        position=[(-60, 0, 0), "mm"],
    ),
    "Cylinder": magpy.magnet.Cylinder(
        polarization=[(0, 0, 10000), "µT"],
        dimension=[(10, 10), "mm"],
        position=[(-5, 0, 0), "cm"],
    ),
    "CylinderSegment": magpy.magnet.CylinderSegment(
        polarization=[(0, 0, 0.01), "T"],
        dimension=(0.003, 0.01, 0.01, 0, 140),
        position=(-0.03, 0, 0),
    ),
    "Sphere": magpy.magnet.Sphere(
        polarization=[(0, 0, 10), "mT"],
        diameter="0.01m",
        position=[(-1, 0, 0), "cm"],
    ),
    "Tetrahedron": magpy.magnet.Tetrahedron(
        polarization=[(0, 0, 10), "mT"],
        vertices=[((-1, 0, 0), (1, 0, 0), (0, -1, 0), (0, -1, -1)), "cm"],
        position=[(-4, 0, 4), "cm"],
    ),
    "TriangularMesh": magpy.magnet.TriangularMesh(
        polarization=[(0, 0, 10), "mT"],
        vertices=[pv.Dodecahedron(radius=1).triangulate().points, "cm"],
        faces=pv.Dodecahedron().triangulate().faces.reshape(-1, 4)[:, 1:],
        position=[(-0.01, 0, 0.04), "m"],
    ),
    "Circle": magpy.current.Circle(
        current="1A",
        diameter="0.01m",
        position=[(40, 0, 0), "mm"],
    ),
    "Polyline": magpy.current.Polyline(
        current="1000mA",
        vertices=[
            [
                (1, 0, 0),
                (0, 1, 0),
                (-1, 0, 0),
                (0, -1, 0),
                (1, 0, 0),
            ],
            "cm",
        ],
        position=[(1, 0, 0), "cm"],
    ),
    "Dipole": magpy.misc.Dipole(
        moment=[(0, 0, 1), "A*m**2"],
        position=[(0.03, 0, 0), "m"],
    ),
    "Triangle": magpy.misc.Triangle(
        polarization=[(0, 0, 10), "mT"],
        vertices=[((-0.01, 0, 0), (0.01, 0, 0), (0, 0.01, 0)), "m"],
        position=[(2, 0, 4), "cm"],
    ),
    "Sensor": magpy.Sensor(
        pixel=[[(0, 0, z) for z in (-5, 0, 5)], "mm"],
        position=[(0, -30, 0), "mm"],
    ),
}

objects["Circle"].move(np.linspace((0, 0, 0), (0, 0, 0.05), 20))
objects["Cuboid"].rotate_from_angax(np.linspace(0, 90, 20), "z", anchor=0)

magpy.show(*objects.values())

import magpylib as magpy

class MyUnitsHandler(magpy.units.UnitsHandler, package="package_name"):

    def is_quantity(self, inp):
        ...

    def to_quantity(self, inp, unit):
        ...

    def to_unit(self, inp, unit):
        ...

    def get_unit(self, inp):
        ...

    def get_magnitude(self, inp):
        ...
   
mapgy.units.package="package_name"