import os
import json

os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15
import tensorflow as tf
from tensorflow.keras.layers import (
    Activation,
    Add,
    AveragePooling2D,
    BatchNormalization,
    Bidirectional,
    Conv1D,
    Conv2D,
    Dense,
    Dropout,
    Embedding,
    Flatten,
    Input,
    Layer,
    LayerNormalization,
    LSTM,
    MaxPooling2D,
    MultiHeadAttention,
    Permute,
    Reshape,
    UpSampling2D,
    ZeroPadding2D,
    StringLookup,
    add,
    concatenate
)
from tensorflow.keras.models import Model
from tensorflow.keras.regularizers import l2
from tensorflow.keras.backend import ctc_batch_cost

from ..patch_encoder import Patches, PatchEncoder

##mlp_head_units = [512, 256]#[2048, 1024]
###projection_dim = 64
##transformer_layers = 2#8
##num_heads = 1#4
RESNET50_WEIGHTS_PATH = './pretrained_model/resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5'
RESNET50_WEIGHTS_URL = ('https://github.com/fchollet/deep-learning-models/releases/download/v0.2/'
                        'resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5')

IMAGE_ORDERING = 'channels_last'
MERGE_AXIS = -1


class CTCLayer(Layer):
    def call(self, y_true, y_pred):
        batch_len = tf.cast(tf.shape(y_true)[0], dtype="int64")
        input_length = tf.cast(tf.shape(y_pred)[1], dtype="int64")
        label_length = tf.cast(tf.shape(y_true)[1], dtype="int64")

        input_length = input_length * tf.ones(shape=(batch_len, 1), dtype="int64")
        label_length = label_length * tf.ones(shape=(batch_len, 1), dtype="int64")
        loss = ctc_batch_cost(y_true, y_pred, input_length, label_length)
        self.add_loss(loss)

        # At test time, just return the computed predictions.
        return y_pred

class CTCDecoder(Layer):
    def call(self, inputs):
        n_samples = tf.shape(inputs)[0]
        n_steps = inputs.shape[1]
        n_classes = inputs.shape[2]
        lengths = tf.ones(n_samples, dtype=tf.int32) * n_steps
        ## Keras beam search seems to mess with double letters
        ## but Keras greedy sometimes removes arbitrary letters
        # outputs, logits = tf.keras.backend.ctc_decode(inputs,
        #                                               lengths,
        #                                               beam_width=20,
        #                                               greedy=False, # True,
        #                                               # backend does not allow these kwargs
        #                                               #merge_repeated=False,
        #                                               #mask_index=inputs.shape[2]-1,
        # )
        # tf.nn.ctc_*_decoder (in contrast to tf.keras.backend.ctc_decode)
        # needs logits instead of probs and time-major (batch 2nd dim)
        inputs = tf.math.log(
            tf.transpose(inputs, perm=[1, 0, 2]) + tf.keras.backend.epsilon()
        )
        # tf.compat.v1.nn.ctc_beam_search_decoder() also needs merge_repeated=False
        # tf.nn.ctc_beam_search_decoder() is not supported by ONNX, yet
        # tf.nn.ctc_greedy_decoder() is not as precise, though:
        decoded, logits = tf.nn.ctc_beam_search_decoder(
            inputs,
            lengths,
            beam_width=10,
            top_paths=1
        )
        # get top path for all sequences in batch
        decoded = decoded[0]
        logits = logits[:, 0]
        probs = tf.exp(-logits / n_steps)
        # convert to dense
        outputs = tf.SparseTensor(decoded.indices, decoded.values,
                                  (n_samples, n_steps))
        outputs = tf.sparse.to_dense(sp_input=outputs, default_value=n_classes-1)
        # # drop non-tokens (-1) and OOV (0)
        # result = []
        # for output in outputs:
        #     result.append(tf.gather(output, tf.where(output > 0)))
        # outputs = tf.stack(result)
        return outputs, probs
    
def mlp(x, hidden_units, dropout_rate):
    for units in hidden_units:
        x = Dense(units, activation=tf.nn.gelu)(x)
        x = Dropout(dropout_rate)(x)
    return x

def one_side_pad(x):
    x = ZeroPadding2D(((1, 0), (1, 0)), data_format=IMAGE_ORDERING)(x)
    return x

