import random
import torch
import numpy as np
import warp as wp

from optimizer import GradientOptimizer
from simulator import MPMSimulator

def set_all_seeds(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

wp.init()
set_all_seeds(42)

if __name__ == '__main__':
    name = 'single_push_rope'
    gripper_type = 'push'

    data_path = f'./data/phystwin_data/{name}'
    config_path = f'./config/{name}_config.json'
    output_path = './output/output_stage1'
    optim_output_path = './optim_stage1'

    simulator = MPMSimulator(data_path, config_path, output_path, gripper_type)
    optimizer = GradientOptimizer(simulator,optim_output_path)
    optimizer.optimize_stage1()
