import os
import numpy as np
import matplotlib.pyplot as plt
import sys

from pathlib import Path
from collections import Counter

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

class CTCLayer(layers.Layer):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.loss_fn = keras.backend.ctc_batch_cost

    def call(self, y_true, y_pred):
        # Compute the training-time loss value and add it
        # to the layer using `self.add_loss()`.
        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 = self.loss_fn(y_true, y_pred, input_length, label_length)
        self.add_loss(loss)

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


print("loading");
model = tf.keras.saving.load_model("captcha_75_25.tf", custom_objects={'CTCLayer': CTCLayer});

print("extracting");
pmodel = keras.models.Model(
    model.get_layer(name="image").input, model.get_layer(name="dense2").output
)
print("saving");

tf.keras.saving.save_model(pmodel, "pcaptcha_75_25.keras")
tf.keras.saving.save_model(pmodel, "pcaptcha_75_25.h5")

print("done");