"""
Base classes for propulsion components.
"""
#  This file is part of FAST-OAD : A framework for rapid Overall Aircraft Design
#  Copyright (C) 2021 ONERA & ISAE-SUPAERO
#  FAST is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <https://www.gnu.org/licenses/>.

from abc import ABC, abstractmethod
from typing import Union

import numpy as np
import pandas as pd
from openmdao import api as om
from openmdao.core.component import Component

from fastoad.model_base import FlightPoint


class IPropulsion(ABC):
    """
    Interface that should be implemented by propulsion models.

    Using this class allows to delegate to the propulsion model the management of
    propulsion-related data when computing performances.

    The performance model calls :meth:`compute_flight_points` by providing one or
    several flight points. The method will feed these flight points with results
    of the model (e.g. thrust, SFC, ..).

    The performance model will then be able to call :meth:`get_consumed_mass` to
    know the mass consumption for each flight point.

    Note::

        If the propulsion model needs fields that are not among defined fields
        of the :class`FlightPoint class`, these fields can be made authorized by
        :class`FlightPoint class`. Please see part about extensibility in
        :class`FlightPoint class` documentation.
    """

    @abstractmethod
    def compute_flight_points(self, flight_points: Union[FlightPoint, pd.DataFrame]):
        """
        Computes Specific Fuel Consumption according to provided conditions.

        See :class:`~fastoad.model_base.flight_point.FlightPoint` for available fields that may be
        used for computation.
        If a DataFrame instance is provided, it is expected that its columns match
        field names of FlightPoint (actually, the DataFrame instance should be
        generated from a list of FlightPoint instances).

        .. note:: **About thrust_is_regulated, thrust_rate and thrust**

            :code:`thrust_is_regulated` tells if a flight point should be computed using
            :code:`thrust_rate` (when False) or :code:`thrust` (when True) as input. This way,
            the method can be used in a vectorized mode, where each point can be set to respect
            a **thrust** order or a **thrust rate** order.

            - if :code:`thrust_is_regulated` is not defined, the considered input will be the
              defined one between :code:`thrust_rate` and :code:`thrust` (if both are provided,
              :code:`thrust_rate` will be used)

            - if :code:`thrust_is_regulated` is :code:`True` or :code:`False` (i.e., not a sequence),
              the considered input will be taken accordingly, and should of course be defined.

            - if there are several flight points, :code:`thrust_is_regulated` is a sequence or array,
              :code:`thrust_rate` and :code:`thrust` should be provided and have the same shape as
              :code:`thrust_is_regulated:code:`. The method will consider for each element which input
              will be used according to :code:`thrust_is_regulated`.


        :param flight_points: FlightPoint or DataFram instance
        :return: None (inputs are updated in-place)
        """

    @abstractmethod
    def get_consumed_mass(self, flight_point: FlightPoint, time_step: float) -> float:
        """
        Computes consumed mass for provided flight point and time step.

        This method should rely on FlightPoint fields that are generated by
        :meth: `compute_flight_points`.

        :param flight_point:
        :param time_step:
        :return: the consumed mass in kg
        """


class IOMPropulsionWrapper:
    """
    Interface for wrapping a :class:`IPropulsion` subclass in OpenMDAO.

    The implementation class defines the needed input variables for instantiating the
    :class:`IPropulsion` subclass in :meth:`setup` and use them for instantiation in
    :meth:`get_model`

    See
    :class:`~fastoad.models.propulsion.fuel_propulsion.rubber_engine.openmdao.OMRubberEngineWrapper`
    for an example of implementation.
    """

    @abstractmethod
    def setup(self, component: Component):
        """
        Defines the needed OpenMDAO inputs for propulsion instantiation as done in :meth:`get_model`

        Use `add_inputs` and `declare_partials` methods of the provided `component`

        :param component:
        """

    @staticmethod
    @abstractmethod
    def get_model(inputs) -> IPropulsion:
        """
        This method defines the used :class:`IPropulsion` subclass instance.

        :param inputs: OpenMDAO input vector where the parameters that define the
                       propulsion model are
        :return: the propulsion model instance
        """


