import os
from argparse import ArgumentParser, RawTextHelpFormatter
from utils.common import set_seed, linearIdx2matIdx
import numpy as np
from src.dataset import Dataset
import pickle

########################################################################################################################

parser = ArgumentParser(description="Examples: \n", formatter_class=RawTextHelpFormatter)
parser.add_argument(
    '--input_folder', type=str, required=True, help='Path of the dataset folder'
)
parser.add_argument(
    '--output_folder', type=str, required=True, help='Path of the output static dataset folder'
)
parser.add_argument(
    '--seed', type=int, required=False, default=19, help='Seed value'
)
args = parser.parse_args()

########################################################################################################################

input_folder = args.input_folder
output_folder = args.output_folder
seed = args.seed

set_seed(seed=seed)

########################################################################################################################

# Read the dataset
data = Dataset(path=input_folder, normalize=False, verbose=False, seed=seed)

if not os.path.exists(output_folder):
    os.makedirs(output_folder)

with open(os.path.join(output_folder, "unweighted.edges"), 'w') as f:
    for pair in data.get_pairs():
        f.write(f"{pair[0]} {pair[1]}\n")

with open(os.path.join(output_folder, "weighted.edges"), 'w') as f:
    for pair, pair_events in zip(data.get_pairs(), data.get_events()):
        f.write(f"{pair[0]} {pair[1]} {len(pair_events)}\n")