package zenith.zov.base.rotation.deeplearnig;

import ai.djl.Model;
import ai.djl.ModelException;
import ai.djl.inference.Predictor;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.Shape;
import ai.djl.nn.Activation;
import ai.djl.nn.Blocks;
import ai.djl.nn.SequentialBlock;
import ai.djl.nn.core.Linear;
import ai.djl.nn.norm.BatchNorm;
import ai.djl.training.DefaultTrainingConfig;
import ai.djl.training.EasyTrain;
import ai.djl.training.Trainer;
import ai.djl.training.TrainingConfig;
import ai.djl.training.dataset.ArrayDataset;
import ai.djl.training.dataset.ArrayDataset.Builder;
import ai.djl.training.initializer.XavierInitializer;
import ai.djl.training.listener.TrainingListener;
import ai.djl.training.listener.TrainingListenerAdapter;
import ai.djl.training.listener.TrainingListener.Defaults;
import ai.djl.training.loss.Loss;
import ai.djl.training.optimizer.Adam;
import ai.djl.training.tracker.Tracker;
import ai.djl.translate.TranslateException;
import ai.djl.translate.Translator;
import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Path;

public abstract class ModelWrapper<I, O> implements Closeable {
    private static final int NUM_EPOCH = 500;
    private static final int BATCH_SIZE = 32;
    private final Translator<I, O> translator;
    private final long outputs;
    private final Model model;
    private final Predictor<I, O> predictor;
    private final String name;

    public ModelWrapper(String name, Translator<I, O> translator, long outputs) {
        this.name = name;
        this.translator = translator;
        this.outputs = outputs;
        this.model = Model.newInstance(name);
        this.model.setBlock(createMlpBlock(outputs));
        this.predictor = this.model.newPredictor(translator);
    }

    public O predict(I input) throws TranslateException {
        return (O)this.predictor.predict(input);
    }

    public void train(float[][] features, float[][] labels) throws ModelException, IOException {
        if (features.length == labels.length && features.length != 0) {
            long inputs = (long)features[0].length;
            TrainingConfig trainingConfig = new DefaultTrainingConfig(Loss.l2Loss())
                .optInitializer(new XavierInitializer(), "weight")
                .optOptimizer(Adam.builder().optLearningRateTracker(Tracker.fixed(0.001F)).build())
                .addTrainingListeners(Defaults.logging("train"))
                .addTrainingListeners(new TrainingListener[]{new TrainingListenerAdapter() {
                    public void onEpoch(Trainer trainer) {
                        System.out.println("Epoch " + trainer.getTrainingResult().getEpoch() + " - Loss: " + trainer.getTrainingResult().getTrainLoss());
                    }

                    public void onTrainingBegin(Trainer trainer) {
                    }

                    public void onTrainingEnd(Trainer trainer) {
                    }
                }});

            try {
                Trainer trainer = this.model.newTrainer(trainingConfig);

                try {
                    NDManager manager = NDManager.newBaseManager();

                    try {
                        ArrayDataset trainingSet = ((Builder)new Builder()
                                .setData(new NDArray[]{manager.create(features)})
                                .optLabels(new NDArray[]{manager.create(labels)})
                                .setSampling(32, true))
                            .build();
                        trainer.initialize(new Shape[]{new Shape(new long[]{32L, inputs})});
                        EasyTrain.fit(trainer, 500, trainingSet, null);
                    } catch (Throwable var12) {
                        if (manager != null) {
                            try {
                                manager.close();
                            } catch (Throwable var11) {
                                var12.addSuppressed(var11);
                            }
                        }

                        throw var12;
                    }

                    if (manager != null) {
                        manager.close();
                    }
                } catch (Throwable var13) {
                    if (trainer != null) {
                        try {
                            trainer.close();
                        } catch (Throwable var10) {
                            var13.addSuppressed(var10);
                        }
                    }

                    throw var13;
                }

                if (trainer != null) {
                    trainer.close();
                }
            } catch (TranslateException var14) {
            }
        } else {
            throw new IllegalArgumentException("Features and labels must have the same size and be non-empty");
        }
    }

    public void load(InputStream stream) throws IOException, ModelException {
        this.model.load(stream);
    }

    public void load(Path path) throws IOException, ModelException {
        this.model.load(path, "tf");
    }

    public void load() {
        try {
            try (InputStream stream = this.getClass().getResourceAsStream("/assets/zenith/models/" + this.name + ".params")) {
                this.load(stream);
            }
        } catch (Exception var6) {
            throw new RuntimeException(var6);
        }
    }

    public void save(Path path) throws IOException, ModelException {
        this.model.save(path, "tf");
    }

    @Override
    public void close() {
        this.predictor.close();
        this.model.close();
    }

    private static SequentialBlock createMlpBlock(long outputs) {
        return new SequentialBlock()
            .add(Linear.builder().setUnits(128L).build())
            .add(Blocks.batchFlattenBlock())
            .add(BatchNorm.builder().build())
            .add(Activation.reluBlock())
            .add(Linear.builder().setUnits(64L).build())
            .add(Blocks.batchFlattenBlock())
            .add(BatchNorm.builder().build())
            .add(Activation.reluBlock())
            .add(Linear.builder().setUnits(32L).build())
            .add(Blocks.batchFlattenBlock())
            .add(BatchNorm.builder().build())
            .add(Activation.reluBlock())
            .add(Linear.builder().setUnits(outputs).build());
    }
}
