import tensorflow as tf
import numpy as np

def select_DeepONet_tx(t_bdry, x_bdry, t_order, initial_t, net_layers, net_units, constraint, setup_boundaries, N_iv):
    """
  Main function called by solve class to select which model is needed in training, and calls that build_model function.

  Args:
    t_bdry (list): List of two elements, the interval of t to be solved on.
    x_bdry (list): List of two elements, the interval of x to be solved on.
    t_order (int): Order of t in equation (highest derivative of t used).
    initial_t (lambda): Inital function for t=t0, as a python lambda funciton, with t0 being inital t in t_bdry.
    net_layers (int): Number of internal layers of PINN.
    net_units (int): Number of units in each internal layer.
    constraint (string): Determines hard or soft constraints.
    setup_boundaries (boundary): boundary conditions set up from return of pde_Boundaries_2var call.
    N_iv (int): Number of randomly sampled collocation points along initial t.

  Returns:
    model (DeepONet): Model generated by build_DeepONet function called in this function.
  """
    boundary_type = setup_boundaries[0]

    if boundary_type == "periodic_timeDependent":
        if constraint == "soft":
            model = build_DeepONet_periodic_tx(net_layers, net_units, N_iv, t_bdry, x_bdry)
        elif constraint == "hard":
            if t_order == 1:
                model = build_DeepONet_periodic_tx_hard1(net_layers, net_units, N_iv, t_bdry, x_bdry, initial_t[0])
            elif t_order == 2:
                model = build_DeepONet_periodic_tx_hard2(net_layers, net_units, N_iv, t_bdry, x_bdry, initial_t[0], initial_t[1])
            elif t_order == 3:
                model = build_DeepONet_periodic_tx_hard3(net_layers, net_units, N_iv, t_bdry, x_bdry, initial_t[0], initial_t[1], initial_t[2])
                      
    elif boundary_type == "dirichlet_timeDependent":
        if constraint == "soft":
            model = build_DeepONet_standard(net_layers, net_units, N_iv, t_bdry, x_bdry)
    #     elif constraint=="hard":
    #         if t_order == 1:
    #             model = build_model_dirichlet_hardconstraint1_tx(t_bdry, x_bdry, inital_t[0], net_layers, net_units, setup_boundaries[3], setup_boundaries[4])

    elif boundary_type == "neumann_timeDependent":
        if constraint == "soft":
            model = build_DeepONet_standard(net_layers, net_units, N_iv, t_bdry, x_bdry)
      
    return model

def select_DeepONet_xy(x_bdry, y_bdry, net_layers, net_units, constraint, setup_boundaries, N_bc):
  """
  Main function called by solve class to select which model is needed in training, and calls that build_model function.

  Args:
    x_bdry (list): List of two elements, the interval of x to be solved on.
    y_bdry (list): List of two elements, the interval of y to be solved on.
    net_layers (int): Number of internal layers of PINN.
    net_units (int): Number of units in each internal layer.
    constraint (string): Determines hard or soft constraints.
    setup_boundaries (boundary): boundary conditions set up from return of pde_Boundaries_2var call.
    N_bc (int): Number of randomly sampled collocation points along boundaries.

  Returns:
    model (DeepONet): Model generated by build_model function called in this function.
  """
  boundary_type = setup_boundaries[0]

  if boundary_type == "dirichlet_timeIndependent":
    if constraint == "soft":
      model = build_DeepONet_standard(net_layers, net_units, N_bc, x_bdry, y_bdry)
    elif constraint == "hard":
        model = build_DeepONet_dirichlet_hardconstraint_xy(net_layers, net_units, N_bc, x_bdry, y_bdry,setup_boundaries[3], setup_boundaries[4],
                                             setup_boundaries[5], setup_boundaries[6])
  elif boundary_type == "neumann_timeIndependent":
    if constraint == "soft":
      model = build_DeepONet_standard(net_layers, net_units, N_bc, x_bdry, y_bdry)
         
  elif boundary_type == "periodic_timeIndependent":
    model = build_DeepONet_periodic_xy(net_layers, net_units, N_bc, x_bdry, y_bdry)
      
  return model