def identity_block(input_tensor, kernel_size, filters, stage, block):
    """The identity block is the block that has no conv layer at shortcut.
    # Arguments
        input_tensor: input tensor
        kernel_size: defualt 3, the kernel size of middle conv layer at main path
        filters: list of integers, the filterss of 3 conv layer at main path
        stage: integer, current stage label, used for generating layer names
        block: 'a','b'..., current block label, used for generating layer names
    # Returns
        Output tensor for the block.
    """
    filters1, filters2, filters3 = filters

    if IMAGE_ORDERING == 'channels_last':
        bn_axis = 3
    else:
        bn_axis = 1

    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    x = Conv2D(filters1, (1, 1), data_format=IMAGE_ORDERING, name=conv_name_base + '2a')(input_tensor)
    x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x)
    x = Activation('relu')(x)

    x = Conv2D(filters2, kernel_size, data_format=IMAGE_ORDERING,
               padding='same', name=conv_name_base + '2b')(x)
    x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x)
    x = Activation('relu')(x)

    x = Conv2D(filters3, (1, 1), data_format=IMAGE_ORDERING, name=conv_name_base + '2c')(x)
    x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x)

    x = add([x, input_tensor])
    x = Activation('relu')(x)
    return x


def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)):
    """conv_block is the block that has a conv layer at shortcut
    # Arguments
        input_tensor: input tensor
        kernel_size: defualt 3, the kernel size of middle conv layer at main path
        filters: list of integers, the filterss of 3 conv layer at main path
        stage: integer, current stage label, used for generating layer names
        block: 'a','b'..., current block label, used for generating layer names
    # Returns
        Output tensor for the block.
    Note that from stage 3, the first conv layer at main path is with strides=(2,2)
    And the shortcut should have strides=(2,2) as well
    """
    filters1, filters2, filters3 = filters

    if IMAGE_ORDERING == 'channels_last':
        bn_axis = 3
    else:
        bn_axis = 1

    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    x = Conv2D(filters1, (1, 1), data_format=IMAGE_ORDERING, strides=strides,
               name=conv_name_base + '2a')(input_tensor)
    x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x)
    x = Activation('relu')(x)

    x = Conv2D(filters2, kernel_size, data_format=IMAGE_ORDERING, padding='same',
               name=conv_name_base + '2b')(x)
    x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x)
    x = Activation('relu')(x)

    x = Conv2D(filters3, (1, 1), data_format=IMAGE_ORDERING, name=conv_name_base + '2c')(x)
    x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x)

    shortcut = Conv2D(filters3, (1, 1), data_format=IMAGE_ORDERING, strides=strides,
                      name=conv_name_base + '1')(input_tensor)
    shortcut = BatchNormalization(axis=bn_axis, name=bn_name_base + '1')(shortcut)

    x = add([x, shortcut])
    x = Activation('relu')(x)
    return x

def resnet50(inputs, weight_decay=1e-6, pretraining=False):
    if IMAGE_ORDERING == 'channels_last':
        bn_axis = 3
    else:
        bn_axis = 1

    x = ZeroPadding2D((3, 3), data_format=IMAGE_ORDERING)(inputs)
    x = Conv2D(64, (7, 7), data_format=IMAGE_ORDERING, strides=(2, 2), kernel_regularizer=l2(weight_decay),
               name='conv1')(x)
    f1 = x

    x = BatchNormalization(axis=bn_axis, name='bn_conv1')(x)
    x = Activation('relu')(x)
    x = MaxPooling2D((3, 3), data_format=IMAGE_ORDERING, strides=(2, 2))(x)

    x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1))
    x = identity_block(x, 3, [64, 64, 256], stage=2, block='b')
    x = identity_block(x, 3, [64, 64, 256], stage=2, block='c')
    f2 = one_side_pad(x)

    x = conv_block(x, 3, [128, 128, 512], stage=3, block='a')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='b')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='c')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='d')
    f3 = x

    x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f')
    f4 = x

    x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a')
    x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b')
    x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c')
    f5 = x

    if pretraining:
        model = Model(inputs, x).load_weights(RESNET50_WEIGHTS_PATH)

    return f1, f2, f3, f4, f5

