from modules import network

from .resnet import ResNetGenerator
from .densenet import DenseNetGenerator
from .pixel_wise_resnet import PixelWiseResNetGenerator
from .channel_wise_resnet import ChannelWiseResNetGenerator
from .global_wise_resnet import GlobalWiseResNetGenerator


def create_generator(generator_type, input_nc, output_nc, generator_channels_base, norm, num_downs, init_type, init_gain, gpu_ids):
    if generator_type == 'resnet':
        net = ResNetGenerator(
            input_nc=input_nc,
            output_nc=output_nc,
            generator_channels_base=generator_channels_base,
            norm_layer=network.get_norm_layer(norm_type=norm),
            num_downs=num_downs
        )
    elif generator_type == 'densenet':
        net = DenseNetGenerator(
            input_nc=input_nc,
            output_nc=output_nc,
            generator_channels_base=generator_channels_base,
            norm_layer=network.get_norm_layer(norm_type=norm),
            num_downs=num_downs
        )
    elif generator_type == 'pixel-wise-resnet':
        net = PixelWiseResNetGenerator(
            input_nc=input_nc,
            output_nc=output_nc,
            generator_channels_base=generator_channels_base,
            norm_layer=network.get_norm_layer(norm_type=norm),
            num_downs=num_downs
        )
    elif generator_type == 'channel-wise-resnet':
        net = ChannelWiseResNetGenerator(
            input_nc=input_nc,
            output_nc=output_nc,
            generator_channels_base=generator_channels_base,
            norm_layer=network.get_norm_layer(norm_type=norm),
            num_downs=num_downs
        )
    elif generator_type == 'global-wise-resnet':
        net = GlobalWiseResNetGenerator(
            input_nc=input_nc,
            output_nc=output_nc,
            generator_channels_base=generator_channels_base,
            norm_layer=network.get_norm_layer(norm_type=norm),
            num_downs=num_downs
        )
    return network.init_net(net, init_type, init_gain, gpu_ids)