class PeriodicBCs(tf.keras.layers.Layer):
    """
    Class which describes a Periodic layer for PINN. Used in periodic models
    """
    def __init__(self, xmin, xmax, name=None, **kwargs):
        super(PeriodicBCs, self).__init__(name=name, **kwargs)
        self.xmin = xmin
        self.xmax = xmax

    def call(self, inputs):
        return tf.concat((tf.cos(2*np.pi*(inputs/(self.xmax - self.xmin))), tf.sin(2*np.pi*(inputs/(self.xmax - self.xmin)))), axis=1)


# Define the normalization layer
class Normalize(tf.keras.layers.Layer):
    """
    Class which describes a normalize layer for PINN. Returns input data
    normalized to interval [-1, 1].

    Models for solving solvePDE_tx equations
    --------------------------------------
    """
    def __init__(self, xmin, xmax, name=None, **kwargs):
        super(Normalize, self).__init__(name=name)
        self.xmin = xmin
        self.xmax = xmax
        super(Normalize, self).__init__(**kwargs)

    def call(self, inputs):
        return 2.0*(inputs-self.xmin)/(self.xmax-self.xmin)-1.0

    def get_config(self):
        config = super(Normalize, self).get_config()
        config.update({'xmin': self.xmin, 'xmax': self.xmax})
        return config
    
def mlp_network(inp, n_layers, n_units):
    """
    Function which creates simple mlp to be used as basis of trunk and branch nets in DeepONets

    Args:
        inp (tensor): Input layer data to mlp
        n_layers (int): Number of network internal layers
        n_units (int): Number of units per internal network layer

    Returns:
        out (tensor): Output data from output layer of network

    Models for solving solvePDE_DeepONet_tx equations
    --------------------------------------
    """

    h = inp
    for i in range(n_layers-1):
        h = tf.keras.layers.Dense(n_units, activation='tanh')(h)
        out = tf.keras.layers.Dense(n_units, activation='linear')(h)
    return out

def build_DeepONet_periodic_tx_hard1(n_layers, n_units, N_iv, t_bdry, x_bdry, u0):
    """
    Used in solving:
    <ul>
    <li>solvePDE_DeepONet_tx with periodic boundaries, hard constraints of a 
        single order equation (i.e, a single initial condtion)</li>
    </ul>

    Args:
        n_layers (int): Number of network internal layers
        n_units (int): Number of units per internal network layer
        N_iv (int): Number of randomly sampled collocation points along initial t
        t_bdry (list): Boundary to train t on  
        u0 (lambda): Function for initial condition u(t0, x)

    Returns:
        model (DeepONet): Constructed model

    """

    # Branch net
    inp_u = tf.keras.layers.Input(shape=(N_iv,))
    branch_net = mlp_network(inp_u, n_layers, n_units)

    # Trunk net
    inp_t = tf.keras.layers.Input(shape=(1,))
    b_t = Normalize(t_bdry[0], t_bdry[1])(inp_t)

    inp_x = tf.keras.layers.Input(shape=(1,))
    periodic_x = PeriodicBCs(x_bdry[0], x_bdry[1])(inp_x)

    inp_tx = tf.keras.layers.Concatenate()([b_t, periodic_x])

    # Actual trunk net
    trunk_net = mlp_network(inp_tx, n_layers, n_units)

    out = tf.keras.layers.Multiply()([branch_net, trunk_net])
    out = tf.keras.layers.Dense(1)(out)

    out = tf.keras.layers.Lambda(lambda x: u0(x[1])*(1-((x[0]-t_bdry[0])/(t_bdry[1]-t_bdry[0]))) + 
                                 (x[0]-t_bdry[0])/(t_bdry[1]-t_bdry[0])*x[2], output_shape=(1,))([inp_t, inp_x, out])

    model = tf.keras.models.Model([inp_t, inp_x, inp_u], out)
    model.summary()
    return model

