from defense.jpeg_defense import JPEGDefense
from defense.HiFiC_defense import HiFiCDefense
from defense.ELIC_defense import ELICDefense
from defense.CRDR_defense import CRDRDefense
from defense.hyperprior_defense import HyperpriorDefense
from defense.base_defense import BaseDefense
from defense.diffusion_defense import DiffusionDefense

def create_defense(defense, param):
    try:
        defenseparam = dict([p.split(':') for p in param])
    except ValueError:
        print(f'No parameters passed for {defense} defense. Using default parameters.')
        defenseparam = {}
    print(f'Creating {defense} defense with parameters {defenseparam}')

    if defense == 'jpeg':
        return JPEGDefense(**defenseparam)
    elif defense == 'HiFiC':
        return HiFiCDefense(**defenseparam)
    elif defense == 'ELIC':
        return ELICDefense(**defenseparam)
    elif defense == 'CRDR':
        return CRDRDefense(**defenseparam)
    elif defense == 'hyperprior':
        return HyperpriorDefense(**defenseparam)
    elif defense == 'base':
        return BaseDefense(**defenseparam)
    elif defense == 'diffusion':
        return DiffusionDefense(**defenseparam)

    #############  Add your defense here ##################
    else:
        raise NotImplementedError(f'Defense {defense} not implemented. You can implement a new defense in create_defense.py')