from beartype.typing import Callable
from beartype import beartype

from setting import Settings

@beartype
def get_check_abstractor_params() -> dict:
    return {
        'crown_batch_size': Settings.backward_batch_size,
        'forward_max_dim': Settings.forward_max_dim,
        'dynamic_forward': Settings.forward_dynamic,
        'optimize_bound_args': {
            'iteration': 2,
            'stop_criterion_func': lambda x: False,
            'enable_alpha_crown': True,
            'enable_beta_crown': False,
            'use_shared_alpha': Settings.share_alphas, 
            'init_alpha': False,
            'fix_interm_bounds': True,
        }
    }

@beartype
def get_branching_opt_params() -> dict:
    return {
        'crown_batch_size': Settings.backward_batch_size,
        'optimize_bound_args': {
            'enable_beta_crown': False, 
            'fix_interm_bounds': True,
        }
    }   
    
    
@beartype
def get_initialize_opt_params(stop_criterion_func: Callable) -> dict:
    return {
        'crown_batch_size': Settings.backward_batch_size,
        'optimize_bound_args': {
            'enable_alpha_crown': True,
            'enable_beta_crown': False, 
            'use_shared_alpha': Settings.share_alphas, 
            'init_alpha': False,
            'fix_interm_bounds': True,
            'stop_criterion_func': stop_criterion_func,
            'iteration': 100, 
            'lr_alpha': 0.1, 
            'lr_decay': 0.98, 
        }
    }
    
    
@beartype
def get_beta_opt_params(stop_criterion_func: Callable) -> dict:
    return {
        'crown_batch_size': Settings.backward_batch_size,
        'optimize_bound_args': {
            'enable_alpha_crown': True,
            'enable_beta_crown': True, 
            'use_shared_alpha': Settings.share_alphas, 
            'fix_interm_bounds': True, 
            'iteration': 20,
            'lr_alpha': 0.1, 
            'lr_beta': 0.1,
            'lr_decay': 0.98, 
            'stop_criterion_func': stop_criterion_func,
        }
    }

    
@beartype
def get_input_opt_params(stop_criterion_func: Callable) -> dict:
    return {
        'crown_batch_size': Settings.backward_batch_size,
        'forward_max_dim': Settings.forward_max_dim,
        'dynamic_forward': Settings.forward_dynamic,
        'optimize_bound_args': {
            'enable_beta_crown': False, 
            'fix_interm_bounds': True, 
            'iteration': 10,
            'lr_alpha': 0.1, 
            'lr_decay': 0.98, 
            'stop_criterion_func': stop_criterion_func,
        }
    }