def build_DeepONet_periodic_tx_hard2(n_layers, n_units, N_iv, t_bdry, x_bdry, u0, ut0):
    """
    Used in solving:
    <ul>
    <li>solvePDE_DeepONet_tx with periodic boundaries, hard constraints of a 
        second order equation (i.e, two initial condtions)</li>
    </ul>

    Args: 
        n_layers (int): Number of network internal layers
        n_units (int): Number of units per internal network layer
        N_iv (int): Number of randomly sampled collocation points along initial t
        t_bdry (list): Boundary to train t on 
        u0 (lambda): Function for initial condition, u(t0, x)
        ut0 (lambda): Function for initial condition of derivative, ut(t0, x)

    Returns:
        model (DeepONet): Constructed model

    """

    # Branch net
    inp_u = tf.keras.layers.Input(shape=(N_iv,))
    branch_net = mlp_network(inp_u, n_layers, n_units)

    # Trunk net
    inp_t = tf.keras.layers.Input(shape=(1,))
    b_t = Normalize(t_bdry[0], t_bdry[1])(inp_t)

    inp_x = tf.keras.layers.Input(shape=(1,))
    periodic_x = PeriodicBCs(x_bdry[0], x_bdry[1])(inp_x)

    inp_tx = tf.keras.layers.Concatenate()([b_t, periodic_x])

    # Actual trunk net
    trunk_net = mlp_network(inp_tx, n_layers, n_units)

    out = tf.keras.layers.Multiply()([branch_net, trunk_net])
    out = tf.keras.layers.Dense(1)(out)

    out = tf.keras.layers.Lambda(lambda x: u0(x[1]) + ut0(x[1])*(x[0]-t_bdry[0]) + 
                                 (((x[0]-t_bdry[0])/(t_bdry[1]-t_bdry[0]))**2)*x[2], output_shape=(1,))([inp_t, inp_x, out])

    model = tf.keras.models.Model([inp_t, inp_x, inp_u], out)
    model.summary()
    return model

def build_DeepONet_periodic_tx_hard3(n_layers, n_units, N_iv, t_bdry, x_bdry, u0, ut0, utt0):
    """
    Used in solving:
    <ul>
    <li>solvePDE_DeepONet_tx with periodic boundaries, hard constraints of a 
        third order equation (i.e, three initial condtions)</li>
    </ul>

    Args:
        n_layers (int): Number of network internal layers
        n_units (int): Number of units per internal network layer
        N_iv (int): Number of randomly sampled collocation points along initial t
        t_bdry (list): Boundary to train t on 
        u0 (lambda): Function for initial condition, u(t0, x)
        ut0 (lambda): Function for initial condition of derivative, ut(t0, x)
        utt0 (lambda): Function for initial condition of second derivative, utt(t0, x)

    Returns:
        model (DeepONet): Constructed model

    Models for solving solvePDE_DeepONet_xy equations
    --------------------------------------
    """

    # Branch net
    inp_u = tf.keras.layers.Input(shape=(N_iv,))
    branch_net = mlp_network(inp_u, n_layers, n_units)

    # Trunk net
    inp_t = tf.keras.layers.Input(shape=(1,))
    b_t = Normalize(t_bdry[0], t_bdry[1])(inp_t)

    inp_x = tf.keras.layers.Input(shape=(1,))
    periodic_x = PeriodicBCs(x_bdry[0], x_bdry[1])(inp_x)

    inp_tx = tf.keras.layers.Concatenate()([b_t, periodic_x])

    # Actual trunk net
    trunk_net = mlp_network(inp_tx, n_layers, n_units)

    out = tf.keras.layers.Multiply()([branch_net, trunk_net])
    out = tf.keras.layers.Dense(1)(out)

    out = tf.keras.layers.Lambda(lambda x: u0(x[1]) + ut0(x[1])*(x[0]-t_bdry[0]) + 
                                 utt0(x[1])*((x[0]-t_bdry[0])**2) + 
                                 (((x[0]-t_bdry[0])/(t_bdry[1]-t_bdry[0]))**3)*x[2], output_shape=(1,))([inp_t, inp_x, out])

    model = tf.keras.models.Model([inp_t, inp_x, inp_u], out)
    model.summary()
    return model

