# Copyright 2023 Euratom
# Copyright 2023 United Kingdom Atomic Energy Authority
# Copyright 2023 Centro de Investigaciones Energéticas, Medioambientales y Tecnológicas
#
# Licensed under the EUPL, Version 1.1 or – as soon they will be approved by the
# European Commission - subsequent versions of the EUPL (the "Licence");
# You may not use this work except in compliance with the Licence.
# You may obtain a copy of the Licence at:
#
# https://joinup.ec.europa.eu/software/page/eupl5
#
# Unless required by applicable law or agreed to in writing, software distributed
# under the Licence is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR
# CONDITIONS OF ANY KIND, either express or implied.
#
# See the Licence for the specific language governing permissions and limitations
# under the Licence.
"""Module for loading magnetic field data from the equilibrium IDS."""

from dataclasses import dataclass
from warnings import warn

import numpy as np

from imas.ids_defs import EMPTY_INT
from imas.ids_struct_array import IDSStructArray

RECTANGULAR_GRID = 1

__all__ = ["load_magnetic_field_data"]


@dataclass
class MagneticField2DData:
    """Dataclass to hold 2D magnetic field data."""

    r: np.ndarray
    z: np.ndarray
    b_field_r: np.ndarray
    b_field_z: np.ndarray
    b_field_phi: np.ndarray


def load_magnetic_field_data(profiles_2d: IDSStructArray) -> MagneticField2DData:
    """Load 2D profiles of the magnetic field components from equilibrium IDS.

    The magnetic field components are extracted from the ``profiles_2d`` IDS structure,
    assuming that the profiles are defined on a rectangular grid.

    Parameters
    ----------
    profiles_2d
        The ``profiles_2d`` structure of the equilibrium IDS.

    Returns
    -------
    MagneticField2DData
        Dataclass containing the magnetic field data:

        :r: ``(N,)`` ndarray with R coordinates of rectangular grid.
        :z: ``(M,)`` ndarray with Z coordinates of rectangular grid.
        :b_field_r: ``(N, M)`` ndarray with R component of the magnetic field.
        :b_field_z: ``(N, M)`` ndarray with Z component of the magnetic field.
        :b_field_phi: ``(N, M)`` ndarray with toroidal component of the magnetic field.

    Raises
    ------
    RuntimeError
        If unable to read the magnetic field due to unsupported grid type or
        mismatched array shapes.
    RuntimeError
        If the toroidal component of the magnetic field is not found in the IDS structure.
    """
    rectangular_grid = False
    for prof2d in profiles_2d:
        if prof2d.grid_type.index == RECTANGULAR_GRID or prof2d.grid_type.index == EMPTY_INT:
            rectangular_grid = True
            break

    if not rectangular_grid:
        raise RuntimeError(
            "Unable to read magnetic field: "
            + "rectangular grid for 2D profiles is not found and other grid types are not supported."
        )

    r = np.asarray_chkfinite(prof2d.grid.dim1)
    z = np.asarray_chkfinite(prof2d.grid.dim2)
    shape = (r.size, z.size)

    b_field_r = np.asarray_chkfinite(prof2d.b_field_r)
    b_field_z = np.asarray_chkfinite(prof2d.b_field_z)

    # TODO: Remove fallback to `b_field_tor` after confirming that all relevant IDS structures have `b_field_phi`
    if hasattr(prof2d, "b_field_phi"):
        b_field_phi = np.asarray_chkfinite(prof2d.b_field_phi)
    elif hasattr(prof2d, "b_field_tor"):
        warn(
            "The 'b_field_tor' field is deprecated since DD v3.42.0; "
            "Support for 'b_field_tor' fallback will be removed in a future release.",
            DeprecationWarning,
            stacklevel=2,
        )
        b_field_phi = np.asarray_chkfinite(prof2d.b_field_tor)
    else:
        raise RuntimeError(
            "Unable to read magnetic field: "
            + "toroidal component of the magnetic field is not found in the IDS structure."
        )

    if b_field_r.shape != shape or b_field_phi.shape != shape or b_field_z.shape != shape:
        raise RuntimeError(
            "Unable to read magnetic field: "
            + "the shape of the magnetic field components does not match the grid shape."
        )

    return MagneticField2DData(
        r=r,
        z=z,
        b_field_r=b_field_r,
        b_field_z=b_field_z,
        b_field_phi=b_field_phi,
    )
