"""Decodes household data using metadata mappings.

Provides functionality to resolve metadata versions, map codes 
to metadata attributes, and join decoded attributes onto data tables.

Classes
-------
CommodityDecoderSettings - Settings for commodity code decoding.
CommodityDecoder - Decodes commodity codes using metadata.
IDDecoderSettings - Settings for household ID decoding.  
IDDecoder - Decodes household IDs using metadata.

Functions
--------- 
read_classification_info - Reads classified metadata by name.
create_classification_table - Creates table from classified metadata.

The decoders resolve metadata versions, map codes to attributes, 
and add decoded columns to the input tables.

The functions provide helpers for resolving metadata and reading
classification info from the raw metadata.

"""
from itertools import product
from typing import Callable, Iterable, Literal, Any

import pandas as pd
from pydantic import BaseModel

from . import utils, metadata_reader, external_data
from .metadata_reader import metadata, _Attribute


_Default = Any


def read_classification_info(
    name: str, year: int, classification_type: Literal["commodity", "occupation"]
) -> dict:
    """Reads classification metadata by name.

    Retrieves the versioned classification metadata for the given
    classification name and year, resolves the version, categorizes it,
    and returns the resulting dictionary.

    Parameters
    ----------
    name : str
        Name of the classification to get info for.

    year : int
        Year to retrieve metadata for.

    Returns
    -------
    dict
        Dictionary containing categorized classification info for the
        given classification name and year.

    Examples
    --------
    >>> info = read_classification_info('original', 1380)

    See Also
    --------
    utils.resolve_metadata : Resolves metadata versions.

    """
    if classification_type == "commodity":
        versioned_metadata = metadata.commodities[name]
    else:
        versioned_metadata = metadata.occupations[name]
    classification_info = utils.resolve_metadata(
        versioned_metadata, year, categorize=True
    )
    return classification_info


def create_classification_table(
    name: str,
    years: Iterable[int],
    classification_type: Literal["commodity", "occupation"],
) -> pd.DataFrame:
    """Creates classification table for given years.

    Loops through the provided years, reads the classification metadata,
    converts to annual DataFrames, concatenates the results, and returns
    the final classification table.

    Parameters
    ----------
    name : str
        Name of classification to create table for.

    years : Iterable[int]
        Years to include in the resulting table.

    Returns
    -------
    DataFrame
        Classification table with a row for each year.

    See Also
    --------
    _create_annual_classification_table : Converts metadata to DataFrame.

    """
    table_list = []
    for year in years:
        classification_info = read_classification_info(name, year, classification_type)
        annual_table = _create_annual_classification_table(classification_info)
        annual_table.loc[:, "Year"] = year
        table_list.append(annual_table)
    table = pd.concat(table_list, ignore_index=True)
    return table


def _create_annual_classification_table(classification_info: dict) -> pd.DataFrame:
    """Creates annual DataFrame from classification metadata.

    Converts the classification info dictionary for a single year
    into a DataFrame. Applies a helper to extract year ranges from
    'code' values then drops the 'code' column.

    This is used by create_classification_table() to generate the
    annual tables that are concatenated.

    Parameters
    ----------
    classification_info : dict
        Classification metadata for a single year.

    Returns
    -------
    DataFrame
        Annual classification data as a DataFrame.

    See Also
    --------
    create_classification_table : Creates full table by concatting
        annual DataFrames generated by this function.

    """
    table = pd.DataFrame(classification_info["items"])
    table["code_range"] = table["code"].apply(
        utils.Argham,  # type: ignore
        default_start=metadata_reader.defaults.first_year,
        default_end=metadata_reader.defaults.last_year + 1,
        keywords=["code"],
    )
    table = table.drop(columns=["code"])
    return table


def extract_column(table: pd.DataFrame, column_name: str) -> pd.Series:
    """Extracts a column from a DataFrame as a Series.

    Checks table columns, index name(s) for the given column
    name and returns the matching column as a Series.

    Raises KeyError if column not found.

    Parameters
    ----------
    table : DataFrame
        DataFrame to extract column from.

    column_name : str
        Name of column to extract.

    Returns
    -------
    Series
        Extracted column as a Series.

    Raises
    ------
    KeyError
        If column not found in DataFrame.

    """
    if column_name in table.columns:
        column = table.loc[:, column_name].copy()
    elif isinstance(table.index, pd.Index) and table.index.name == column_name:
        column = table.index.to_series()
    elif isinstance(table.index, pd.MultiIndex) and column_name in table.index.names:
        column = table.index.to_frame().loc[:, column_name].copy()
    else:
        raise KeyError
    return column


