"""
Extender Node with abstract backend.
"""
import typing
import logging
from abc import abstractmethod
import numpy as np

import tfields
import w7x
from w7x.core import Code, Backend


LOGGER = logging.getLogger(__name__)


class ExtenderBackend(Backend):
    """
    Backend base class to be subclassed for Field Line Tracer Nodes
    """

    @staticmethod
    @abstractmethod
    def _plasma_field(state, **kwargs) -> tfields.TensorFields:
        """
        Return the plasma field only (without background field)

        Examples:
            see test_extender
        """

    @staticmethod
    @abstractmethod
    def _field(state, **kwargs) -> tfields.TensorFields:
        """
        Return the extended field

        Examples:
            see test_extender
        """


def explicit_points(points):
    """
    Call provide explicit poitns
    """
    if isinstance(points, tfields.TensorGrid):
        if points.is_empty():
            points = points.explicit()
    if isinstance(points, tfields.Tensors):
        points.transform(tfields.bases.CYLINDER)
    return points


def to_tensor_field(points, field):
    """
    Convert points and field to a tensor field
    """
    if isinstance(points, tfields.TensorGrid):
        tensor_field = points.copy()
        tensor_field.fields = [field]
    else:
        tensor_field = tfields.TensorFields(points, field)
    return tensor_field


@w7x.node
def merge_magnetic_fields_to_equilibrium(
    points: tfields.TensorFields,
    fields: typing.List[typing.Union[tfields.TensorFields, tfields.Tensors]],
    attribute: str,
):
    """
    Merge parts of a magnetic fields to one field. Passing a grid

    Args:
        points: base points of fields
        fields: sub-fields belonging to points in the order given
        attribute: name of the equilibrium attribute to set the field to.
    """
    if len(fields) == 1:
        field = fields[0]
    else:
        field = tfields.Tensors.merged(*fields)

    tensor_field = to_tensor_field(points, field)
    kwargs = {attribute: tensor_field}
    equilibrium = w7x.model.Equilibrium(**kwargs)
    return equilibrium


class Extender(Code):
    """
    High level Field Line Tracer object.
    """

    STRATEGY_TYPE = ExtenderBackend
    STRATEGY_DEFAULT = w7x.config.get_option("extender", "backend")

    def __execute(self, name, state, **kwargs):
        """
        Args:
            name: both, the backend method name (with added underscore before) and attribute of
                w7x.model.Equilibrium
            state
            chunk_size (int): The points fields can become quite large so for data transfer,
                in some situations (e.g. web_service) it is required to split the points in
                parts. The size of the parts is the chunk_size (minimal size).
            **kwargs: see originating method

        Note:
            Joachim Geiger has a trick to select all grid points only in w7x vessel (+ eps)
        """
        # pylint:disable=protected-access
        points = kwargs.pop("points")
        # TODO-2(@dboe) make chunk_size dynamic, depending chunk_size = kwargs.pop("chunk_size")
        chunk_size = kwargs.pop("chunk_size")
        n_proc = max(round(len(points) / chunk_size), 1)
        LOGGER.info("Building %s chunks of points to evaluate with EXTENDER.", n_proc)
        if len(points) == 0:
            raise ValueError("Points to be evaluated are empty.")

        backend_method = getattr(self.backend, "_" + name)
        if n_proc == 1:
            fields = [
                w7x.node(backend_method)(
                    state, points=explicit_points(points), **kwargs
                )
            ]
        else:
            fields = []
            for patch_points in np.array_split(explicit_points(points), n_proc):
                patch_field = backend_method(state, points=patch_points, **kwargs)
                fields.append(patch_field)

        equilibrium = merge_magnetic_fields_to_equilibrium(points, fields, name)
        return equilibrium

    @w7x.dependencies(
        w7x.config.CoilSets.Ideal,
        w7x.model.Equilibrium,  # w7x.config.Equilibria.StandardCaseVacuum,
        points=w7x.config.get_mgrid_options(return_type=tfields.TensorGrid),
        chunk_size=int(w7x.config.get_option("extender", "chunk_size")),
    )
    def field(self, state, **kwargs):
        """
        Calculate the full magnetic field (i.e. generated by plasma currents + external coils).

        Args:
            points: points to evaluate the field on

        """
        equilibrium = self.__execute("field", state, **kwargs)
        return equilibrium

    @w7x.node
    @w7x.dependencies(
        w7x.model.Equilibrium,  # w7x.config.Equilibria.StandardCaseVacuum,
        points=w7x.config.get_mgrid_options(return_type=tfields.TensorGrid),
        chunk_size=int(w7x.config.get_option("extender", "chunk_size")),
    )
    def plasma_field(self, state, **kwargs):
        """
        Calculate the plasma magnetic field, i.e. the field (from plasma currents) added to the
        background field (generated by external coils).
        """
        equilibrium = self.__execute("plasma_field", state, **kwargs)
        return equilibrium
