from .ace_params import data_root_dir, get_ace_file_paths

# exec check
exec_check = {
    "method": "random",
    "metric": {"name": "test_loss_epoch", "goal": "minimize"},
    "parameters": {
        "max_epochs": {"values": [1]},
    }
}


# Generate file paths for different number of months for ace
ace_file_paths_1_month = get_ace_file_paths(data_root_dir, n_months=1)
ace_file_paths_3_months = get_ace_file_paths(data_root_dir, n_months=3)
ace_file_paths_12_months = get_ace_file_paths(data_root_dir, n_months=12)
}

tet_with_ortho_weight = { 
    "method": "grid",
    "metric": {"name": "test_loss_epoch", "goal": "minimize"},
    "parameters" : {
        "seed" : {"values" : [1]},
        "datamodule_args": {
          "parameters": {
            "mode" : {"values" : ["spatio_temporal_interpolation"]},
            "file_paths" : {"value" : ace_file_paths_12_months},
            "train_fraction": {"values": [0.01]},
            "val_fraction": {"values": [0.01]},
            "test_fraction": {"values": [0.01]},
          }
        },
        "locationencoder_args": {
            "parameters": {
                "ortho_weight": {"values": [
                  0,
                  1e0
                  ]},
                "time_embedding_type": {"values": 
                                        [
                                            "legendre",
                                            "fourier",
                                            "constant",
                                            "triangle",
                                            "monomial",
                                            "time_copy",
                                        ]
                    },
                "time_embedding_dim": {"values": [40]},
                "legendre_polys": {"values": [20]},
                "normality_flag": {"values": [False]},
                },
            },
        "max_epochs": {"values": [
            500
        ]},
        "arch": {
          "values": [
            "baseline_arch_v1_large",
          ]
        }
    }
}

ace_metric = {
    "name": "test_RMSE_epoch",
    "goal": "minimize",
}

ace_config  = tet_with_ortho_weight

ace_config["metric"]