# Copyright (C) 2024 Tobias Bode
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.

import time
import sys

import jax
from jax import lax, config
import jax.numpy as jnp
import flax

from autopdex import seeder, geometry, solver, solution_structures, plotter, utility, models, spaces, assembler

config.update("jax_enable_x64", True)
config.update("jax_compilation_cache_dir", './cache')


### Definition of geometry and boundary conditions
# Cook's membrane, see e.g. https://link.springer.com/article/10.1007/s00466-017-1405-4

### Definition of geometry and boundary conditions
# Positive smooth distance functions that select the specific boundary segments
pts = [[0.,0.], [48.,44.], [48.,60.], [0.,44.]]
psdf_bottom = lambda x: geometry.psdf_trimmed_line(x, pts[0], pts[1])
psdf_right = lambda x: geometry.psdf_trimmed_line(x, pts[1], pts[2])
psdf_top = lambda x: geometry.psdf_trimmed_line(x, pts[2], pts[3])
psdf_left = lambda x: geometry.psdf_trimmed_line(x, pts[3], pts[0])
psdfs = [psdf_bottom, psdf_right, psdf_top, psdf_left]

# Positive smooth distance funcion of full domain
sdf_cooks = lambda x: geometry.sdf_convex_polygon_2d(x, pts)
psdf_cooks = lambda x: geometry.only_positive(sdf_cooks(x))
psdf = lambda x: jnp.asarray([psdf_cooks(x), psdf_cooks(x)])

# Boundary condition functions
homogeneous = lambda x: 0.
normals = lambda x: geometry.normals_from_normalized_sdf(x, psdf_cooks)

# Field selectors for boundary conditions
selection_dirichlet_1 = lambda x: jnp.asarray([1.,0.,0.,0.,0.,0.])
selection_dirichlet_2 = lambda x: jnp.asarray([0.,1.,0.,0.,0.,0.])
def selection_neumann_1(x):
  normal_x = normals(x)[0]
  normal_y = normals(x)[1]
  return jnp.asarray([0.,0.,normal_x,normal_y,0.,0.])
def selection_neumann_2(x):
  normal_x = normals(x)[0]
  normal_y = normals(x)[1]
  return jnp.asarray([0.,0.,0.,0.,normal_x,normal_y])

# Assignment of boundary conditions (same ordering as in psdfs)
load = lambda x: 4.0
bc_1 = [homogeneous, homogeneous, homogeneous, homogeneous]
bc_2 = [homogeneous, load, homogeneous, homogeneous]
a_1 = [selection_neumann_1, selection_neumann_1, selection_neumann_1, selection_dirichlet_1]
a_2 = [selection_neumann_2, selection_neumann_2, selection_neumann_2, selection_dirichlet_2]

### Transfinite interpolation of boundary segments
b_coeff = lambda x: jnp.asarray([geometry.transfinite_interpolation(x, psdfs, a_1),
                                   geometry.transfinite_interpolation(x, psdfs, a_2)])
b_cond = lambda x: jnp.asarray([geometry.transfinite_interpolation(x, psdfs, bc_1),
                                geometry.transfinite_interpolation(x, psdfs, bc_2)])


### Initialization and generation of nodes and integration points
# Regular node placement
# nod_spacing = 0.675
nod_spacing = 2.
(x_nodes, n_nodes, region_size) = seeder.regular_in_psdf(sdf_cooks, pts[0], pts[2], nod_spacing, atol=-1e-8)

# Gauss integration on background grid
int_spacing = nod_spacing
(x_int, w_int, n_int) = seeder.gauss_points_in_psdf(sdf_cooks, pts[0], pts[2], int_spacing, order=2)

# Support radius for moving least square ansatz
support_radius = 2.71 * nod_spacing

# # ### Visualizer for boundary functions
# fun = lambda x: b_cond(x)[1]
# plotter.pv_function_plot(x_int, fun)
# sys.exit()