class DecoderSettings(BaseModel):
    """Settings for decoding commodity codes.

    Attributes
    ----------
    name : str
        Name of classification metadata to use.

    code_column_name : str
        Column name for commodity codes.

    year_column_name : str
        Column name for year.

    labels : tuple[str]
        Labels to extract as output columns.

    levels : tuple[int]
        Hierarchy levels to extract.

    output_column_names : tuple[str]
        Names for extracted output columns.

    See Also
    --------
    CommodityDecoder : Uses these settings to decode codes.

    """

    classification_type: Literal["commodity", "occupation"] = "commodity"
    name: str = "original"
    code_column_name: str = _Default
    year_column_name: str = metadata_reader.defaults.columns.year
    versioned_info: dict = {}
    defaults: dict = {}
    labels: tuple[str, ...] = ()
    levels: tuple[int, ...] = ()
    drop_value: bool = False
    output_column_names: tuple[str, ...] = ()
    required_columns: tuple[str, ...] | None = None
    missing_value_replacements: dict[str, str] | None = None

    def model_post_init(self, __contex=None) -> None:
        if self.classification_type == "commodity":
            if self.code_column_name == _Default:
                self.code_column_name = metadata_reader.defaults.columns.commodity_code
            self.versioned_info = metadata.commodities[self.name]
        else:
            if self.code_column_name == _Default:
                self.code_column_name = metadata_reader.defaults.columns.job_code
            self.versioned_info = metadata.occupations[self.name]
        if "defaults" in self.versioned_info:
            self.defaults = self.versioned_info["defaults"]
        for key, value in self.defaults.items():
            if isinstance(value, list):
                value = tuple(value)
            if (getattr(self, key) is None) or (len(getattr(self, key)) == 0):
                setattr(self, key, value)
        if len(self.labels) == 0:
            self.labels = ("item_key",)
        if len(self.levels) == 0:
            self.levels = (1,)
        self._resolve_column_names()
        super().model_post_init(None)

    def _resolve_column_names(self) -> None:
        if len(self.output_column_names) == 0:
            names = [
                f"{label}_{level}" for label, level in product(self.labels, self.levels)
            ]
            self.output_column_names = tuple(names)
        elif len(self.output_column_names) == len(self.labels) * len(self.levels):
            pass
        elif len(self.output_column_names) == len(self.labels):
            names = [
                f"{label}_{level}"
                for label, level in product(self.output_column_names, self.levels)
            ]
            self.output_column_names = tuple(names)

    @property
    def rename_dict(self):
        """Mapping of label-level keys to output column names.

        Returns a dictionary mapping each combination of labels
        and levels to the corresponding output column name.

        This can be used to rename the multi-level columns in the
        decoded mapping table to the configured output names.

        Returns
        -------
        dict
            Mapping of label-level tuples to output column names.

        """
        label_level = product(self.labels, self.levels)
        return dict(zip(label_level, self.output_column_names))


