# AUTOGENERATED! DO NOT EDIT! File to edit: ../03_ConvertFEEMSResultToProto.ipynb.

# %% auto 0
__all__ = ["logger", "ch", "formatter", "FEEMSResultConverter"]

# %% ../03_ConvertFEEMSResultToProto.ipynb 3
from functools import cached_property
from typing import Union, List, Dict, Optional, cast

from feems.components_model.component_electric import FuelCellSystem
from feems.fuel import FuelSpecifiedBy
from feems.types_for_feems import FEEMSResult, TypeComponent, EmissionType
from feems.system_model import (
    MechanicalPropulsionSystemWithElectricPowerSystem,
    ElectricPowerSystem,
    HybridPropulsionSystem,
    MechanicalPropulsionSystem,
    FEEMSResultForMachinerySystem,
)
import pandas as pd
import numpy as np
import MachSysS.feems_result_pb2 as proto
import MachSysS.gymir_result_pb2 as proto_gymir
import logging

# %% ../03_ConvertFEEMSResultToProto.ipynb 4
# Define logger
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

ch = logging.StreamHandler()
ch.setLevel(logging.INFO)

formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
ch.setFormatter(formatter)

logger.addHandler(ch)

_COLUMN_NAMES = {
    "multi fuel consumption [kg]": "multi_fuel_consumption_kg",
    "electric energy consumption [MJ]": "electric_energy_consumption_mj",
    "mechanical energy consumption [MJ]": "mechanical_energy_consumption_mj",
    "energy_stored [MJ]": "energy_stored_mj",
    "running hours [h]": "running_hours_h",
    "CO2 emission [kg]": "co2_emissions_kg",
    "NOx emission [kg]": "nox_emissions_kg",
    "component type": "component_type",
    "rated capacity": "rated_capacity",
    "rated capacity unit": "rated_capacity_unit",
    "switchboard id": "switchboard_id",
    "shaftline id": "shaftline_id",
}


