import tensorflow as tf
import numpy as np

def select_model_tx(t_bdry, x_bdry, t_order, inital_t, net_layers, net_units, constraint, setup_boundaries):
  """
  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).
    inital_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.

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

  if boundary_type == "periodic_timeDependent":
    if constraint == "soft":
      model = build_model_periodic_tx(t_bdry, x_bdry, net_layers, net_units)
    elif constraint == "hard":
      if t_order == 1:
        model = build_model_periodic_hardconstraint1_tx(t_bdry, x_bdry, inital_t[0], net_layers, net_units)
      elif t_order == 2:
        model = build_model_periodic_hardconstraint2_tx(t_bdry, x_bdry, inital_t[0], inital_t[1], net_layers, net_units)
      elif t_order == 3:
        model = build_model_periodic_hardconstraint3_tx(t_bdry, x_bdry, inital_t[0], inital_t[1], inital_t[2], net_layers, net_units)
                      
  elif boundary_type == "dirichlet_timeDependent":
    if constraint == "soft":
      model = build_model_standard(t_bdry, x_bdry, net_layers, net_units)
    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 t_order == 2:
        model = build_model_dirichlet_hardconstraint2_tx(t_bdry, x_bdry, inital_t[0], inital_t[1], net_layers, net_units, setup_boundaries[3], setup_boundaries[4])

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

def select_model_xy(x_bdry, y_bdry, net_layers, net_units, constraint, setup_boundaries):
  """
  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.

  Returns:
    model (PINN): 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_model_standard(x_bdry, y_bdry, net_layers, net_units)
    elif constraint == "hard":
      model = build_model_hardconstraint_xy(x_bdry, y_bdry, net_layers, net_units, setup_boundaries[3], setup_boundaries[4],
                                             setup_boundaries[5], setup_boundaries[6])
  elif boundary_type == "neumann_timeIndependent":
    if constraint == "soft":
      model = build_model_standard(x_bdry, y_bdry, net_layers, net_units)
         
  elif boundary_type == "periodic_timeIndependent":
    model = build_model_periodic_xy(net_layers, net_units, x_bdry, y_bdry)
      
  return model

class Periodic(tf.keras.layers.Layer):
  """
  Class which describes a Periodic layer for PINN. Used in periodic models
  """
  def __init__(self, xmin, xmax):
    super(Periodic, self).__init__()
    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

# Define the network
def build_model_periodic_tx(t, x_bdry, n_layers, n_units):
  """
  Used in solving:
  <ul>
  <li>solvePDE_tx with periodic boundaries and soft constraints</li>
  </ul>

  Args:
    t (list): t_bdry, boundary to train t on 
    n_layers (int): Number of network internal layers
    n_units (int): Number of units per internal network layer

  Returns:
    model (PINN): Constructed model

  """

  # Define the network
  inp1 = tf.keras.layers.Input(shape=(1,))
  b1 = Normalize(t[0], t[1])(inp1)

  inp2 = tf.keras.layers.Input(shape=(1,))
  b2 = Periodic(x_bdry[0], x_bdry[1])(inp2)
  b = tf.keras.layers.Concatenate()([b1, b2])

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

  model = tf.keras.models.Model([inp1, inp2], out)
  model.summary()

  return model

def build_model_standard(x_bdry, y_bdry, n_layers, n_units):
  """
  Model used by both tx and xy equations, standard PINN with 2 variables, no periodic
  layers and no hard constrainting.

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

  Args:
    x_bdry (list): Used as t or x boundary
    y_bdry (list): Used as x or y boundary
    n_layers (int): Number of network internal layers
    n_units (int): Number of units per internal network layer

  Returns:
    model (PINN): Constructed model

  """

  # Define the network
  inp1 = tf.keras.layers.Input(shape=(1,))
  b1 = Normalize(x_bdry[0], x_bdry[1])(inp1)

  inp2 = tf.keras.layers.Input(shape=(1,))
  b2 = Normalize(y_bdry[0], y_bdry[1])(inp2)
  b = tf.keras.layers.Concatenate()([b1, b2])

  for i in range(n_layers):
    b = tf.keras.layers.Dense(n_units, activation='tanh', kernel_initializer='glorot_normal')(b)
  out = tf.keras.layers.Dense(1, activation='linear')(b)

  model = tf.keras.models.Model([inp1, inp2], out)
  model.summary()

  return model


