# 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')


### First order least square Laplace equation with coupled Neumann boundary conditions on circular cutout



# Positive smooth distance functions that select the specific boundary segments
psdf_bottom = lambda x: geometry.psdf_trimmed_line(x, jnp.stack([0.,0.]), jnp.stack([1.,0.]))
psdf_right = lambda x: geometry.psdf_trimmed_line(x, jnp.stack([1.,0.]), jnp.stack([1.,1.]))
psdf_top = lambda x: geometry.psdf_trimmed_line(x, jnp.stack([1.,1.]), jnp.stack([0.,1.]))
psdf_left = lambda x: geometry.psdf_trimmed_line(x, jnp.stack([0.,1.]), jnp.stack([0.,0.]))
psdf_cutout = lambda x: geometry.psdf_nd_sphere_cutout(x, jnp.stack([0.5,0.5]), 0.25)
psdf_rect = lambda x: geometry.psdf_parallelogram(x, jnp.stack([0.,0.]), jnp.stack([0.,1.]), jnp.stack([1.,0.]))
psdfs = [psdf_bottom, psdf_right, psdf_top, psdf_left, psdf_cutout]

# Positive smooth distance funcion of full domain (used for seeding nodes and integration points)
psdf_domain = lambda x: geometry.r_equivalence(psdf_rect(x), psdf_cutout(x))

# Boundary condition functions
def homogeneous(x): # Bottom, right and top boundaries
  return 0
def bc_0(x): # Left boundary
  return jax.lax.sin(jnp.pi * x[0])
def neumann_at_circ(x): # Normal derivative at circular cutout for a specific analytic solution, may be used for convergence study
  return (2. * jnp.pi / jnp.sinh(jnp.pi)) * ((-1.+2.*x[1])*jnp.cosh(jnp.pi-jnp.pi*x[1]) + (1.-2.*x[0])*jnp.cos(jnp.pi*x[0])*jnp.sinh(jnp.pi-jnp.pi*x[1]))

# Field selectors for boundary conditions
def dirichlet(x):
  return jnp.asarray([1., 0., 0.])
def coupled_neumann(x):
  centroid = jnp.asarray([0.5,0.5])
  dxc = x - centroid
  normal = - dxc / jnp.sqrt(jnp.dot(dxc, dxc))
  return jnp.concatenate([jnp.asarray([0.]), normal])

# Assignment of boundary conditions
dirichlet_conditions = [bc_0, homogeneous, homogeneous, homogeneous, neumann_at_circ]
a_i = [dirichlet, dirichlet, dirichlet, dirichlet, coupled_neumann]

### Transfinite interpolation of boundary segments
b_coeff = lambda x: jnp.asarray([geometry.transfinite_interpolation(x, psdfs, a_i)])
b_cond = lambda x: jnp.asarray([geometry.transfinite_interpolation(x, psdfs, dirichlet_conditions)])
psdf = lambda x: jnp.asarray([psdf_domain(x)])



### Initialization and generation of nodes and integration points
# Regular node placement
nod_spacing = 0.025
(x_nodes, n_nodes, region_size) = seeder.regular_in_psdf(psdf_domain, [0.,0.], [1.,1.], nod_spacing)

# Integration points
# # Regular spacing
# int_spacing = nod_spacing
# (x_int, w_int, n_int) = seeder.gauss_points_in_psdf(psdf_domain, [0.,0.], [1.,1.], int_spacing, order=2)
# Collocation
x_int = x_nodes
n_int = n_nodes
w_int = (region_size / n_int) * jnp.ones(n_int)
print("Estimated region size: ", w_int.sum())



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

# Definition of pde
pde = models.poisson_fos(support_radius)

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



### Settings
n_fields = 3
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': (1,),
  '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': 'linear',
  'solver backend': 'pardiso',
  'solver': 'lu',
  'verbose': 1,
  'connectivity': (utility.jnp_to_tuple(neighbor_list),),
})
settings = {
  'beta': (3,),
  'node coordinates': x_nodes,
  'integration coordinates': (x_int,),
  'integration weights': (w_int,),
  'support radius': (support_radius,),
}

### Optional: Precompute shape functions and pass them to e.g. static_settings...
settings = solution_structures.precompile(dofs0, settings, static_settings)


### Call solver
dofs, _ = solver.solver(dofs0, settings, static_settings, tol=1e-8)


### 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)

# Plotting with pyvista
plotter.pv_plot(x_int, data[:,0], export_vtk=False, show=True)