class FEEMSResultConverter:
    feems_result: Union[FEEMSResult, FEEMSResultForMachinerySystem, None]
    time_series_input: Optional[proto_gymir.TimeSeriesResult] = None
    fuel_specified_by: FuelSpecifiedBy = FuelSpecifiedBy.IMO

    _system_feems: Union[
        MechanicalPropulsionSystemWithElectricPowerSystem,
        ElectricPowerSystem,
        HybridPropulsionSystem,
        None,
    ] = None
    _time_series_data_for_electric_component: Union[
        List[Dict[str, proto.TimeSeriesResultForComponent]], None
    ] = None
    _time_series_data_for_mechanical_component: Union[
        List[Dict[str, proto.TimeSeriesResultForComponent]], None
    ] = None

    def __init__(
        self,
        feems_result: Union[FEEMSResult, FEEMSResultForMachinerySystem],
        system_feems: [
            MechanicalPropulsionSystemWithElectricPowerSystem,
            ElectricPowerSystem,
            HybridPropulsionSystem,
        ],
        time_series_input: Optional[proto_gymir.TimeSeriesResult] = None,
        fuel_specified_by: FuelSpecifiedBy = FuelSpecifiedBy.IMO,
    ):
        self.feems_result = feems_result
        self.time_series_input = time_series_input
        self.system_feems = system_feems
        self.fuel_specified_by = fuel_specified_by

    @property
    def system_feems(self):
        return self._system_feems

    @system_feems.setter
    def system_feems(self, value):
        self._system_feems = value
        self._time_series_data_for_electric_component = []
        self._time_series_data_for_mechanical_component = []

    @property
    def isSystemElectric(self) -> bool:
        return isinstance(self.system_feems, ElectricPowerSystem)

    @property
    def electric_system(self) -> ElectricPowerSystem:
        return (
            self.system_feems
            if self.isSystemElectric
            else self.system_feems.electric_system
        )

    @property
    def mechanical_system(self) -> MechanicalPropulsionSystem:
        return (
            self.system_feems.mechanical_system if not self.isSystemElectric else None
        )

    @cached_property
    def _time_interval_to_time_array(self) -> np.ndarray:
        number_points = self.electric_system.power_sources[0].power_output.size
        if self.time_series_input is None:
            time_interval_s = self.electric_system.time_interval_s
            if np.isscalar(time_interval_s):
                return np.linspace(
                    0, (number_points - 1) * time_interval_s, number_points
                )
            else:
                return time_interval_s.cumsum()
        else:
            return np.array(
                [
                    power_instance.epoch_s
                    for power_instance in self.time_series_input.propulsion_power_timeseries
                ]
            )[:number_points]

    def _retrieve_time_series_data_from_components(self):
        """Retrieve time series data from components of power sources and energy storages and
        store them in a list of dictionaries. Each dictionary contains the name of the component,
        the data in the proto data type of TimeSeriesResultForComponent and the switchboard id.
        """
        # in the feems result
        power_sources_electric = (
            self.electric_system.power_sources
            + self.electric_system.energy_storage
            + self.electric_system.pti_pto
        )
        if self.mechanical_system is not None:
            power_sources_mechanical = self.mechanical_system.main_engines
        else:
            power_sources_mechanical = []

        for power_source in power_sources_electric:
            result_proto = proto.TimeSeriesResultForComponent(
                time=self._time_interval_to_time_array,
                power_output_kw=power_source.power_output.tolist(),
            )
            if power_source.type == TypeComponent.GENSET:
                genset_result = (
                    power_source.get_fuel_cons_load_bsfc_from_power_out_generator_kw(
                        fuel_specified_by=self.fuel_specified_by,
                    )
                )
                for fuel in genset_result.engine.fuel_flow_rate_kg_per_s.fuels:
                    result_proto.fuel_consumption_kg_per_s.fuels.append(
                        proto.FuelArray(
                            fuel_type=fuel.fuel_type.value,
                            fuel_origin=fuel.origin.value,
                            fuel_specified_by=fuel.fuel_specified_by.value,
                            mass_or_mass_fraction=fuel.mass_or_mass_fraction.tolist(),
                            lhv_mj_per_g=fuel.lhv_mj_per_g,
                        )
                    )
            elif power_source.type in [
                TypeComponent.FUEL_CELL_SYSTEM,
                TypeComponent.FUEL_CELL,
            ]:
                power_source = cast(FuelCellSystem, power_source)
                fuel_cell_result = power_source.get_fuel_cell_run_point(
                    power_source.power_output
                )
                for fuel in fuel_cell_result.fuel_flow_rate_kg_per_s.fuels:
                    result_proto.fuel_consumption_kg_per_s.fuels.append(
                        proto.FuelArray(
                            fuel_type=fuel.fuel_type.value,
                            fuel_origin=fuel.origin.value,
                            fuel_specified_by=fuel.fuel_specified_by.value,
                            mass_or_mass_fraction=fuel.mass_or_mass_fraction.tolist(),
                            lhv_mj_per_g=fuel.lhv_mj_per_g,
                        )
                    )
            elif power_source.type in [
                TypeComponent.BATTERY,
                TypeComponent.BATTERY_SYSTEM,
            ]:
                pass
            elif power_source.type == [TypeComponent.PTI_PTO_SYSTEM]:
                pass
            else:
                raise NotImplementedError(
                    f"Retrieving time series data for component type {power_source.type} "
                    f"is not implemented."
                )
            self._time_series_data_for_electric_component.append(
                dict(
                    name=power_source.name,
                    data=result_proto,
                    node_id=power_source.switchboard_id,
                )
            )
        for power_source in power_sources_mechanical:
            result_proto = proto.TimeSeriesResultForComponent(
                time=self._time_interval_to_time_array,
                power_output_kw=power_source.power_output.tolist(),
            )
            engine_run_result = power_source.get_engine_run_point_from_power_out_kw(
                power_source.power_output
            )
            for fuel in engine_run_result.fuel_flow_rate_kg_per_s.fuels:
                result_proto.fuel_consumption_kg_per_s.fuels.append(
                    proto.FuelArray(
                        fuel_type=fuel.fuel_type.value,
                        fuel_origin=fuel.origin.value,
                        fuel_specified_by=fuel.fuel_specified_by.value,
                        mass_or_mass_fraction=fuel.mass_or_mass_fraction.tolist(),
                    )
                )
            self._time_series_data_for_mechanical_component.append(
                dict(
                    name=power_source.name,
                    data=result_proto,
                    node_id=power_source.shaft_line_id,
                )
            )

    def _get_feems_result_proto_for_subsystem(
        self,
        feems_result: FEEMSResult,
        feems_system: Union[ElectricPowerSystem, MechanicalPropulsionSystem],
        include_time_series_for_components: bool = False,
        verbose: bool = False,
    ) -> proto.FeemsResult:
        """Convert the result of FEEMS calculation to the proto data type of FeemsResult.
        Args:
            feems_result: The result of FEEMS calculation.
            feems_system: The system for which the FEEMS calculation is performed.
        Returns:
            The result of FEEMS calculation in the proto data type of FeemsResult.
        """
        is_electric = isinstance(feems_system, ElectricPowerSystem)
        if not is_electric and not isinstance(feems_system, MechanicalPropulsionSystem):
            raise TypeError(
                f"feems_system must be either ElectricPowerSystem or MechanicalPropulsionSystem, "
                f"but {type(feems_system)} is given."
            )
        if include_time_series_for_components:
            self._retrieve_time_series_data_from_components()
            time_series_data = (
                self._time_series_data_for_electric_component
                if is_electric
                else self._time_series_data_for_mechanical_component
            )
        else:
            time_series_data = None
        result = proto.FeemsResult()
        for key, value in feems_result.__dict__.items():
            if key == "detail_result":
                continue
            if key == "multi_fuel_consumption_total_kg":
                for fuel in value.fuels:
                    result.multi_fuel_consumption_total_kg.fuels.append(
                        proto.FuelScalar(
                            fuel_type=fuel.fuel_type.value,
                            fuel_origin=fuel.origin.value,
                            fuel_specified_by=fuel.fuel_specified_by.value,
                            mass_or_mass_fraction=fuel.mass_or_mass_fraction,
                            lhv_mj_per_g=fuel.lhv_mj_per_g,
                        )
                    )
                continue
            if key == "total_emission_kg":
                result.nox_emission_total_kg = value[EmissionType.NOX]
                continue
            if hasattr(result, key):
                if value is not None:
                    setattr(result, key, value)
                else:
                    if verbose:
                        logger.warning(f"Value for {key} is None.")
            else:
                if verbose:
                    logger.warning(
                        f"There is no matching key in Protobuf message for {key}"
                    )
        if feems_result.detail_result is not None:
            for (
                component_name,
                component_result,
            ) in feems_result.detail_result.iterrows():
                result_per_component = proto.ResultPerComponent(
                    component_name=component_name
                )
                for key, value in component_result.items():
                    key_result_per_component = _COLUMN_NAMES.get(key)
                    if hasattr(
                        result_per_component, key_result_per_component
                    ) or result_per_component.HasField(key_result_per_component):
                        if key_result_per_component == "multi_fuel_consumption_kg":
                            for fuel in value.fuels:
                                result_per_component.multi_fuel_consumption_kg.fuels.append(
                                    proto.FuelScalar(
                                        fuel_type=fuel.fuel_type.value,
                                        fuel_origin=fuel.origin.value,
                                        fuel_specified_by=fuel.fuel_specified_by.value,
                                        mass_or_mass_fraction=fuel.mass_or_mass_fraction,
                                        lhv_mj_per_g=fuel.lhv_mj_per_g,
                                    )
                                )
                            continue
                        if value is not None:
                            setattr(
                                result_per_component, key_result_per_component, value
                            )
                        else:
                            if verbose:
                                logger.warning(f"Value for {key} is None.")
                    else:
                        if verbose:
                            logger.warning(
                                f"There is no matching key in Protobuf message for {key} in "
                                f"{component_name}: {key_result_per_component}"
                            )
                if include_time_series_for_components:
                    try:
                        switchboard_shaft_line_id = (
                            component_result["switchboard id"]
                            if is_electric
                            else component_result["shaftline id"]
                        )
                        result_per_component.result_time_series.CopyFrom(
                            next(
                                filter(
                                    lambda time_series_item: (
                                        time_series_item["name"] == component_name
                                    )
                                    and (
                                        time_series_item["node_id"]
                                        == switchboard_shaft_line_id
                                    ),
                                    time_series_data,
                                )
                            ).get("data")
                        )
                    except StopIteration:
                        logger.warning(
                            f"No time-series data found for {component_name}"
                        )
                result.detailed_result.append(result_per_component)
        return result

    def get_feems_result_proto(
        self, include_time_series_for_components: bool = False, verbose: bool = False
    ) -> proto.FeemsResultForMachinerySystem:
        if self.isSystemElectric:
            feems_result_electric = self._get_feems_result_proto_for_subsystem(
                feems_result=self.feems_result,
                feems_system=self.electric_system,
                include_time_series_for_components=include_time_series_for_components,
                verbose=verbose,
            )
            return proto.FeemsResultForMachinerySystem(
                electric_system=feems_result_electric
            )
        else:
            feems_result_electric = self._get_feems_result_proto_for_subsystem(
                feems_result=self.feems_result.electric_system,
                feems_system=self.electric_system,
                include_time_series_for_components=include_time_series_for_components,
                verbose=verbose,
            )
            feems_result_mechanical = self._get_feems_result_proto_for_subsystem(
                feems_result=self.feems_result.mechanical_system,
                feems_system=self.mechanical_system,
                include_time_series_for_components=include_time_series_for_components,
                verbose=verbose,
            )
            return proto.FeemsResultForMachinerySystem(
                electric_system=feems_result_electric,
                mechanical_system=feems_result_mechanical,
            )

    def get_timeseries_for_power_sources_and_energy_storage(self):
        df = pd.DataFrame()
        df.index = pd.Series(
            data=self._time_interval_to_time_array,
            name="time",
        )
        for each_time_series in [
            *self._time_series_data_for_electric_component,
            *self._time_series_data_for_mechanical_component,
        ]:
            name = each_time_series["name"]
            data = each_time_series["data"]
            df[f"{name}-power_output_kw"] = list(data.power_output_kw)
            if len(data.fuel_consumption_kg_per_s) > 0:
                df[f"{name}-fuel_consumption_kg_per_s"] = list(
                    data.fuel_consumption_kg_per_s
                )
        return df