def build_model_periodic_xy(n_layers, n_units, x, y):
  """
  Used in solving:
  <ul>
  <li>solvePDE_xy with periodic boundaries and soft constraints</li>
  <li>solvePDE_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

  Returns:
    model (PINN): 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.

  """

  # Define the network
  inp1 = tf.keras.layers.Input(shape=(1,))
  b1 = Periodic(x[0], x[1])(inp1)

  inp2 = tf.keras.layers.Input(shape=(1,))
  b2 = Periodic(y[0], y[1])(inp2)
  b = tf.keras.layers.Concatenate()([b1, b2])

  for i in range(n_layers):
    b = tf.keras.layers.Dense(n_units, activation='tanh', kernel_initializer='glorot_normal')(b)
  out = tf.keras.layers.Dense(1, activation='linear')(b)

  model = tf.keras.models.Model([inp1, inp2], out)
  model.summary()

  return model


# Define the network
def build_model_periodic_hardconstraint1_tx(t, x, u0, n_layers, n_units):
  """
  Used in solving:
  <ul>
  <li>solvePDE_tx with periodic boundaries, hard constraints of a 
    single order equation (i.e, a single initial condtion)</li>
  </ul>

  Args:
    t (list): t_bdry, boundary to train t on  
    u0 (lambda): Function for initial condition u(t0, x)
    n_layers (int): Number of network internal layers
    n_units (int): Number of units per internal network layer

  Returns:
    model (PINN): Constructed model

  """

  # Define the network
  inp1 = tf.keras.layers.Input(shape=(1,))
  b1 = Normalize(t[0], t[1])(inp1)

  inp2 = tf.keras.layers.Input(shape=(1,))
  b2 = Periodic(x[0], x[1])(inp2)
  b = tf.keras.layers.Concatenate()([b1, b2])

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

  #Scales u0
  #out = tf.keras.layers.Lambda(lambda x: u0(x[1])*(1-(x[0]-t[0])/(t[1]-t[0])) + (x[0]-t[0])/(t[1]-t[0])*x[2])([inp1, inp2, out])
  #still works and quicker
  out = tf.keras.layers.Lambda(lambda x: u0(x[1]) + (x[0]-t[0])/(t[1]-t[0])*x[2])([inp1, inp2, out])

  model = tf.keras.models.Model([inp1, inp2], out)
  model.summary()

  return model

def build_model_periodic_hardconstraint2_tx(t, x, u0, ut0, n_layers, n_units):
  """
  Used in solving:
  <ul>
  <li>solvePDE_tx with periodic boundaries, hard constraints of a 
    second order equation (i.e, two initial condtions)</li>
  </ul>

  Args:
    t (list): t_bdry, 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) 
    n_layers (int): Number of network internal layers
    n_units (int): Number of units per internal network layer

  Returns:
    model (PINN): Constructed model

  """

  # Define the network
  inp1 = tf.keras.layers.Input(shape=(1,))
  b1 = Normalize(t[0], t[1])(inp1)

  inp2 = tf.keras.layers.Input(shape=(1,))
  b2 = Periodic(x[0], x[1])(inp2)
  b = tf.keras.layers.Concatenate()([b1, b2])

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

  #better hard constraint to just not scale and let nn learn
  out = tf.keras.layers.Lambda(lambda x: u0(x[1]) + ut0(x[1])*(x[0] - t[0])
                               +  (((x[0]-t[0])/(t[1]-t[0]))**2)*x[2])([inp1, inp2, out])
  
  #Scales inital conditions down away from t0, slow/not really working, does work after a while though
  # out = tf.keras.layers.Lambda(lambda x: u0(x[1])*(1-(x[0]-t[0])/(t[1]-t[0])) + 
  #                              ut0(x[1])*((x[0]-t[0])/(t[1]-t[0]))*(1-(x[0]-t[0])/(t[1]-t[0]))
  #                              +  (((x[0]-t[0])/(t[1]-t[0]))**2)*x[2])([inp1, inp2, out])

  model = tf.keras.models.Model([inp1, inp2], out)
  model.summary()

  return model