# Definition of pde
lam = 100.
mu = 40.
Em = mu * (3 * lam + 2 * mu) / (lam + mu)
nu = lam / (2 * (lam + mu))
youngs_mod_fun = lambda x: Em
poisson_ratio_fun = lambda x: nu
# pde = models.linear_elasticity_fos(youngs_mod_fun = youngs_mod_fun, poisson_ratio_fun = poisson_ratio_fun, mode='plain strain', spacing=support_radius)
pde = models.hyperelastic_steady_state_fos(models.neo_hooke, youngs_mod_fun = youngs_mod_fun, poisson_ratio_fun = poisson_ratio_fun, mode='plain strain', spacing=support_radius)


### Neighborhood search
(num_neighbors, max_neighbors, min_neighbors, neighbor_list) = utility.search_neighborhood(x_nodes, x_int, support_radius)


### Setting
n_fields = 6
dofs0 = jnp.zeros((n_nodes, n_fields))
static_settings = flax.core.FrozenDict({
  'solution space': ('mls',),
  'shape function mode': 'compiled',
  'number of fields': (n_fields,),
  'assembling mode': ('sparse',),
  'maximal number of neighbors': (max_neighbors,),
  'order of basis functions': (2,),
  'weight function type': ('gaussian',),
  'variational scheme': ('least square pde loss',),
  'model': (pde,),
  'solution structure': ('first order set',),
  'boundary coefficients': (b_coeff,),
  'boundary conditions': (b_cond,),
  'psdf': (psdf,),
  'solver type': 'newton',
  'solver backend': 'pardiso',
  'solver': 'lu',
  'verbose': 2,
  'connectivity': (utility.jnp_to_tuple(neighbor_list),),
})
settings = {
  'node coordinates': x_nodes,
  'integration coordinates': (x_int,),
  'integration weights': (w_int,),
  'support radius': (support_radius,),
  'beta': (3,),
}


# Precompute shape functions
settings = solution_structures.precompile(dofs0, settings, static_settings)

### Load stepping in order to provide better initial guesses for the newton method
dofs = dofs0
multiplier = 0.0
n_step = 5
for i in range(n_step):
  # Update boundary conditions
  print('Start load step ', i+1)
  multiplier = multiplier + 1 / n_step
  load_mult = lambda x: multiplier * load(x)
  bc_2 = [homogeneous, load_mult, homogeneous, homogeneous]
  b_cond = lambda x: jnp.asarray([geometry.transfinite_interpolation(x, psdfs, bc_1),
                                  geometry.transfinite_interpolation(x, psdfs, bc_2)])
  settings['compiled bc'] = (solution_structures.precompute_coupled_boundary(b_cond, b_coeff, x_int, num_diff=1),)
  dofs, _ = solver.solver(dofs, settings, static_settings)



### Preperation of postprocessing data
local_dofs = dofs[neighbor_list]

def post_fun(x, itt, local_dofs, settings, static_settings):
  return solution_structures.solution_structure(x, itt, local_dofs, settings, static_settings, 0)

post_fun_vj = jax.jit(jax.vmap(post_fun, (0,0,0,None,None)), static_argnames=['static_settings'])
data = post_fun_vj(x_int, jnp.arange(n_int), local_dofs, settings, static_settings)

x_vis = x_int + jnp.transpose(jnp.asarray([data[:,0],data[:,1]]))

# Plotting with pyvista
plotter.pv_plot(x_vis, data[:,0], export_vtk=False, show=True, data_on_z=False)
plotter.pv_plot(x_vis, data[:,1], export_vtk=False, show=True, data_on_z=False)
plotter.pv_plot(x_vis, data[:,2], export_vtk=False, show=True, data_on_z=False)
plotter.pv_plot(x_vis, data[:,3], export_vtk=False, show=True, data_on_z=False)
plotter.pv_plot(x_vis, data[:,4], export_vtk=False, show=True, data_on_z=False)
plotter.pv_plot(x_vis, data[:,5], export_vtk=False, show=True, data_on_z=False)
