## @file spec_field.py
#  @brief Setup the SPEC magnetic field system for ODE solver
#  @author Zhisong Qu (zhisong.qu@anu.edu.au)
#
from .spec_problem import SPECProblem
from .toroidal_bfield import ToroidalBfield
import numpy as np

## Class that used to setup the SPEC bfield problem for interfacing Fortran, used in ODE solver.
# See \ref specbfield for more details.
#
# The SPECBfield system of ODEs is given by
# \f[ \frac{ds}{d\zeta} = \frac{B^{s}}{B^{\zeta}}  \f]
# \f[ \frac{d\theta}{d\zeta} = \frac{B^{\theta}}{B^{\zeta}}  \f]
class SPECBfield(SPECProblem, ToroidalBfield):

    ## the problem size, 2 for 1.5D/2D Hamiltonian system

    def __init__(self, spec_data, lvol):
        """! Set up the equilibrium for use of the fortran module
        @param spec_data the SPEC data generated by py_spec.SPECout
        @param lvol which volume we are interested in, from 1 to spec_data.input.Mvol
        Only support SPEC version >=3.0
        """
        super().__init__(spec_data, lvol)
        self.problem_size = 2
        if self.Igeometry == 1:
            self.poincare_plot_type = "yx"
            self.poincare_plot_xlabel = r"$\theta$"
            self.poincare_plot_ylabel = r"R"
        elif self.Igeometry == 2:
            self.poincare_plot_type = "RZ"
            self.poincare_plot_xlabel = "X(m)"
            self.poincare_plot_ylabel = "Y(m)"
        elif self.Igeometry == 3:
            self.poincare_plot_type = "RZ"
            self.poincare_plot_xlabel = "R(m)"
            self.poincare_plot_ylabel = "Z(m)"
        else:
            raise ValueError("Unknown Igeometry!")

        ## The output of B contains the jacobian factor
        self.has_jacobian = True

    def B(self, coords, *args):
        """! Returns magnetic fields
        @param coordinates \f$(s,\theta,\zeta)\f$
        @param *args extra parameters
        @returns the contravariant magnetic fields
        """

        return self.fortran_module.specbfield.get_bfield(coords)

    def dBdX(self, coords, *args):
        """! Returns magnetic fields
        @param coordinates \f$(s,\theta,\zeta)\f$
        @param *args extra parameters
        @returns B, dBdX, the contravariant magnetic fields, the derivatives of them
        """

        B, dB = self.fortran_module.specbfield.get_bfield_tangent(coords)
        return B, dB

    def B_many(self, x1arr, x2arr, x3arr, input1D=True, *args):
        """! Returns magnetic fields, with multipy coordinate inputs
        @param x1arr the first coordinates. Should have the same length as the other two if input1D=True.
        @param x2arr the second coordinates. Should have the same length as the other two if input1D=True.
        @param x3arr the third coordinates. Should have the same length as the other two if input1D=True.
        @param input1D if False, create a meshgrid with sarr, tarr and zarr, if True, treat them as a list of points
        @param *args extra parameters
        @returns the contravariant magnetic fields
        """
        x1arr = np.atleast_1d(x1arr)
        x2arr = np.atleast_1d(x2arr)
        x3arr = np.atleast_1d(x3arr)

        if not input1D:
            size = (x1arr.size, x2arr.size, x3arr.size)
            s = np.broadcast_to(x1arr[:, np.newaxis, np.newaxis], size).flatten()
            t = np.broadcast_to(x2arr[np.newaxis, :, np.newaxis], size).flatten()
            z = np.broadcast_to(x3arr[np.newaxis, np.newaxis, :], size).flatten()
            n = s.size
            Blist = self.fortran_module.specbfield.get_bfield_many_1d(s, t, z, n)
            Blist = np.reshape(Blist, (x1arr.size, x2arr.size, x3arr.size, 3))
        else:
            n = x1arr.size
            Blist = self.fortran_module.specbfield.get_bfield_many_1d(x1arr, x2arr, x3arr, n)

        return Blist

    def dBdX_many(self, x1arr, x2arr, x3arr, input1D=True, *args):
        """! Returns magnetic fields
        @param x1arr the first coordinates. Should have the same length as the other two if input1D=True.
        @param x2arr the second coordinates. Should have the same length as the other two if input1D=True.
        @param x3arr the third coordinates. Should have the same length as the other two if input1D=True.
        @param input1D if False, create a meshgrid with sarr, tarr and zarr, if True, treat them as a list of points
        @param *args extra parameters
        @returns B, dBdX, the contravariant magnetic fields, the derivatives of them
        """

        x1arr = np.atleast_1d(x1arr)
        x2arr = np.atleast_1d(x2arr)
        x3arr = np.atleast_1d(x3arr)

        if not input1D:
            size = (x1arr.size, x2arr.size, x3arr.size)
            s = np.broadcast_to(x1arr[:, np.newaxis, np.newaxis], size).flatten()
            t = np.broadcast_to(x2arr[np.newaxis, :, np.newaxis], size).flatten()
            z = np.broadcast_to(x3arr[np.newaxis, np.newaxis, :], size).flatten()
            n = s.size
            Blist, dBlist = self.fortran_module.specbfield.get_bfield_tangent_many_1d(s, t, z, n)
            Blist = np.reshape(Blist, (x1arr.size, x2arr.size, x3arr.size, 3))
            dBlist = np.reshape(dBlist, (x1arr.size, x2arr.size, x3arr.size, 3, 3))

        else:
            n = x1arr.size
            Blist, dBlist = self.fortran_module.specbfield.get_bfield_tangent_many_1d(x1arr, x2arr, x3arr, n)
            
        return Blist, dBlist

    def convert_coords(self, stz):
        """! Python wrapper for getting the xyz coordinates from stz
        @param stz  the stz coordinate
        @returns the xyz coordinates
        """
        xyz = self.fortran_module.speccoords.get_xyz(stz)

        # depending on the geometry, return RZ or yx
        if self.Igeometry == 1:
            # for a slab, return x=R, y=theta, z=zeta
            return np.array(
                [
                    xyz[0],
                    np.mod(stz[1], 2 * np.pi) * self.rpol,
                    np.mod(stz[2], 2 * np.pi) * self.rtor,
                ],
                dtype=np.float64,
            )
        if self.Igeometry == 2:
            # for cylinderical geometry, return x=r*cos theta, y=zeta*rtor, z=sin theta
            return np.array(
                [xyz[0] * np.cos(stz[1]), stz[2] * self.rtor, xyz[0] * np.sin(stz[1])],
                dtype=np.float64,
            )
        if self.Igeometry == 3:
            # for toroidal geometry, return x=R, y=zeta, z=Z
            return xyz