def build_model_periodic_hardconstraint3_tx(t, x, u0, ut0, utt0, n_layers, n_units):
  """
  Used in solving:
  <ul>
  <li>solvePDE_tx with periodic boundaries, hard constraints of a 
    third order equation (i.e, three initial condtions)</li>
  </ul>

  Args:
    t (list): t_bdry, 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)
    n_layers (int): Number of network internal layers
    n_units (int): Number of units per internal network layer

  Returns:
    model (PINN): Constructed model

  Models for solving solvePDE_xy equations
  --------------------------------------
  """

  # Define the network
  inp1 = tf.keras.layers.Input(shape=(1,))
  b1 = Normalize(t[0], t[1])(inp1)

  inp2 = tf.keras.layers.Input(shape=(1,))
  b2 = Periodic(x[0], x[1])(inp2)
  b = tf.keras.layers.Concatenate()([b1, b2])

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

  #better hard constraint to just not scale and let nn learn
  out = tf.keras.layers.Lambda(lambda x: u0(x[1]) + ut0(x[1])*(x[0] - t[0]) 
                               + utt0(x[1])*((x[0] - t[0])**2)
                               +  (((x[0]-t[0])/(t[1]-t[0]))**3)*x[2])([inp1, inp2, out])
  
  #would work but significantly slower even though it scales inital conds down
  # out = tf.keras.layers.Lambda(lambda x: u0(x[1])*(1-(x[0]-t[0])/(t[1]-t[0])) + ut0(x[1])*(x[0] - t[0])*(1-(x[0]-t[0])/(t[1]-t[0])) 
  #                              + utt0(x[1])*((x[0] - t[0])**2)*(1-(x[0]-t[0])/(t[1]-t[0]))
  #                              +  (((x[0]-t[0])/(t[1]-t[0]))**3)*x[2])([inp1, inp2, out])

  model = tf.keras.models.Model([inp1, inp2], out)
  model.summary()

  return model

# Define the network
def build_model_dirichlet_hardconstraint1_tx(t, x_bdry, u0, n_layers, n_units, xleft_bound, xright_bound):

  # Define the network
  inp1 = tf.keras.layers.Input(shape=(1,))
  b1 = Normalize(t[0], t[1])(inp1)

  inp2 = tf.keras.layers.Input(shape=(1,))
  b2 = Normalize(x_bdry[0], x_bdry[1])(inp2)
  b = tf.keras.layers.Concatenate()([b1, b2])

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

  x0 = np.float64(x_bdry[0])
  x1 = np.float64(x_bdry[1])
  t0 = np.float64(t[0])
  t1 = np.float64(t[1])

  #A(t,x) = (1-t)(tleftbound(tleft, x)) + t(trightbound(tright, x)) + 
  # (1-x)[xlowerbound(t, xlower) - ((1-t)(xlowerbound(tleft,xlower)) + t(xlowerbound(tright, xlower)))] + 
  # x[xupperbound(t, xupper) - ((1-t)(xupperbound(tleft, xupper)) + t(xupperbound(tright, xupper)))]

  # Add the input u as hard constraint (u(t,x) = u0(x)*(1-(t-t0)/(tfinal-t0)) + (t-t0)/(tfinal-t0)*NN(t,x)) t,    x,   u
  # out = tf.keras.layers.Lambda(lambda x: (1-((x[1]-x0)/(x1-x0)))*((x[1]-x0)/(x1-x0))*u0(x[1]) + 
  #                              (1-((x[1]-x0)/(x1-x0)))*
  #                             (xleft_bound(x[0]) - ((1-((x[0]-t0)/(t1-t0)))*xleft_bound(t0) + 
  #                              ((x[0]-t0)/(t1-t0)*xleft_bound(t1)))) + 
  #                              ((x[1]-x0)/(x1-x0)) * 
  #                              (xright_bound(x[0]) - ((1-((x[0]-t0)/(t1-t0)))*xright_bound(t0) + 
  #                              ((x[0]-t0)/(t1-t0)*xright_bound(t1)))) 
  #                              + ((x[0]-t0)/(t1-t0))*#(1-((x[0]-t0)/(t1-t0)))*
  #                              ((x[1]-x0)/(x1-x0))*(1-((x[1]-x0)/(x1-x0)))*x[2], output_shape=(1,))([inp1, inp2, out])
  
