from algorithm.hdpview import HDPView

"""
For data synthesis task.
"""

class HDPViewSampler:
    
    @classmethod
    def make_model(cls, prng):
        return cls(prng)
    
    def __init__(self, prng):
        self.prng = prng
        
    def train(self, ct, epsilon, label_attr, ratio=0.9, alpha=1.6, beta=1.2, gamma=0.9):
        is_classification = True
        p_view, _ = HDPView(ct, epsilon, self.prng, ratio, alpha, beta, gamma, is_classification, label_attr)
        self.p_view = p_view
    
    def generate_data(self, n_data):
        data = self.p_view.random_sample(self.prng, n_data)
        return data
    

class NonDPHistogram:
    
    @classmethod
    def make_model(cls, prng):
        return cls(prng)
    
    def __init__(self, prng):
        self.prng = prng
        
    def train(self, ct):
        distribution = ct.pd_table.values.flatten() / ct.pd_table.values.sum()
        self.data = ct.pd_table.index
        self.distribution = distribution
    
    def generate_data(self, size):
        indices = self.prng.choice(a=range(len(self.distribution)), size=size, p=self.distribution, replace=True)
        return self.data[indices].to_frame(index=False)
from algorithm.hdpview import HDPView

"""
For data synthesis task.
"""

class HDPViewSampler:
    
    @classmethod
    def make_model(cls, prng):
        return cls(prng)
    
    def __init__(self, prng):
        self.prng = prng
        
    def train(self, ct, epsilon, label_attr, ratio=0.9, alpha=1.6, beta=1.2, gamma=0.9):
        is_classification = True
        p_view, _ = HDPView(ct, epsilon, self.prng, ratio, alpha, beta, gamma, is_classification, label_attr)
        self.p_view = p_view
    
    def generate_data(self, n_data):
        data = self.p_view.random_sample(self.prng, n_data)
        return data
    

class NonDPHistogram:
    
    @classmethod
    def make_model(cls, prng):
        return cls(prng)
    
    def __init__(self, prng):
        self.prng = prng
        
    def train(self, ct):
        distribution = ct.pd_table.values.flatten() / ct.pd_table.values.sum()
        self.data = ct.pd_table.index
        self.distribution = distribution
    
    def generate_data(self, size):
        indices = self.prng.choice(a=range(len(self.distribution)), size=size, p=self.distribution, replace=True)
        return self.data[indices].to_frame(index=False)
For data synthesis task.