import pathlib
import csv
import numpy as np

def shifted_geometric_mean(values, shift):
    values = np.maximum(1, np.array(values) + shift)
    sgm = np.exp(np.sum(np.log(values)) / len(values)) - shift
    return sgm

def shifted_geometric_std(values, shift, sgm):
    values = np.maximum(1, np.array(values) + shift)
    sg_std = np.exp(np.sum((np.log(values) - np.log(sgm)) ** 2) / len(values))
    return sg_std
    
if __name__ == "__main__":
    task = 'dual'
    filename = 'hard_set_cover_all'
    files = [
        pathlib.Path(f"submissions/{task}/rl/{filename}.csv"),
        pathlib.Path(f"submissions/{task}/gcn/{filename}.csv"),
        pathlib.Path(f"submissions/{task}/expert/{filename}.csv"),
    ]
    
    # Extract wallclock times from results and compute shifted geometric mean
    sgms = []
    sgdev = []
    for file in files:
        wallclock_times = []
        with open(file, mode='r') as csv_file:
            reader = csv.DictReader(csv_file)
            for row in reader:
                wallclock_times.append(float(row['time']))

        # Calculate SGM and scale
        sgm = shifted_geometric_mean(wallclock_times, shift=10)
        sg_std = shifted_geometric_std(wallclock_times, shift=10, sgm=sgm)

        print(f"Shifted Geometric Mean: {sgm:.2f}")
        print(f"Shifted Geometric Standard Deviation: {sg_std:.2f}")
        sgms.append(sgm)
        sgdev.append(sg_std)