import mlatom as ml
import os
import torch

torch.cuda.empty_cache()
print(torch.cuda.is_available())
root_dir = os.getcwd()
col = ml.models.methods(program='columbus', command_line_arguments=['-m', '1700'], directory_with_input_files=f'{root_dir}/columbus', save_files_in_current_directory=False)
model_predict_kwargs = {
        'nstates':2,
        'current_state':0,
        'calculate_energy_gradients':[True] * 2
    }


eqmol = ml.molecule.load('eqmol.json',format='json')


init_sampler_kwargs = {
        'eqmol':eqmol,
        'initial_temperature':0,
        'number_of_initial_conditions': 50
    }

sampler_kwargs = {
        'maximum_propagation_time':60.0,
        'time_step':0.1,
        'nstates':2,
        'initial_state':1,
        'prevent_back_hop':False,
        'dump_trajectory_interval':40,
        'reduce_memory_usage':True,
        'min_surfuq_points': 2, # ->
        'ngapmd_trajs':50,
        'max_hopprobuq_points':15,
        'prob_uq_threshold':0.1,
        'plot_populations':False,
        'nthreads':16,
        'stop_uncertain_md':True,
        'initcond_sampler_kwargs':init_sampler_kwargs,
    }
init_dataset_sampler_kwargs = {
    'molecule':eqmol,
    'initial_temperature': 0,
    'number_of_initial_conditions': 50
}
ml_model_kwargs = {'nstates':2,
                   'gap_weight':1.0,
                  'use_pretrained_model':'OMNIP2x_CV_1.pt'} #
ml.al(molecule = eqmol,
      reference_method = col,
      ml_model = 'vecmsani',
      sampler= 'lzsh',
      initdata_sampler_kwargs=init_dataset_sampler_kwargs,
      initdata_sampler='wigner',
      sampler_kwargs = sampler_kwargs,
      refmethod_kwargs=model_predict_kwargs,
      label_nthreads=16,
      ml_model_kwargs= ml_model_kwargs,
      new_points     = 300,
     )
