
# Composite Bridge Keys and Composite Tables Generator for the Hypercube
# ---------------------------------------------------------
# This class automates the process of generating composite keys and link tables
# for a set of relational tables
#
# Purpose:
# - When two or more tables are linked by more than one column (they share multiple columns),
#   we need to combine these columns into a single composite key that uniquely represents the relationship.
# - This class detects such multi-column links, generates a composite key by merging the values
#   of the shared columns, and uses this key to efficiently link tables.
# - The idea is to maintain a schema that is suitable for the queries while keeping the underlying data intact.



import pandas as pd
from typing import Dict, Any, List

class CompositeBridgeGenerator:
    def __init__(
        self,
        tables: Dict[str, pd.DataFrame],
        print_info: bool = False
    ) -> None:
        """Initialize the model with tables and create link tables before adding relationships."""
        self.tables: Dict[str, pd.DataFrame] = tables  # all tables (including link tables)
        self.composite_tables: Dict[str, pd.DataFrame] = {}  # the link tables generated by this class
        self.composite_keys: Dict[str, Any] = {}  # the composite keys for each table
        self.column_combinations: List[Any] = []  # Track shared column combinations
        self.key_mapping: Dict[Any, Any] = {}  # Store the mapping of composite keys to auto-incremented keys
        self.column_table_matrix = self._create_column_table_matrix()

        # Identify and create link tables based on shared columns
        if self._create_link_tables():
            if print_info:
                print("CompositeBridge table/s created")
                for table in self.composite_tables:
                    print(' - ', table, '->', self.tables[table].columns.to_list())

            # After creating link tables, apply composite keys
            self._apply_composite_keys()

            # Replace composite keys with auto-incremented keys
            self._replace_with_autonumbered_keys()
        

    def _create_column_table_matrix(self):
        """Create a column-to-table matrix showing which columns belong to which tables."""
        all_columns = set()
        for table_name, table_data in self.tables.items():
            all_columns.update(table_data.columns)

        # Create an MxN matrix with columns as rows and tables as columns
        matrix = pd.DataFrame(0, index=list(all_columns), columns=self.tables.keys())

        # Fill the matrix
        for table_name, table_data in self.tables.items():
            for column in table_data.columns:
                matrix.at[column, table_name] = 1

        return matrix

    def _create_composite_key(self, df, columns):
        """Create a composite key by concatenating the values in the given columns."""
        return df[columns].astype(str).agg('-'.join, axis=1)

    def _create_link_tables(self):
        """Create link tables per shared column-combination (size >= 2)."""
        # 1) Identify shared columns (appear in more than one table)
        shared_columns_matrix = self.column_table_matrix.loc[(self.column_table_matrix.sum(axis=1) > 1)]
        shared_columns = set(shared_columns_matrix.index.tolist())

        # 2) Discover all combos via pairwise table intersections restricted to shared columns
        table_names = list(self.tables.keys())
        table_cols = {t: set(df.columns) for t, df in self.tables.items()}

        combo_set = set()  # set of tuples of sorted column names
        for i in range(len(table_names)):
            for j in range(i + 1, len(table_names)):
                t1, t2 = table_names[i], table_names[j]
                inter = (table_cols[t1] & table_cols[t2]) & shared_columns
                if len(inter) >= 2:
                    combo_set.add(tuple(sorted(inter)))

        # 3) For each combo, gather participants (tables with all combo columns) and create a composite per combo
        self.column_combinations = []
        created_any = False
        for combo in sorted(combo_set):
            combo_list = list(combo)
            participants = [t for t, cols in table_cols.items() if set(combo_list).issubset(cols)]
            if len(participants) < 2:
                continue

            # Track combo once
            if combo_list not in self.column_combinations:
                self.column_combinations.append(combo_list)

            # Build composite df as union of distinct rows of the combo columns
            comp_df = pd.DataFrame(columns=combo_list)
            for t in participants:
                comp_df = pd.concat([comp_df, self.tables[t][combo_list]], ignore_index=True)
            comp_df = comp_df.drop_duplicates()

            # Name includes participants and combo to avoid collisions when same participants share multiple combos
            composite_table_name = "_composite_" + "_".join(sorted(participants)) + "__by__" + "_".join(combo_list)
            # Ensure uniqueness just in case
            base_name = composite_table_name
            suffix = 1
            while composite_table_name in self.tables:
                suffix += 1
                composite_table_name = f"{base_name}__{suffix}"

            self.tables[composite_table_name] = comp_df
            self.composite_tables[composite_table_name] = comp_df
            created_any = True

        return created_any

    def _apply_composite_keys(self):
        """For all tables, create composite keys based on self.column_combinations."""
        # Apply composite keys to all tables based on the stored column combinations
        for table_name, table_data in self.tables.items():
            for columns in self.column_combinations:
                # Check if the table has all the columns in the combination
                if all(col in table_data.columns for col in columns):
                    # Create composite key
                    composite_key_column_name = f"_composite_key_{'_'.join(columns)}"
                    self.tables[table_name][composite_key_column_name] = self._create_composite_key(self.tables[table_name], columns)

                    # Track the composite key for the table
                    if table_name not in self.composite_keys:
                        self.composite_keys[table_name] = []
                    self.composite_keys[table_name].append(composite_key_column_name)

                    # Rename the original shared columns, only from non-link tables
                    if table_name not in self.composite_tables:
                        for column in columns:
                            self.tables[table_name].rename(columns={column: f'{column} <{table_name}>'}, inplace=True)


    def _replace_with_autonumbered_keys(self):
        """Replace composite keys with auto-incremented integer keys."""
        # Create a unique mapping for composite keys to auto-incremented keys
        key_counter = 1  # Start the key counter
        for table_name, composite_key_columns in self.composite_keys.items():
            for composite_key_column in composite_key_columns:
                unique_keys = self.tables[table_name][composite_key_column].unique()

                for key in unique_keys:
                    if key not in self.key_mapping:
                        self.key_mapping[key] = key_counter
                        key_counter += 1

                # Replace the composite keys with the auto-incremented keys
                self.tables[table_name][composite_key_column] = self.tables[table_name][composite_key_column].map(self.key_mapping)