def build_DeepONet_periodic_tx(n_layers, n_units, N_iv, t_bdry, x_bdry):
    """
    Used in solving:
    <ul>
    <li>solvePDE_DeepONet_tx with periodic boundaries and soft constraints</li>
    </ul>

    Args:
        n_layers (int): Number of network internal layers
        n_units (int): Number of units per internal network layer
        N_iv (int): Number of randomly sampled collocation points along initial t
        t_bdry (list): Boundary to train t on 

    Returns:
        model (DeepONet): Constructed model

    """

    # Branch net
    inp_u = tf.keras.layers.Input(shape=(N_iv,))
    branch_net = mlp_network(inp_u, n_layers, n_units)

    # Trunk net
    inp_t = tf.keras.layers.Input(shape=(1,))
    b_t = Normalize(t_bdry[0], t_bdry[1])(inp_t)

    inp_x = tf.keras.layers.Input(shape=(1,))
    periodic_x = PeriodicBCs(x_bdry[0], x_bdry[1])(inp_x)

    inp_tx = tf.keras.layers.Concatenate()([b_t, periodic_x])

    # Actual trunk net
    trunk_net = mlp_network(inp_tx, n_layers, n_units)

    out = tf.keras.layers.Multiply()([branch_net, trunk_net])
    out = tf.keras.layers.Dense(1)(out)

    model = tf.keras.models.Model([inp_t, inp_x, inp_u], out)
    model.summary()
    return model

def build_DeepONet_standard(n_layers, n_units, N, t_bdry, x_bdry):
    """
    Model used by both tx and xy equations, standard DeepONet with 2 variables, no periodic
    layers and no hard constrainting.

    Used in solving:
    <ul>
    <li>solvePDE_DeepONet_tx with dirichlet boundaries and soft constraints</li>
    <li>solvePDE_DeepONet_tx with neumann boundaries and soft constraints</li>
    <li>solvePDE_DeepONet_xy with dirichlet boundaries and soft constraints</li>
    <li>solvePDE_DeepONet_xy with neumann boundaries and soft constraints</li>
    </ul>

    Args:
        n_layers (int): Number of network internal layers
        n_units (int): Number of units per internal network layer
        N (int): N_iv or N_bc, Number of randomly sampled collocation points along initial t 
            or number of randomly sampled collocation points along boundaries.
        t_bdry (list): Used as t or x boundary
        x_bdry (list): Used as x or y boundary

    Returns:
        model (DeepONet): Constructed model

    """

    # Branch net
    inp_u = tf.keras.layers.Input(shape=(N,))
    branch_net = mlp_network(inp_u, n_layers, n_units)

    # Trunk net
    inp_t = tf.keras.layers.Input(shape=(1,))
    b_t = Normalize(t_bdry[0], t_bdry[1])(inp_t)

    inp_x = tf.keras.layers.Input(shape=(1,))
    b_x = Normalize(x_bdry[0], x_bdry[1])(inp_x)

    inp_tx = tf.keras.layers.Concatenate()([b_t, b_x])

    # Actual trunk net
    trunk_net = mlp_network(inp_tx, n_layers, n_units)

    out = tf.keras.layers.Multiply()([branch_net, trunk_net])
    out = tf.keras.layers.Dense(1)(out)

    model = tf.keras.models.Model([inp_t, inp_x, inp_u], out)
    model.summary()
    return model