class Decoder:
    """Decodes commodity codes using classification metadata.

    Parameters
    ----------
    table : DataFrame
        Table with code and year columns to decode.

    settings : DecoderSettings
        Decoding configuration settings.

    Attributes
    ----------
    classification_table : DataFrame
        Resolved classification metadata.

    year_code_pairs : DataFrame
        Unique year and code combinations.

    Methods
    -------
    create_mapping_table()
        Maps codes to metadata based on year.

    add_classification()
        Adds decoded columns to the input table.

    See Also
    --------
    DecoderSettings : Decoding configuration.

    """

    def __init__(self, table: pd.DataFrame, settings: DecoderSettings) -> None:
        self.table = table
        self.settings = settings
        self.code_column = extract_column(table, settings.code_column_name)
        self.year_column = extract_column(table, settings.year_column_name)
        self.classification_table = create_classification_table(
            name=self.settings.name,
            years=self.year_column.drop_duplicates().to_list(),
            classification_type=settings.classification_type,
        )
        self.year_code_pairs = self._create_year_code_pairs()

    def _create_year_code_pairs(self) -> pd.DataFrame:
        years = self.year_column.drop_duplicates()
        yc_pair_list = []
        for year in years:
            filt = self.year_column == year
            codes = self.code_column.loc[filt].drop_duplicates()
            yc_pair = codes.to_frame()
            yc_pair[self.settings.year_column_name] = year
            yc_pair_list.append(yc_pair)
        return pd.concat(yc_pair_list, ignore_index=True)

    def _build_year_code_table(
        self, year_code_pairs: pd.DataFrame, row: pd.Series
    ) -> pd.DataFrame:
        filt = year_code_pairs[self.settings.code_column_name].apply(
            lambda x: x in row["code_range"]
        )
        filt = filt & (
            year_code_pairs[self.settings.year_column_name]
            == row[self.settings.year_column_name]
        )
        matched_codes = year_code_pairs.loc[filt].set_index(
            [self.settings.year_column_name, self.settings.code_column_name]
        )
        columns = row.drop(["code_range", self.settings.year_column_name]).index
        code_table = pd.DataFrame(
            data=[row.loc[columns]] * len(matched_codes.index),
            index=matched_codes.index,
            columns=columns,
        )
        return code_table

    def create_mapping_table(self) -> pd.DataFrame:
        """Creates code mapping table from metadata.

        Loops through classification table and builds a mapping
        table linking codes to metadata based on year.

        Multi-index columns are renamed using the settings.
        Table is validated before returning.

        Returns
        -------
        DataFrame
            Mapping table with year, code and decoded columns.

        See Also
        --------
        _build_year_code_table : Builds table for a single row.

        _validate_mapping_table : Validates mapping integrity.

        """
        code_table_list = []
        for _, row in self.classification_table.iterrows():
            code_table = self._build_year_code_table(self.year_code_pairs, row)
            if not code_table.empty:
                code_table_list.append(code_table)
        mapping_table = pd.concat(code_table_list)
        mapping_table = mapping_table.set_index("level", append=True)
        self._validate_mapping_table(mapping_table)
        mapping_table = mapping_table.unstack(-1)
        mapping_table = mapping_table.loc[:, self.settings.rename_dict.keys()]  # type: ignore
        mapping_table.columns = self.settings.rename_dict.values()
        return mapping_table

    def _validate_mapping_table(self, mapping_table: pd.DataFrame):
        filt = mapping_table.index.duplicated(keep=False)
        if filt.sum() > 0:
            invalid_case_sample = (
                mapping_table.loc[filt]
                .sort_values([self.settings.code_column_name, "level"])
                .head(10)
            )
            raise ValueError(f"Classification is not valid \n{invalid_case_sample}")

    def _fill_missing_values(self):
        if "missing_value_replacements" not in self.settings.defaults:
            return
        for column, default in self.settings.defaults[
            "missing_value_replacements"
        ].items():
            if column not in self.table.columns:
                continue
            filt = self.table.loc[:, column].isna()
            self.table.loc[filt, column] = default  # type: ignore

    def add_classification(self):
        """Adds decoded columns to the input table.

        Joins the mapping table to the input table using
        the year and code columns.

        Also fills in any missing values based on defaults.

        Returns
        -------
        DataFrame
            Input table with decoded columns added.

        """
        mapping = self.create_mapping_table()
        self.table = self.table.join(
            mapping, on=[self.settings.year_column_name, self.settings.code_column_name]
        )
        self._fill_missing_values()
        return self.table


class IDDecoderSettings(BaseModel):
    """Settings for decoding household IDs.

    Attributes
    ----------
    name : Attribute
        Name of household attribute to decode.

    id_column_name : str
        Column name for household IDs.

    year_column_name : str
        Column name for year.

    labels : tuple[str]
        Labels to extract as output columns.

    output_column_names : tuple[str]
        Names of columns to add to output table.

    """

    name: _Attribute
    id_column_name: str = "ID"
    year_column_name: str = "Year"
    labels: tuple[str, ...] = ("name",)
    output_column_names: tuple[str, ...] = tuple()

    def model_post_init(self, __contex=None) -> None:
        self._resolve_output_column_names()
        super().model_post_init(None)

    def _resolve_output_column_names(self) -> None:
        if len(self.output_column_names) != len(self.labels):
            if len(self.labels) == 1:
                names = [self.name]
            else:
                names = [f"{self.name}_{label}" for label in self.labels]
            self.output_column_names = tuple(names)


