doczyai-pipelines/fieldExtraction/crosswalk/crosswalk_utils.py

import csv
import json

import pandas as pd
from src.config import STATE


class CrosswalkBuilder:
    def __init__(self):
        self.mapping: dict[str, str] = {}
        self.metadata: dict[str, str | None] = {}

    def from_excel(
        self, path: str, from_col: str, to_col: str, sheet_name: str | None = None
    ) -> "CrosswalkBuilder":
        """Load mapping from an excel file

        Args:
            path (str): location for the excel file
            from_col (str): Source column (for keys)
            to_col (str): Destination column (for values)
            sheet_name (str|None, optional): Sheet name in the excel file. Defaults to None.

        Returns:
            CrosswalkBuilder: the builder instance for method chaining
        """
        if ".xlsx" in path:
            if sheet_name is None:
                # If sheet_name is not provided, read the first sheet
                df = pd.read_excel(path, sheet_name=0)
            else:
                df = pd.read_excel(path, sheet_name=sheet_name)
        elif ".csv" in path:
            df = pd.read_csv(path)
        else:
            raise ValueError(f"Extension '{path.split(".")[-1]}' not a valid file extension.")

        if from_col not in df.columns:
            raise ValueError(f"Column '{from_col}' not found in the excel file")
        elif to_col not in df.columns:
            raise ValueError(f"Column '{to_col}' not found in the excel file")

        self.metadata = {
            "from_col": from_col,
            "to_col": to_col,
        }

        self.mapping = dict(zip(df[from_col], df[to_col]))

        return self

    def from_dict(
        self,
        mapping: dict[str, str],
        from_col: str | None = None,
        to_col: str | None = None,
    ) -> "CrosswalkBuilder":
        """Load mapping from a dictionary

        Args:
            mapping (dict[str, str]): Mapping dictionary

        Returns:
            CrosswalkBuilder: the builder instance for method chaining
        """
        self.metadata = {
            "from_col": from_col,
            "to_col": to_col,
        }
        self.mapping = mapping
        return self

    def from_df(
        self,
        df: pd.DataFrame,
        from_col: str,
        to_col: str
    ) -> "CrosswalkBuilder":
        """Load mapping from a DataFrame

        Args:
            df (pd.DataFrame): DataFrame containing the mapping
            from_col (str): Column name for keys
            to_col (str): Column name for values

        Returns:
            CrosswalkBuilder: the builder instance for method chaining
        """
        self.metadata = {
            "from_col": from_col,
            "to_col": to_col,
        }
        self.mapping = dict(zip(df[from_col], df[to_col]))
        return self


    def from_json(self, path: str) -> "CrosswalkBuilder":
        """Load mapping from a JSON file

        Args:
            path (str): Path to the JSON file

        Returns:
            CrosswalkBuilder: the builder instance for method chaining
        """
        with open(path, "r") as f:
            data = json.load(f)

        self.metadata = data.get("metadata", {})

        mapping = data.get("mapping", {})
        for state in STATE:
            mapping.update(data.get("state_mapping", {}).get(state, {}))

        self.mapping = mapping

        return self

    def save(self, output_path: str) -> None:
        """Save the crosswalk to JSON.

        Args:
            output_path (str): Path to save the crosswalk
        """
        output_dict = {"metadata": self.metadata, "mapping": self.mapping}
        with open(output_path, "w") as f:
            json.dump(output_dict, f, indent=4)

    def to_csv(self, output_path: str) -> None:
        """Save the crosswalk to CSV.

        Args:
            output_path (str): Path to save the crosswalk
        """
        from_col = (
            self.metadata.get("from_col") if self.metadata.get("from_col") else "source"
        )
        to_col = (
            self.metadata.get("to_col") if self.metadata.get("to_col") else "target"
        )

        df = pd.DataFrame(self.mapping.items(), columns=[from_col, to_col])
        df.to_csv(output_path, index=False, quoting=csv.QUOTE_ALL, encoding="utf-8-sig")

    def create_reverse_mapping(self) -> dict[str, str]:
        """Create a reverse mapping where values that map to the same key are concatenated

        Example:
            {"A": "1", "B": "1", "C": "2"} becomes {"1": "A, B", "2": "C"}

        Returns:
            dict[str, str]: Reverse mapping with concatenated values
        """
        reverse_dict: dict[str, list[str]] = {}

        # First collect all sources for each target
        for source, target in self.mapping.items():
            if target in reverse_dict:
                reverse_dict[target].append(source)
            else:
                reverse_dict[target] = [source]

        # Then join the lists with commas
        return {target: ", ".join(sources) for target, sources in reverse_dict.items()}


def apply_crosswalk(
    val: str, mapping: dict[str, str], default: str | None = None
) -> str:
    """Apply a crosswalk to a value

    Args:
        val (str): Input value
        mapping (dict[str, str]): Mapping dictionary
        default (str | None, optional): Default value; if None and val is not found in the mapping, val is returned. Defaults to None.

    Returns:
        str: Crosswalked value
    """
    if "," in val:
        values = val.split(",") + [val]  # Split by comma if present
    else:
        values = [val]
    mapped_values = [mapping.get(v.strip(), default if default is not None else v.strip()) for v in values]
    return ",".join(mapped_values)  # Join back into a single string