def build_DeepONet_periodic_xy(n_layers, n_units, N_bc, x_bdry, y_bdry):
    """
    Used in solving:
    <ul>
    <li>solvePDE_DeepONet_xy with periodic boundaries and soft constraints</li>
    <li>solvePDE_DeepONet_xy with periodic boundaries and hard constraints</li>
    </ul>

    Args:
        n_layers (int): Number of network internal layers
        n_units (int): Number of units per internal network layer
        N_bc (int): Number of randomly sampled collocation points along boundaries.

    Returns:
        model (DeepONet): Constructed model

    If using periodic boundaries in a time independent equation, all boundaries are periodic
    and therefore there is nothing to be soft/hard constrainted, as periodic implemented on a 
    network layer level.

    """

    # Branch net
    inp_u = tf.keras.layers.Input(shape=(N_bc,))
    branch_net = mlp_network(inp_u, n_layers, n_units)

    # Trunk net
    inp_x = tf.keras.layers.Input(shape=(1,))
    periodic_x = PeriodicBCs(x_bdry[0], x_bdry[1])(inp_x)

    inp_y = tf.keras.layers.Input(shape=(1,))
    periodic_y = PeriodicBCs(y_bdry[0], y_bdry[1])(inp_y)

    inp_tx = tf.keras.layers.Concatenate()([periodic_x, periodic_y])

    # Actual trunk net
    trunk_net = mlp_network(inp_tx, n_layers, n_units)

    out = tf.keras.layers.Multiply()([branch_net, trunk_net])
    out = tf.keras.layers.Dense(1)(out)

    model = tf.keras.models.Model([inp_x, inp_y, inp_u], out)
    model.summary()
    return model

def build_DeepONet_dirichlet_hardconstraint_xy(n_layers, n_units, N_bc, x_bdry, y_bdry, xleft_bound, 
                                               xright_bound, ylower_bound, yupper_bound):
    """
    Used in solving:
    <ul>
    <li>solvePDE_DeepONet_xy with dirichlet boundaries and hard constraints</li>
    </ul>

    Args:
        n_layers (int): Number of network internal layers
        n_units (int): Number of units per internal network layer
        N_bc (int): Number of randomly sampled collocation points along boundaries
        x_bdry (list): boundary to train x on 
        y_bdry (list): boundary to train y on 
        xleft_bound (lambda): Function which describes x left boundary condition, lambda of 2 variables
        xright_bound (lambda): Function which describes x right boundary condition, lambda of 2 variables
        ylower_bound (lambda): Function which describes y lower boundary condition, lambda of 2 variables
        yupper_bound (lambda): Function which describes y upper boundary condition, lambda of 2 variables

    Returns:
        model (DeepONet): Constructed model

    """

    # Branch net
    inp_u = tf.keras.layers.Input(shape=(N_bc,))
    branch_net = mlp_network(inp_u, n_layers, n_units)

    # Trunk net
    inp_x = tf.keras.layers.Input(shape=(1,))
    b_x = Normalize(x_bdry[0], x_bdry[1])(inp_x)

    inp_y = tf.keras.layers.Input(shape=(1,))
    b_y = Normalize(y_bdry[0], y_bdry[1])(inp_y)

    inp_xy = tf.keras.layers.Concatenate()([b_x, b_y])

    # Actual trunk net
    trunk_net = mlp_network(inp_xy, n_layers, n_units)

    out = tf.keras.layers.Multiply()([branch_net, trunk_net])
    out = tf.keras.layers.Dense(1)(out)

    x0 = np.float64(x_bdry[0])
    x1 = np.float64(x_bdry[1])
    y0 = np.float64(y_bdry[0])
    y1 = np.float64(y_bdry[1])

    out = tf.keras.layers.Lambda(lambda x: (1-((x[0]-x0)/(x1-x0)))*xleft_bound(x0, x[1]) + 
                              ((x[0]-x0)/(x1-x0))*xright_bound(x1, x[1]) + (1-((x[1]-y0)/(y1-y0)))*
                              (ylower_bound(x[0], y0) - ((1-((x[0]-x0)/(x1-x0)))*ylower_bound(x0, y0) + 
                               ((x[0]-x0)/(x1-x0)*ylower_bound(x1, y0)))) + 
                               ((x[1]-y0)/(y1-y0)) * 
                               (yupper_bound(x[0], y1) - ((1-((x[0]-x0)/(x1-x0)))*yupper_bound(x0, y1) + 
                               ((x[0]-x0)/(x1-x0)*yupper_bound(x1, y1)))) 
                               + ((x[0]-x0)/(x1-x0))*(1-((x[0]-x0)/(x1-x0)))*
                               ((x[1]-y0)/(y1-y0))*(1-((x[1]-y0)/(y1-y0)))*x[2], output_shape=(1,))([inp_x, inp_y, out])

    model = tf.keras.models.Model([inp_x, inp_y, inp_u], out)
    model.summary()
    return model