class IDDecoder:
    """Decodes household IDs using metadata mappings.

    Parameters
    ----------
    table : DataFrame
        Table with ID and year columns to decode.

    settings : IDDecoderSettings
        Configuration settings for decoding.

    Methods
    -------
    construct_mapping_table()
        Builds mapping table from metadata.

    add_attribute()
        Adds decoded columns to the input table.
    """

    def __init__(
        self,
        table: pd.DataFrame,
        settings: IDDecoderSettings,
    ) -> None:
        self.table = table
        self.settings = settings
        self.id_column = extract_column(table, settings.id_column_name)
        self.year_column = extract_column(table, settings.year_column_name)

    def construct_mapping_table(self) -> pd.DataFrame:
        """Constructs metadata mapping table for household IDs.

        Maps ID column to decoded labels based on year. Concatenates
        the ID, year, and label columns into a mapping table.

        Returns
        -------
        DataFrame
            Mapping table with year, ID and decoded columns.

        """
        mapped_columns = [self.year_column, self.id_column]
        for label in self.settings.labels:
            mapped_column = self._map_id_to_label(label)
            mapped_columns.append(mapped_column)
        year_and_id = [self.settings.year_column_name, self.settings.id_column_name]
        columns = year_and_id + list(self.settings.output_column_names)
        mapping_table = pd.concat(mapped_columns, axis="columns", keys=columns)
        mapping_table = mapping_table.drop_duplicates().set_index(year_and_id)
        return mapping_table

    def _create_code_builder(
        self, household_metadata: dict
    ) -> Callable[[pd.Series], pd.Series]:
        ld_len = household_metadata["ID_Length"]
        attr_dict = household_metadata[self.settings.name]["code"]

        if ("position" in attr_dict) and attr_dict["position"] is not None:
            start, end = attr_dict["position"]["start"], attr_dict["position"]["end"]

            def builder(household_id_column: pd.Series) -> pd.Series:
                return (
                    household_id_column
                    % pow(10, (ld_len - start))
                    // pow(10, (ld_len - end))
                )

        elif "external_file" in attr_dict:
            code_builer_file = external_data.load_table("regions", reset_index=False)
            code_series = code_builer_file.loc[household_metadata["year"]].iloc[:, 0]
            assert isinstance(code_series, pd.Series)
            mapping_dict = code_series.to_dict()

            def builder(household_id_column: pd.Series) -> pd.Series:
                codes = household_id_column.map(mapping_dict)
                assert codes.isna().sum() == 0
                return codes

        else:
            raise ValueError("Code position is not available")

        return builder

    def _create_code_mapper(
        self, label: str, year: int
    ) -> Callable[[pd.Series], pd.Series]:
        household_metadata = utils.resolve_metadata(metadata.household, year)

        if label == "code":
            return self._create_code_builder(household_metadata)

        if not isinstance(household_metadata, dict):
            raise ValueError
        # pylint: disable=unsubscriptable-object
        mapping = household_metadata[self.settings.name][label]
        code_builder = self._create_code_builder(household_metadata)

        def mapper(household_id_column: pd.Series) -> pd.Series:
            mapped = code_builder(household_id_column).map(mapping).astype("category")
            mapped.name = label
            return mapped

        return mapper

    def _map_id_to_label(self, label: str):
        years = self.year_column.drop_duplicates()
        attribute_column = pd.Series(index=self.table.index, dtype="object")
        for year in years:
            filt = self.year_column == year
            attribute_column.loc[filt] = self._create_code_mapper(label, year)(
                self.id_column.loc[filt]
            )
        return attribute_column

    def add_attribute(self):
        """Adds decoded columns to the input table.

        Joins the mapping table to the input table using the ID and
        year columns.

        Returns
        -------
        DataFrame
            Input table with decoded columns added.

        """
        mapping_table = self.construct_mapping_table()
        year_and_id = [self.settings.year_column_name, self.settings.id_column_name]
        self.table = self.table.join(mapping_table, year_and_id)
        return self.table
