import gradio as gr
from modules import shared


section = ('segment_anything', "Segment Anything")
prefix_id = 'segment_anything_cn_preprocessor_'
prefix_label = "CN Preprocessor: "


def getSegmentAnythingModel(sam_model_list):
    if not sam_model_list:
        raise Exception("There are no sam models")

    res : str = shared.opts.data.get(prefix_id + 'sam_model', 'sam_hq_vit_l.pth')
    if res not in sam_model_list:
        res = sam_model_list[0]
    return res


def getAutoSamOptions():
    res = []
    def append(id, default):
        value = shared.opts.data.get(prefix_id + id, default)
        if value % 1 == 0:
            value = int(value)
        res.append(value)
    append('points_per_side', 32)
    append('points_per_batch', 64)
    append('pred_iou_thresh', 0.88)
    append('stability_score_thresh', 0.95)
    append('stability_score_offset', 1)
    append('box_nms_thresh', 0.7)
    append('crop_n_layers', 0)
    append('crop_nms_thresh', 0.7)
    append('crop_overlap_ratio', 512/1500)
    append('crop_n_points_downscale_factor', 1)
    append('min_mask_region_area', 0)
    return res


def needAutoUnloadModels():
    opt = shared.opts.data.get(prefix_id + "_always_unload_models", 'Enabled')

    if opt == 'Enabled':
        return True
    if opt == 'Disabled':
        return False
    if opt == 'Only SDXL':
        return shared.sd_model.is_sdxl

    return shared.cmd_opts.lowvram or shared.cmd_opts.medvram or (shared.sd_model.is_sdxl and shared.cmd_opts.medvram_sdxl)


def avoidOOM():
    res : bool = shared.opts.data.get(prefix_id + 'avoid_oom', True)
    return res

def getTemplate(sam_model_list):
    if not sam_model_list:
        sam_model_list = ['not found']
    defaultModel = 'sam_hq_vit_l.pth'
    if defaultModel not in sam_model_list:
        defaultModel = sam_model_list[0]

    options = {
        prefix_id + 'sam_model': shared.OptionInfo(
            defaultModel,
            prefix_label + 'segment anything model',
            gr.Dropdown,
            {
                'choices' : sam_model_list,
            },
            section=section,
        ),
        prefix_id + 'avoid_oom': shared.OptionInfo(
            True,
            prefix_label + 'Try to avoid OOM',
            section=section,
        ).info("Automatic SAM segmentation requires much more VRAM them regular segment anything. "
               "With this option in case of OOM it will try to unload SD model, if don't help, use CPU"),
        prefix_id + "always_unload_models": shared.OptionInfo(
            'Enabled',
            prefix_label + 'Always unload models',
            gr.Radio,
            {
                'choices' : ['Automatic', 'Enabled', 'Only SDXL', 'Disabled'],
            },
            section=section,
        ).info("Automatic means enable only for --lowvram and --medvram mode. "),
    }

    def addNumberOption(label: str, value):
        args = {}
        if value % 1 == 0:
            args['step'] = 1
        options[prefix_id + label] = shared.OptionInfo(
            value,
            prefix_label + label,
            gr.Number,
            args,
            section=section,
        )

    addNumberOption('points_per_side', 32)
    addNumberOption('points_per_batch', 64)
    addNumberOption('pred_iou_thresh', 0.88)
    addNumberOption('stability_score_thresh', 0.95)
    addNumberOption('stability_score_offset', 1)
    addNumberOption('box_nms_thresh', 0.7)
    addNumberOption('crop_n_layers', 0)
    addNumberOption('crop_nms_thresh', 0.7)
    addNumberOption('crop_overlap_ratio', 512/1500)
    addNumberOption('crop_n_points_downscale_factor', 1)
    addNumberOption('min_mask_region_area', 0)

    return shared.options_section(section, options)