import pandas as pd
import numpy as np
from data_utils.data_simulation import UpliftSimulationReady

class EDASimulationReady:

    def __init__(self, files_path):
        self.files_path = files_path

    def load_conversions(self, file_name):

        uplift_simulation = UpliftSimulationReady(self.files_path)
        df = uplift_simulation.load_uplift_data(file_name)

        sum_conversions = df.pivot_table(values=['conversion','discounted_price','benefit'],
                index='treatment_group_key',
                aggfunc=[np.sum],
                margins=False)
        
        mean_conversions = df.pivot_table(values=['conversion','discounted_price','benefit'],
                index='treatment_group_key',
                aggfunc=[np.mean],
                margins=False)
        
        # save to csv
        sum_conversions.to_csv(self.files_path + 'sum_conversions.csv')
        mean_conversions.to_csv(self.files_path + 'mean_conversions.csv')

        return sum_conversions, mean_conversions