import torchvision
torchvision.disable_beta_transforms_warning()

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_stage2'
    optim_output_path = './optim_stage2'

    optimized_stage1_params_path = None

    simulator = MPMSimulator(data_path, config_path, output_path, gripper_type, True, 'cuda:0',
                             optimized_stage1_params_path)
    optimizer = GradientOptimizer(simulator,optim_output_path)
    optimizer.optimize_stage1()

