import numpy as np

DATASETS = {
    "cifa": {
        "img_height": 512,
        "img_width": 512,
        "img_channels": 1,
        "n_classes": 2,
    },
    "lizard": {
        "img_height": 512,
        "img_width": 512,
        "img_channels": 3,
        "n_classes": 2,
    },
    "vaid": {
        "img_height": 512,
        "img_width": 512,
        "img_channels": 3,
        "n_classes": 3,
    },
}

MODELS = {
    "cifa": {
        "path": "models/cifa/LMG_unet_cifa_20233210.h5",
        "counterpoint_param_grid": {
            "filter_size": np.linspace(12, 24, 7, dtype="int"),
            "min_distance": np.linspace(3, 15, 17),
            "threshold": np.linspace(0.3, 0.6, 16),
        },
        "counterpoint_params": {
            "class_0": {"filter_size": 12, "min_distance": 3.0, "threshold": 0.52},
            "class_1": {"filter_size": 12, "min_distance": 3.0, "threshold": 0.30},
        },
    },
    "lizard": {
        "path": "models/lizard/LMG_unet_lizard_20234715.h5",
        "counterpoint_param_grid": {
            "filter_size": np.linspace(2, 10, 9, dtype="int"),
            "min_distance": np.linspace(2, 6, 11),
            "threshold": np.linspace(0.2, 0.6, 21),
        },
        "counterpoint_params": {
            "class_0": {"filter_size": 5, "min_distance": 3.2, "threshold": 0.50},
            "class_1": {"filter_size": 3, "min_distance": 2.4, "threshold": 0.36},
        },
    },
    "vaid": {
        "path": "models/vaid/LMG_unet_vaid_202321113.h5",
        "counterpoint_param_grid": {
            "filter_size": np.linspace(4, 24, 11, dtype="int"),
            "min_distance": np.linspace(5, 15, 17),
            "threshold": np.linspace(0.2, 0.6, 21),
        },
        "counterpoint_params": {
            "class_0": {
                "filter_size": 4,
                "min_distance": 5.0,
                "threshold": 0.46,
            },
            "class_1": {
                "filter_size": 4,
                "min_distance": 5.0,
                "threshold": 0.40,
            },
            "class_2": {"filter_size": 6, "min_distance": 5.0, "threshold": 0.42},
        },
    },
}