def unet_decoder(img, f1, f2, f3, f4, f5, n_classes, light=False, task="segmentation", weight_decay=1e-6):
    if IMAGE_ORDERING == 'channels_last':
        bn_axis = 3
    else:
        bn_axis = 1

    o = Conv2D(512 if light else 1024, (1, 1), padding='same',
               data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(f5)
    o = BatchNormalization(axis=bn_axis)(o)
    o = Activation('relu')(o)

    if light:
        f4 = Conv2D(512, (1, 1), padding='same',
                    data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(f4)
        f4 = BatchNormalization(axis=bn_axis)(f4)
        f4 = Activation('relu')(f4)

    o = UpSampling2D((2, 2), data_format=IMAGE_ORDERING, interpolation="bilinear")(o)
    o = concatenate([o, f4], axis=MERGE_AXIS)
    o = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(o)
    o = Conv2D(512, (3, 3), padding='valid',
               data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o)
    o = BatchNormalization(axis=bn_axis)(o)
    o = Activation('relu')(o)

    o = UpSampling2D((2, 2), data_format=IMAGE_ORDERING, interpolation="bilinear")(o)
    o = concatenate([o, f3], axis=MERGE_AXIS)
    o = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(o)
    o = Conv2D(256, (3, 3), padding='valid',
               data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o)
    o = BatchNormalization(axis=bn_axis)(o)
    o = Activation('relu')(o)

    o = UpSampling2D((2, 2), data_format=IMAGE_ORDERING, interpolation="bilinear")(o)
    o = concatenate([o, f2], axis=MERGE_AXIS)
    o = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(o)
    o = Conv2D(128, (3, 3), padding='valid',
               data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o)
    o = BatchNormalization(axis=bn_axis)(o)
    o = Activation('relu')(o)

    o = UpSampling2D((2, 2), data_format=IMAGE_ORDERING, interpolation="bilinear")(o)
    o = concatenate([o, f1], axis=MERGE_AXIS)
    o = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(o)
    o = Conv2D(64, (3, 3), padding='valid',
               data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o)
    o = BatchNormalization(axis=bn_axis)(o)
    o = Activation('relu')(o)

    o = UpSampling2D((2, 2), data_format=IMAGE_ORDERING, interpolation="bilinear")(o)
    o = concatenate([o, img], axis=MERGE_AXIS)
    o = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(o)
    o = Conv2D(32, (3, 3), padding='valid',
               data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o)
    o = BatchNormalization(axis=bn_axis)(o)
    o = Activation('relu')(o)

    o = Conv2D(n_classes, (1, 1), padding='same',
               data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o)
    if task == "segmentation":
        o = BatchNormalization(axis=bn_axis)(o)
        o = Activation('softmax')(o)
    else:
        o = Activation('sigmoid')(o)

    return Model(img, o)
    
def resnet50_unet_light(n_classes, input_height=224, input_width=224, task="segmentation", weight_decay=1e-6, pretraining=False):
    assert input_height % 32 == 0
    assert input_width % 32 == 0

    img_input = Input(shape=(input_height, input_width, 3))

    features = resnet50(img_input, weight_decay=weight_decay, pretraining=pretraining)

    return unet_decoder(img_input, *features, n_classes, light=True, task=task, weight_decay=weight_decay)

def resnet50_unet(n_classes, input_height=224, input_width=224, task="segmentation", weight_decay=1e-6, pretraining=False):
    assert input_height % 32 == 0
    assert input_width % 32 == 0

    img_input = Input(shape=(input_height, input_width, 3))

    features = resnet50(img_input, weight_decay=weight_decay, pretraining=pretraining)

    return unet_decoder(img_input, *features, n_classes, light=False, task=task, weight_decay=weight_decay)

def transformer_block(img,
                      num_patches,
                      patchsize_x,
                      patchsize_y,
                      mlp_head_units,
                      n_layers,
                      num_heads,
                      projection_dim):
    patches = Patches(patchsize_x, patchsize_y)(img)
    # Encode patches.
    encoded_patches = PatchEncoder(num_patches, projection_dim)(patches)

    for _ in range(n_layers):
        # Layer normalization 1.
        x1 = LayerNormalization(epsilon=1e-6)(encoded_patches)
        # Create a multi-head attention layer.
        attention_output = MultiHeadAttention(num_heads=num_heads,
                                              key_dim=projection_dim,
                                              dropout=0.1)(x1, x1)
        # Skip connection 1.
        x2 = Add()([attention_output, encoded_patches])
        # Layer normalization 2.
        x3 = LayerNormalization(epsilon=1e-6)(x2)
        # MLP.
        x3 = mlp(x3, hidden_units=mlp_head_units, dropout_rate=0.1)
        # Skip connection 2.
        encoded_patches = Add()([x3, x2])

    encoded_patches = Reshape((img.shape[1],
                               img.shape[2],
                               projection_dim // (patchsize_x * patchsize_y)),
                              name="reshape_patches")(encoded_patches)
    return encoded_patches

def vit_resnet50_unet(num_patches,
                      n_classes,
                      transformer_patchsize_x,
                      transformer_patchsize_y,
                      transformer_mlp_head_units=None,
                      transformer_layers=8,
                      transformer_num_heads=4,
                      transformer_projection_dim=64,
                      input_height=224,
                      input_width=224,
                      task="segmentation",
                      weight_decay=1e-6,
                      pretraining=False):
    if transformer_mlp_head_units is None:
        transformer_mlp_head_units = [128, 64]
    inputs = Input(shape=(input_height, input_width, 3))
    
    features = list(resnet50(inputs, weight_decay=weight_decay, pretraining=pretraining))

    features[-1] = transformer_block(features[-1],
                                     num_patches,
                                     transformer_patchsize_x,
                                     transformer_patchsize_y,
                                     transformer_mlp_head_units,
                                     transformer_layers,
                                     transformer_num_heads,
                                     transformer_projection_dim)

    return unet_decoder(inputs, *features, n_classes, task=task, weight_decay=weight_decay)

def vit_resnet50_unet_transformer_before_cnn(num_patches,
                                             n_classes,
                                             transformer_patchsize_x,
                                             transformer_patchsize_y,
                                             transformer_mlp_head_units=None,
                                             transformer_layers=8,
                                             transformer_num_heads=4,
                                             transformer_projection_dim=64,
                                             input_height=224,
                                             input_width=224,
                                             task="segmentation",
                                             weight_decay=1e-6,
                                             pretraining=False):
    if transformer_mlp_head_units is None:
        transformer_mlp_head_units = [128, 64]
    inputs = Input(shape=(input_height, input_width, 3))

    encoded_patches = transformer_block(inputs,
                                        num_patches,
                                        transformer_patchsize_x,
                                        transformer_patchsize_y,
                                        transformer_mlp_head_units,
                                        transformer_layers,
                                        transformer_num_heads,
                                        transformer_projection_dim)
    encoded_patches = Conv2D(3, (1, 1), padding='same',
                             data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay),
                             name='convinput')(encoded_patches)

    features = resnet50(encoded_patches, weight_decay=weight_decay, pretraining=pretraining)

    return unet_decoder(inputs, *features, n_classes, task=task, weight_decay=weight_decay)

def resnet50_classifier(n_classes,input_height=224,input_width=224,weight_decay=1e-6,pretraining=False):
    include_top=True
    assert input_height%32 == 0
    assert input_width%32 == 0

    
    img_input = Input(shape=(input_height,input_width , 3 ))

    _, _, _, _, x = resnet50(img_input, weight_decay, pretraining)

    x = AveragePooling2D((7, 7), name='avg_pool')(x)
    x = Flatten()(x)
    
    ##
    x = Dense(256, activation='relu', name='fc512')(x)
    x=Dropout(0.2)(x)
    ##
    x = Dense(n_classes, activation='softmax', name='fc1000')(x)
    model = Model(img_input, x)
    
    return model

def machine_based_reading_order_model(n_classes,input_height=224,input_width=224,weight_decay=1e-6,pretraining=False):
    assert input_height%32 == 0
    assert input_width%32 == 0

    img_input = Input(shape=(input_height,input_width , 3 ))

    _, _, _, _, x = resnet50(img_input, weight_decay, pretraining)
    
    x = AveragePooling2D((7, 7), name='avg_pool1')(x)
    flattened = Flatten()(x)
    
    o = Dense(256, activation='relu', name='fc512')(flattened)
    o=Dropout(0.2)(o)
    
    o = Dense(256, activation='relu', name='fc512a')(o)
    o=Dropout(0.2)(o)

    o = Dense(n_classes, activation='sigmoid', name='fc1000')(o)
    model = Model(img_input , o)

    return model

def cnn_rnn_ocr_model(input_height=None, input_width=None, n_classes=None, max_len=None, inference=False, characters_txt_file=None):
    inputs = Input(shape=(input_height, input_width, 3), name="image")
    labels = Input(name="label", shape=(None,))

    x = Conv2D(64,kernel_size=(3,3),padding="same")(inputs)
    x = BatchNormalization(name="bn1")(x)
    x = Activation("relu", name="relu1")(x)
    x = Conv2D(64,kernel_size=(3,3),padding="same")(x)
    x = BatchNormalization(name="bn2")(x)
    x = Activation("relu", name="relu2")(x)
    x = MaxPooling2D(pool_size=(1,2),strides=(1,2))(x)

    x = Conv2D(128,kernel_size=(3,3),padding="same")(x)
    x = BatchNormalization(name="bn3")(x)
    x = Activation("relu", name="relu3")(x)
    x = Conv2D(128,kernel_size=(3,3),padding="same")(x)
    x = BatchNormalization(name="bn4")(x)
    x = Activation("relu", name="relu4")(x)
    x = MaxPooling2D(pool_size=(1,2),strides=(1,2))(x)

    x = Conv2D(256,kernel_size=(3,3),padding="same")(x)
    x = BatchNormalization(name="bn5")(x)
    x = Activation("relu", name="relu5")(x)
    x = Conv2D(256,kernel_size=(3,3),padding="same")(x)
    x = BatchNormalization(name="bn6")(x)
    x = Activation("relu", name="relu6")(x)
    x = MaxPooling2D(pool_size=(2,2),strides=(2,2))(x)

    x = Conv2D(input_width,kernel_size=(3,3),padding="same")(x)
    x = BatchNormalization(name="bn7")(x)
    x = Activation("relu", name="relu7")(x)
    x = Conv2D(input_width,kernel_size=(16,1))(x)
    x = BatchNormalization(name="bn8")(x)
    x = Activation("relu", name="relu8")(x)
    x2d = MaxPooling2D(pool_size=(1,2),strides=(1,2))(x)
    x4d = MaxPooling2D(pool_size=(1,2),strides=(1,2))(x2d)

    new_shape = (x.shape[1]*x.shape[2], x.shape[3])
    new_shape2 = (x2d.shape[1]*x2d.shape[2], x2d.shape[3])
    new_shape4 = (x4d.shape[1]*x4d.shape[2], x4d.shape[3])

    x = Reshape(new_shape, name="reshape")(x)
    x2d = Reshape(new_shape2, name="reshape2")(x2d)
    x4d = Reshape(new_shape4, name="reshape4")(x4d)

    xrnnorg = Bidirectional(LSTM(input_width, return_sequences=True, dropout=0.25))(x)
    xrnn2d = Bidirectional(LSTM(input_width, return_sequences=True, dropout=0.25))(x2d)
    xrnn4d = Bidirectional(LSTM(input_width, return_sequences=True, dropout=0.25))(x4d)

    xrnn2d = Reshape((1, xrnn2d.shape[1], xrnn2d.shape[2]), name="reshape6")(xrnn2d)
    xrnn4d = Reshape((1, xrnn4d.shape[1], xrnn4d.shape[2]), name="reshape8")(xrnn4d)

    xrnn2dup = UpSampling2D(size=(1, 2), interpolation="nearest")(xrnn2d)
    xrnn4dup = UpSampling2D(size=(1, 4), interpolation="nearest")(xrnn4d)

    xrnn2dup = Reshape((xrnn2dup.shape[2], xrnn2dup.shape[3]), name="reshape10")(xrnn2dup)
    xrnn4dup = Reshape((xrnn4dup.shape[2], xrnn4dup.shape[3]), name="reshape12")(xrnn4dup)

    addition = Add()([xrnnorg, xrnn2dup, xrnn4dup])

    addition_rnn = Bidirectional(LSTM(input_width, return_sequences=True, dropout=0.25))(addition)

    #out = Conv1D(max_len, 1, data_format="channels_first")(addition_rnn)
    out = Permute((2, 1))(addition_rnn)
    out = Conv1D(max_len, 1, data_format="channels_last")(out)
    out = Permute((2, 1))(out)
    out = BatchNormalization(name="bn9")(out)
    out = Activation("relu", name="relu9")(out)

    out = Dense(n_classes, activation="softmax", name="dense2")(out)

    if inference:
        # add second path for binarization
        inputs_bin = Input(shape=(input_height, input_width, 3), name="image_bin")
        out_bin = Model(inputs, out)(inputs_bin)
        # ensemble raw results
        out = 0.5 * (out + out_bin)
        # get tf.string batch
        out, prob = CTCDecoder()(out)
        # decode int to str
        with open(characters_txt_file, "r") as voc_file:
            voc = json.load(voc_file)
        char2num = StringLookup(vocabulary=voc)
        voc = char2num.get_vocabulary()
        num2char = StringLookup(vocabulary=voc, invert=True)
        output = num2char(out)

        return Model((inputs, inputs_bin), (output, prob))

    # Add CTC layer for calculating CTC loss at each step.
    out = CTCLayer(name="ctc_loss")(labels, out)
    
    return Model((inputs, labels), out)

def cnn_rnn_ocr_model4inference(model, model_path):
    """convert trained cnn-rnn-ocr model to inference model post-hoc"""
    try:
        model.get_layer(name='ctc_loss')
    except ValueError:
        # likely already converted
        return model
    else:
        inputs = model.get_layer(name='image').input
        output = model.get_layer(name='dense2').output
        inputs_bin = Input(inputs.shape[1:], name='image_bin')
        output_bin = Model(inputs, output)(inputs_bin)
        output = 0.5 * (output + output_bin)
        output, prob = CTCDecoder()(output)
        with open(model_path / "characters_org.txt", "r") as voc_file:
            voc = json.load(voc_file)
        char2num = StringLookup(vocabulary=voc)
        voc = char2num.get_vocabulary()
        num2char = StringLookup(vocabulary=voc, invert=True)
        output = num2char(output)
        inputs = (inputs, inputs_bin)
        outputs = (output, prob)
        return Model(inputs, outputs)

def get_model(config, logger):
    from sacred.config import create_captured_function

    task = config['task']
    if task in ["segmentation", "enhancement", "binarization"]:
        if config['backbone_type'] == 'nontransformer':
            builder = resnet50_unet
        else:
            num_patches_x, num_patches_y = config['transformer_num_patches_xy']
            num_patches = num_patches_x * num_patches_y

            if config['transformer_cnn_first']:
                builder = vit_resnet50_unet
                multiple = 32
            else:
                builder = vit_resnet50_unet_transformer_before_cnn
                multiple = 1

            assert config['input_height'] == (
                num_patches_y * config['transformer_patchsize_y'] * multiple), (
                "transformer_patchsize_y or transformer_num_patches_xy height value error: "
                "input_height should be equal to "
                "(transformer_num_patches_xy height value * transformer_patchsize_y * %d)" % multiple)
            assert config['input_width'] == (
                num_patches_x * config['transformer_patchsize_x'] * multiple), (
                    "transformer_patchsize_x or transformer_num_patches_xy width value error: "
                    "input_width should be equal to "
                    "(transformer_num_patches_xy width value * transformer_patchsize_x * %d)" % multiple)
            assert 0 == (config['transformer_projection_dim'] %
                         (config['transformer_patchsize_y'] *
                          config['transformer_patchsize_x'])), (
                             "transformer_projection_dim error: "
                             "The remainder when parameter transformer_projection_dim is divided by "
                             "(transformer_patchsize_y*transformer_patchsize_x) should be zero")

            config['num_patches'] = num_patches
    elif task == "cnn-rnn-ocr":
        builder = cnn_rnn_ocr_model
    elif task=='classification':
        builder = resnet50_classifier
    elif task=='reading_order':
        builder = machine_based_reading_order_model
    else:
        raise ValueError("unknown model task '%s'" % task)

    builder = create_captured_function(builder)
    builder.config = config
    builder.logger = logger
    return builder()
