from sklearn.model_selection import train_test_split
from tensorflow.keras import Model, Input
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.losses import CategoricalCrossentropy
from tensorflow.keras.optimizers import Adam

import numpy as np
import onnxruntime as ort
import onnx
import tensorflow as tf
import tf2onnx


class BaseModel:
    def __init__(self, dataset, seed):
        self.dataset = dataset
        self.seed = seed

    def load(self, path):
        self.model = onnx.load(path)
        self.session = ort.InferenceSession(path)

    def build(self, hidden_layers, activation, learning_rate):
        inputs = Input(shape=self.dataset.sample_shape)
        x = inputs

        # Only flatten if input is multi-dimensional (like digits or mnist)
        if len(self.dataset.sample_shape) > 1:
            x = Flatten()(x)

        for units in hidden_layers:
            x = Dense(units, activation=activation)(x)

        # Use softmax activation for classification
        outputs = Dense(self.dataset.num_classes, activation='softmax')(x)

        self.model = Model(inputs=inputs, outputs=outputs, name=self.name)

        self.model.compile(
            loss=CategoricalCrossentropy(),
            optimizer=Adam(learning_rate=learning_rate),
            metrics=['accuracy']
        )
        self.model.summary()

    def train(self, x, y, batch_size, epochs, patience):
        # Use 10% of the data for validation
        VAL_SIZE = 0.1

        x_train, x_val, y_train, y_val = train_test_split(x, y, test_size=VAL_SIZE, random_state=self.seed)

        early_stop = EarlyStopping(
            monitor='val_loss',
            patience=patience,
            restore_best_weights=True
        )

        self.model.fit(
            x_train, y_train,
            batch_size=batch_size,
            epochs=epochs,
            validation_data=(x_val, y_val),
            callbacks=[early_stop],
            verbose=1
        )

    def evaluate(self, x, y, batch_size):
        score = self.model.evaluate(
            x=x,
            y=y,
            batch_size=batch_size,
            verbose=1
        )

        return score

    def predict_proba(self, input_data):
        input_name = self.session.get_inputs()[0].name
        output_name = self.session.get_outputs()[0].name

        input_data = np.array(input_data, dtype=np.float32)
        probs = self.session.run([output_name], {input_name: input_data})[0]

        return probs

    def classify(self, input_data):
        probs = self.predict_proba(input_data)
        return np.argmax(probs, axis=1)

    def filter(self, samples, label):
        """Returns a list of samples that match the label."""
        predictions = self.classify(samples)
        return [sample for sample, prediction in zip(samples, predictions) if prediction == label]

    def get_samples(self, label, num_samples):
        samples = self.dataset.get_samples(label, 100000)
        samples = self.filter(samples, label)

        return samples[:num_samples]

    def save(self, path):
        if isinstance(self.model, Model):
            input_shape = self.dataset.sample_shape
            if len(input_shape) == 1:
                input_spec = (None,) + input_shape
            else:
                input_spec = (None,) + tuple(input_shape)

            # Create a function from the model
            @tf.function
            def model_function(x):
                return self.model(x)

            # Convert Keras to ONNX using the function approach
            self.model, _ = tf2onnx.convert.from_function(
                model_function,
                input_signature=[tf.TensorSpec(shape=input_spec, dtype=tf.float32)]
            )

            self.session = ort.InferenceSession(
                self.model.SerializeToString(),
                providers=['CPUExecutionProvider']
            )

        onnx.save(self.model, path)

    @property
    def name(self):
        return self.dataset.name

    @property
    def path(self):
        return self.dataset.path

    @property
    def last_node(self):
        return self.model.graph.node[-1]
