# : generate_data.py ()
import os
import sys
import torch
import numpy as np
import argparse 

# =============================================================================
PROJECT_ROOT = '' # <<< MODIFIED
# =============================================================================
sys.path.append(PROJECT_ROOT)
from LIN.utils import getData

def generate_custom_data(seed): # <<< MODIFIED: Accept seed as an argument
    """
     Edgeexist 。
    """
    # <<< MODIFIED: Set seed at the beginning of the function
    torch.manual_seed(seed)
    np.random.seed(seed)

    # 1.  ()
    d = 5
    P = 2
    # e_den=4*2/(P*(d+1))
    e_den=0.5
    T = 1000
    K_true = 3
    num_tgt_per_intervention = 2
    graph_type = 'ER'
    
    print(f" Edgeexist ...")
    print(f": d={d}, P={P}, K_true={K_true}, T={T}, seed={seed}")

    # 2.  ()
    save_path = './data'
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    
    # 3. 
    # <<< MODIFIED: Include seed in the filename to avoid overwriting (optional but good practice)
    file_name_prefix = f"synth_d{d}_P{P}_T{T}_K{K_true}_seed{seed}"
    
    # 4. 
    G, K, Data_generator = getData(
        P=P, 
        d=d, 
        e_den=e_den,
        num_additional_interv=K_true - 1, 
        graph_type=graph_type,
        num_tgt=num_tgt_per_intervention
    )
    
    # 5.  ()
    weight = torch.ones(K_true).float()
    data, intervention_type = Data_generator.generate(T=T, weight=weight)
    
    # 6. 
    torch.save(G, os.path.join(save_path, file_name_prefix + '_G.pt'))
    torch.save(data, os.path.join(save_path, file_name_prefix + '_data.pt'))
    torch.save(intervention_type, os.path.join(save_path, file_name_prefix + '_intvs.pt'))
    
    print(f"\n: {os.path.abspath(save_path)}")
    print("-" * 40)
    # <<< MODIFIED: Return the filename prefix so run_single_trial knows which data to load
    return file_name_prefix

if __name__ == '__main__':
    # <<< MODIFIED: Setup argument parser to receive the seed >>>
    parser = argparse.ArgumentParser(description="Generate synthetic data for experiments.")
    parser.add_argument('--seed', type=int, default=42, help='Random seed for data generation.')
    args = parser.parse_args()

    print(f" {args.seed} 。")
    generate_custom_data(seed=args.seed)

    print("。")
