from typing import Dict
from pandas import DataFrame
from lib.case_line import convert_cases_to_time_series
from lib.cast import safe_int_cast
from lib.data_source import DataSource
from lib.utils import table_rename


class NewZealandDataSource(DataSource):
    def parse_dataframes(
        self, dataframes: Dict[str, DataFrame], aux: Dict[str, DataFrame], **parse_opts
    ) -> DataFrame:
        cases = table_rename(
            dataframes[0],
            {
                "Report Date": "date_new_confirmed",
                # "CaseStatus": "_status",
                "Sex": "sex",
                "Age group": "age",
                # "DHB_label": "",
                # "Overseas": "",
            },
            drop=True,
        )
        cases["key"] = "NZ"
        cases["age"] = cases["age"].str.slice(0, 2).str.replace(" ", "").apply(safe_int_cast)
        data = convert_cases_to_time_series(cases, ["key"])
        return data