class BaseOMPropulsionComponent(om.ExplicitComponent, ABC):
    """
    Base class for creating an OpenMDAO component from subclasses of :class:`IOMPropulsionWrapper`.

    Classes that implements this interface should add their own inputs in setup()
    and implement :meth:`get_wrapper`.
    """

    def setup(self):
        self.add_input("data:propulsion:mach", np.nan, shape_by_conn=True)
        self.add_input("data:propulsion:altitude", np.nan, shape_by_conn=True, units="m")
        self.add_input("data:propulsion:engine_setting", np.nan, shape_by_conn=True)
        self.add_input("data:propulsion:use_thrust_rate", np.nan, shape_by_conn=True)
        self.add_input("data:propulsion:required_thrust_rate", np.nan, shape_by_conn=True)
        self.add_input("data:propulsion:required_thrust", np.nan, shape_by_conn=True, units="N")

        self.add_output(
            "data:propulsion:SFC", copy_shape="data:propulsion:mach", units="kg/s/N", ref=1e-4
        )
        self.add_output(
            "data:propulsion:thrust_rate", copy_shape="data:propulsion:mach", lower=0.0, upper=1.0
        )
        self.add_output(
            "data:propulsion:thrust", copy_shape="data:propulsion:mach", units="N", ref=1e5
        )

    def setup_partials(self):
        self.declare_partials("*", "*", method="fd")

    def compute(self, inputs, outputs, discrete_inputs=None, discrete_outputs=None):
        wrapper = self.get_wrapper().get_model(inputs)
        flight_point = FlightPoint(
            mach=inputs["data:propulsion:mach"],
            altitude=inputs["data:propulsion:altitude"],
            engine_setting=inputs["data:propulsion:engine_setting"],
            thrust_is_regulated=np.logical_not(
                inputs["data:propulsion:use_thrust_rate"].astype(int)
            ),
            thrust_rate=inputs["data:propulsion:required_thrust_rate"],
            thrust=inputs["data:propulsion:required_thrust"],
        )
        wrapper.compute_flight_points(flight_point)
        outputs["data:propulsion:SFC"] = flight_point.sfc
        outputs["data:propulsion:thrust_rate"] = flight_point.thrust_rate
        outputs["data:propulsion:thrust"] = flight_point.thrust

    @staticmethod
    @abstractmethod
    def get_wrapper() -> IOMPropulsionWrapper:
        """
        This method defines the used :class:`~fastoad.model_base.propulsion.IOMPropulsionWrapper`
        instance.

        :return: an instance of OpenMDAO wrapper for propulsion model
        """


class AbstractFuelPropulsion(IPropulsion, ABC):
    """
    Propulsion model that consume any fuel should inherit from this one.

    In inheritors, :meth:`compute_flight_points` is expected to define
    "sfc" and "thrust" in computed FlightPoint instances.
    """

    def get_consumed_mass(self, flight_point: FlightPoint, time_step: float) -> float:
        return time_step * flight_point.sfc * flight_point.thrust


class FuelEngineSet(AbstractFuelPropulsion):
    def __init__(self, engine: IPropulsion, engine_count):
        """
        Class for modelling an assembly of identical fuel engines.

        Thrust is supposed equally distributed among them.

        :param engine: the engine model
        :param engine_count:
        """
        self.engine = engine
        self.engine_count = engine_count

    def compute_flight_points(self, flight_points: Union[FlightPoint, pd.DataFrame]):
        if flight_points.thrust is not None:
            flight_points.thrust = flight_points.thrust / self.engine_count

        self.engine.compute_flight_points(flight_points)
        flight_points.thrust = flight_points.thrust * self.engine_count
