import datamazing.pandas as pdz
import pandas as pd

from .masterdata_manager import MasterdataManager


class UnitManager(MasterdataManager):
    """
    Manager which simplifies the process of getting units from masterdata.
    """

    def __init__(
        self,
        db: pdz.Database,
        time_interval: pdz.TimeInterval,
        resolution: pd.Timedelta,
    ) -> None:
        self.db = db
        self.time_interval = time_interval
        self.resolution = resolution

    def get_units(
        self,
        filters: dict = {},
        columns: list = [
            "masterdata_gsrn",
            "installed_power_MW",
            "plant_id",
            "price_area",
            "valid_from_date_utc",
            "valid_to_date_utc",
            "power_system_resource_type",
        ],
    ) -> pd.DataFrame:
        """Gets the units for a given unit type.
        Filters for units valid at the end of time interval.
        Filters by default for units in operation.
        """
        return self.get_data("masterdataUnit", filters=filters, columns=columns)
