import tensorflow as tf
import numpy as np
import optimctrltf.tf as oc
import math
import os


class ConstrainedAlgebraicMinimization(tf.keras.Model):
    def __init__(self, args):
        if args is not None:
            constraint_method = args[0]
        else:
            constraint_method = None
        super(ConstrainedAlgebraicMinimization, self).__init__()

        # Define the state vector (for inner optimization)
        self.u = tf.Variable(tf.random.uniform((1, 2),
                                               dtype=tf.dtypes.float64))

        self.dim_state = self.u.get_shape()
        self.dim_outputs = self.dim_state[1]

        # Number of dimensions of tensors
        self.dims = tf.range(2)

        # Define target state for constraint equations
        self.X1_TARGET = tf.constant(1., dtype=tf.dtypes.float64)
        self.X2_TARGET = tf.constant(2., dtype=tf.dtypes.float64)

        # Call to model returns the energy to be minimized
        self.energy_formulation = True

        self.constraint_method = constraint_method
        if constraint_method == "penalty":
            self.penalty_factor = 100000.
        elif constraint_method == "auglag":
            self.penalty_factor_list = [100., 1000., 10000.]
            self.penalty_factor = tf.Variable(self.penalty_factor_list[0],
                                              dtype=tf.dtypes.float64)
            # For a single constraint, define a (1, 1) tensor
            self.lagmult = tf.Variable(
                tf.zeros((1, 2), dtype=tf.dtypes.float64))

    def call(self, inputs):
        # Returns the energy x1^2+x2^2 - reference
        refen = tf.constant(math.pow(self.X1_TARGET, 2) +
                            math.pow(self.X2_TARGET, 2),
                            dtype=tf.dtypes.float64)
        energy = refen - tf.math.square(tf.norm(self.u))
        return energy

    def constraint_equation(self):
        # In case of a single constraint, return a (1, 1) tensor
        # return tf.stack([self.u[0, 1] - X2_TARGET])
        # Example: handle multiple constraints using tf.stack
        # MUST RETURN A 1 x # constraints tensor
        return tf.stack(
            [[self.u[0, 0] - self.X1_TARGET, self.u[0, 1] - self.X2_TARGET]])


def test_constrained_penalty_minimization_tf(min_method=""):
    """ This test minimizes the function
        x1^2 + x2^2
        subject to the constraints x1 = X1_TARGET & x2 = X2_TARGET,
        where (x1, x2) are the state variables, using the penalty method.
    """

    problem = oc.OptimizationProblem()

    problem.create_model(ConstrainedAlgebraicMinimization, args=("auglag", ))

    epochs = 10001

    problem.min_method = min_method

    problem.minimize(epochs, inputs=None, use_constraints=True, grad_tol=1e-8)

    print("Model solution (x1, x2) = ", problem.model.u.numpy()[0])
    print("Exact solution (x1, x2) = ",
          [problem.model.X1_TARGET.numpy(),
           problem.model.X2_TARGET.numpy()])


if __name__ == "__main__":
    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
    # tf.config.experimental_run_functions_eagerly(True)
    test_constrained_penalty_minimization_tf(min_method="lbfgs")
