import numpy as np
import tensorflow as tf
from tensorflow import keras


class NeuralSolver:
    def __init__(self, width, height, quotientX=False, quotientY=False):
        self.Width = width
        self.Height = height
        self.QuotientX = quotientX
        self.QuotientY = quotientY
        self._build_forward_model()
        self._build_reverse_model()

    def _build_forward_model(self):
        inputs = keras.Input(shape=(self.Width, self.Height, 1), name="InitialState")
        hidden = keras.Conv2D(32, 3, padding="same", activation="relu")(inputs)
        hidden = keras.layers.Conv2D(32, 3, padding="same", activation="relu")(hidden)
        outputs = keras.Conv2D(1, 1, padding="same", activation="sigmoid")(hidden)
        self.ForwardModel = keras.Model(inputs, outputs, name="ForwardModel")
        self.ForwardModel.compile(
            optimizer=keras.optimizers.Adam(learning_rate=0.001),
            loss=keras.losses.BinaryCrossentropy(),
            metrics=["accuracy"],
        )

    def _build_reverse_model(self):
        inputs = keras.Input(shape=(self.Width, self.Height, 1), name="FinalState")
        hidden = keras.Conv2D(32, 3, padding="same", activation="relu")(inputs)
        hidden = keras.layers.Conv2D(32, 3, padding="same", activation="relu")(hidden)
        outputs = keras.Conv2D(1, 1, padding="same", activation="sigmoid")(hidden)
        self.ReverseModel = keras.Model(inputs, outputs, name="ReverseModel")


    def train_forward(self, dataset, batch_size=8, epochs=10):
        x, y = dataset
        self.ForwardModel.fit(
            x=x,
            y=y,
            batch_size=batch_size,
            epochs=epochs,
            verbose=1
        )

    def train_backward(self, dataset, batch_size=8, epochs=10):
        x, y = dataset
        reverse_inputs = self.ReverseModel.inputs
        reverse_outputs = self.ReverseModel.outputs
        forward_outputs = self.ForwardModel(reverse_outputs)
        composite_model = keras.Model(reverse_inputs, forward_outputs, name="CompositeModel")
        composite_model.compile(
            optimizer=keras.optimizers.Adam(learning_rate=0.001),
            loss=keras.losses.BinaryCrossentropy(),
            metrics=["accuracy"]
        )
        composite_model.fit(
            x = y,
            y = y,
            batch_size=batch_size,
            epochs=epochs,
            verbose=1
        )

    def predict_forward(self, b):
        if len(b.shape) == 2:
            b = b[None, ..., None]
        preds = self.ForwardModel.predict(b)
        return preds > 0.5

    def predict_reverse(self, b):
        if len(b.shape) == 2:
            b = b[None, ..., None]
        preds = self.ReverseModel.predict(b)
        return preds > 0.5