from sdv.metadata import MultiTableMetadata
from sdv.multi_table.base import BaseMultiTableSynthesizer
from sdv.single_table import CopulaGANSynthesizer
from sdv.sampling.independent_sampler import BaseIndependentSampler


class IndependentSynthesizer(BaseIndependentSampler, BaseMultiTableSynthesizer):
    def __init__(self, metadata: MultiTableMetadata, epochs = {}):
        BaseMultiTableSynthesizer.__init__(self, metadata)
        self._table_sizes = {}
        self._table_synthesizers = {
            t: CopulaGANSynthesizer(c, epochs=epochs.get(t, 300)) for t, c in metadata.tables.items()
        }
        BaseIndependentSampler.__init__(self, metadata, self._table_synthesizers, self._table_sizes)
        self._fitted = False

    def fit(self, tables):
        for t, s in self._table_synthesizers.items():
            s.fit(tables[t])
            self._table_sizes[t] = tables[t].shape[0]
        self._fitted = True

    def _add_foreign_key_columns(self, child_table, parent_table, child_name, parent_name):
        fk = None
        for rel in self.metadata.relationships:
            if rel["parent_table_name"] == parent_name and rel["child_table_name"] == child_name:
                fk = rel
                break
        if fk is None:
            raise ValueError()
        ccol = fk["child_foreign_key"]
        if ccol not in child_table.columns:
            child_table[ccol] = parent_table[fk["parent_primary_key"]].sample(
                n=child_table.shape[0], replace=True
            ).reset_index(drop=True)