#*(1-(x[0]-t0)/(t1-t0))*((x[1]-x0)/(x1-x0))*(1-((x[1]-x0)/(x1-x0)))
  out = tf.keras.layers.Lambda(lambda x: u0(x[1]) + 
                               xleft_bound(x[0])*(1-(x[0]-t0)/(t1-t0))*((x[0]-t0)/(t1-t0))*(1-((x[1]-x0)/(x1-x0)))+
                               xright_bound(x[0])*(1-(x[0]-t0)/(t1-t0))*((x[0]-t0)/(t1-t0))*((x[1]-x0)/(x1-x0))+
                    ((x[0]-t0)/(t1-t0))*((x[1]-x0)/(x1-x0))*(1-((x[1]-x0)/(x1-x0)))*x[2], output_shape=(1,))([inp1, inp2, out])


  model = tf.keras.models.Model([inp1, inp2], out)
  model.summary()

  return model

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

  Args:
    x_bdry (list): boundary to train x on 
    y_bdry (list): boundary to train y on 
    n_layers (int): Number of network internal layers
    n_units (int): Number of units per internal network layer
    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 (PINN): Constructed model

  """

  # Define the network
  inp1 = tf.keras.layers.Input(shape=(1,))
  b1 = Normalize(x_bdry[0], x_bdry[1])(inp1)

  inp2 = tf.keras.layers.Input(shape=(1,))
  b2 = Normalize(y_bdry[0], y_bdry[1])(inp2)
  b = tf.keras.layers.Concatenate()([b1, b2])

  for i in range(n_layers):
    b = tf.keras.layers.Dense(n_units, activation='tanh', kernel_initializer='glorot_normal')(b)
  out = tf.keras.layers.Dense(1, activation='linear')(b)

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

  # x = (x-xleft)/(xright-xleft), y = (y-ylower)/(yupper-ylower)
  #A(x,y) + x(1-x)y(1-y)*NN, with
  #A(x,y) = (1-x)(xleftbound(xleft, y)) + x(xrightbound(xright, y)) + 
  # (1-y)[ylowerbound(x, ylower) - ((1-x)(ylowerbound(xleft,ylower)) + x(ylowerbound(xright, ylower)))] + 
  # y[yupperbound(x, yupper) - ((1-x)(yupperbound(xleft, yupper)) + x(yupperbound(xright, yupper)))]

  # when x = xleft (0), y=ylower(0)
  #1(xleftbound(0, 0)) + 0 + 1[ylowerbound(0, 0) - (1(yloerbound(0,0)) + 0)] + 0 = xleftboundary(0,0)
  # when x = xright (1), y=ylower(0)
  # 0 + xrightbound(1,0) + 1[ylowerbound(1,0) - (0 + ylowerbound(1, 0))] +0 = xrightboundary(1,0)
  #when x = xleft (0), y=yupper(1)
  # 1(xleftbound(0,1) + 0 + 0 +1[yupper(0,1) - (1(yupper(0,1)) + 0)]) = xleftbound(0,1) 
  #when x = xright (1), y=yupper (1)
  # 0 + xrightbound(1,1) + 0 + 1[yupperbound(1,1) - (0 + 1(yupperbound(1,1)))] = xrightbound(1,1)

  # when x = xleft (0):


  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,))([inp1, inp2, out])
  
  

  

  model = tf.keras.models.Model([inp1, inp2], out)
  model.summary()

  return model


# Define the network
def build_model_dirichlet_hardconstraint2_tx(t, x_bdry, u0, ut0, n_layers, n_units, xleft_bound, xright_bound):

  # Define the network
  inp1 = tf.keras.layers.Input(shape=(1,))
  b1 = Normalize(t[0], t[1])(inp1)

  inp2 = tf.keras.layers.Input(shape=(1,))
  b2 = Normalize(x_bdry[0], x_bdry[1])(inp2)
  b = tf.keras.layers.Concatenate()([b1, b2])

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

  x0 = np.float64(x_bdry[0])
  x1 = np.float64(x_bdry[1])
  t0 = np.float64(t[0])
  t1 = np.float64(t[1])

  out = tf.keras.layers.Lambda(lambda x: u0(x[1]) + ut0(x[1])*(x[0] - t0) +
                               xleft_bound(x[0])*(1-(x[0]-t0)/(t1-t0))**2*((x[0]-t0)/(t1-t0))**2*(1-((x[1]-x0)/(x1-x0)))**2+
                               xright_bound(x[0])*(1-(x[0]-t0)/(t1-t0))**2*((x[0]-t0)/(t1-t0))**2*((x[1]-x0)/(x1-x0))**2+
                  (((x[0]-t0)/(t1-t0))**2)*((x[1]-x0)/(x1-x0))**2*(1-((x[1]-x0)/(x1-x0)))**2*x[2], output_shape=(1,))([inp1, inp2, out])


  model = tf.keras.models.Model([inp1, inp2], out)
  model.summary()

  return model