from models.feedforward import FeedForward, transpose_layer_defs
from models.lgm import LGM, pLGM, sopLGM
from models import cnnvae2


def init_lgm(gm_type, input_size, encoder_layers, latent_size,
             decoder_layers=None):
    
    encoder_layers += [('linear', [2 * latent_size])]
    encoder = FeedForward(input_size, encoder_layers, flatten=False)

    if decoder_layers is None:
        decoder_layers = encoder_layers[:-1]
        decoder_layers.append(('linear', [latent_size]))

        decoder_layers = transpose_layer_defs(decoder_layers, input_size)

    decoder_layers = decoder_layers.copy()
    if isinstance(input_size, (list, tuple)):
        if not input_size[0] == 1 and decoder_layers[-1][0] == 'tconv':
            decoder_layers[-1] = ('tconv', (input_size[0], *decoder_layers[-1][1][1:]))
    else:
        if decoder_layers[-1][0] == 'linear' and not input_size == decoder_layers[-1][1][0]:
            decoder_layers[-1][1][0] = input_size
    decoder = FeedForward(latent_size, decoder_layers, flatten=False)

    return LGM(latent_size, encoder, decoder)


def init_plgm(encoder_name, decoder_name, input_size, latent_size, decoder_params = None):
    
    doubled_input_size = tuple([input_size[0]*2] + input_size[1:])
    list_encoder_name = ['higgins', 'burgess', 'burgess_v2', 'mpcnn', 'mathieu', 'kim']
    if encoder_name in list_encoder_name:
        encoder_layers = (getattr(cnnvae2, encoder_name).encoder_layers).copy()
        encoder_layers += [('linear', [latent_size])]
        encoder = FeedForward(doubled_input_size, encoder_layers, flatten=False)
    elif 'resnet' in encoder_name:
        from models.resnet import resnet18, resnet34, resnet50, resnet101, resnet152
        if encoder_name == 'resnet18':
            model_class = resnet18
        elif encoder_name == 'resnet34':
            model_class = resnet34
        elif encoder_name == 'resnet50':
            model_class = resnet50
        elif encoder_name == 'resnet101':
            model_class = resnet101
        elif encoder_name == 'resnet152':
            model_class = resnet152
        else:
            raise
        encoder = model_class(in_channel = input_size[0] * 2, num_classes = latent_size)        
    else:
        raise
    if not decoder_params:
        decoder_params = {}
    if decoder_name == 'CycleGAN':
        from models.networks_cyclegan import ResnetGenerator
        decoder = ResnetGenerator(input_size[0] + latent_size, input_size[0], output_activation = None, **decoder_params)
    #decoder = FeedForward(latent_size + input_size[0]//2, decoder_layers, flatten=True)

    return pLGM(latent_size, encoder, decoder)


def init_soplgm(encoder_name, decoder_name, input_size, latent_size, so_n_repre, decoder_params = None):
    
    latent_size_full = latent_size + so_n_repre.latent_num
    latent_size_decoder = latent_size + so_n_repre.decoder_num
    doubled_input_size = tuple([input_size[0]*2] + input_size[1:])
    list_encoder_name = ['higgins', 'burgess', 'burgess_v2', 'mpcnn', 'mathieu', 'kim']
    if encoder_name in list_encoder_name:
        encoder_layers = (getattr(cnnvae2, encoder_name).encoder_layers).copy()
        encoder_layers += [('linear', [latent_size_full])]
        encoder = FeedForward(doubled_input_size, encoder_layers, flatten=False)
    else:
        raise
    if not decoder_params:
        decoder_params = {}
    if decoder_name == 'CycleGAN':
        from models.networks_cyclegan import ResnetGenerator
        decoder = ResnetGenerator(input_size[0] + latent_size_decoder, input_size[0], output_activation = None, **decoder_params)
    #decoder = FeedForward(latent_size + input_size[0]//2, decoder_layers, flatten=True)

    return sopLGM(latent_size, so_n_repre, encoder, decoder)
