# SPDX-FileCopyrightText: 2023 German Aerospace Center <fame@dlr.de>
#
# SPDX-License-Identifier: Apache-2.0

from typing import List, Dict, Iterable, Optional

import pandas as pd
from fameprotobuf.DataStorage_pb2 import DataStorage
from fameprotobuf.Services_pb2 import Output

from fameio.source.results.agent_type import AgentTypeLog
from fameio.source.results.data_transformer import DataTransformer


class OutputDAO:
    """Grants convenient access to content of Output protobuf messages for given DataStorages"""

    def __init__(self, data_storages: List[DataStorage], agent_type_log: AgentTypeLog) -> None:
        self._agent_type_log = agent_type_log
        outputs = self._extract_output_from_data_storages(data_storages)
        self._agent_type_log.update_agents(self._extract_new_agent_types(outputs))
        self._all_series = self._extract_series(outputs)

    @staticmethod
    def _extract_output_from_data_storages(data_storages: List[DataStorage]) -> List[Output]:
        """Returns list of Outputs extracted from given `data_storages`"""
        if data_storages is None:
            data_storages = []
        return [data_storage.output for data_storage in data_storages if data_storage.HasField("output")]

    @staticmethod
    def _extract_new_agent_types(outputs: List[Output]) -> Dict[str, Output.AgentType]:
        """Returns dict of agent names mapped to its type defined in given `outputs`"""
        list_of_agent_type_lists = [output.agentType for output in outputs if len(output.agentType) > 0]
        list_of_agent_types = [item for sublist in list_of_agent_type_lists for item in sublist]
        return {item.className: item for item in list_of_agent_types}

    @staticmethod
    def _extract_series(outputs: List[Output]) -> Dict[str, List[Output.Series]]:
        """Returns series data from associated `outputs` mapped to the className of its agent"""
        list_of_series_lists = [output.series for output in outputs if len(output.series) > 0]
        list_of_series = [series for sublist in list_of_series_lists for series in sublist]

        series_per_class_name = {}
        for series in list_of_series:
            if series.className not in series_per_class_name:
                series_per_class_name[series.className] = []
            series_per_class_name[series.className].append(series)
        return series_per_class_name

    def get_sorted_agents_to_extract(self) -> Iterable[str]:
        """Returns iterator of requested and available agent names in ascending order by count of series"""
        all_series = self._get_agent_names_by_series_count_ascending()
        filtered_series = [agent_name for agent_name in all_series if self._agent_type_log.is_requested(agent_name)]
        return iter(filtered_series)

    def _get_agent_names_by_series_count_ascending(self) -> List[str]:
        """Returns list of agent type names sorted by their amount of series"""
        length_per_agent_types = {agent_name: len(value) for agent_name, value in self._all_series.items()}
        sorted_dict = sorted(length_per_agent_types.items(), key=lambda item: item[1])
        return [agent_name for agent_name, _ in sorted_dict]

    def get_agent_data(self, agent_name: str, data_transformer: DataTransformer) -> Dict[Optional[str], pd.DataFrame]:
        """
        Returns DataFrame(s) containing all data of given `agent` - data is removed after the first call
        Depending on the chosen ResolveOption the dict contains one DataFrame for the simple (and merged columns),
        or, in `SPLIT` mode, additional DataFrames mapped to each complex column's name.
        """
        agent_series = self._all_series.pop(agent_name) if agent_name in self._all_series else []
        agent_type = self._agent_type_log.get_agent_type(agent_name)
        return data_transformer.extract_agent_data(agent_series, agent_type)
