import jax.numpy as jnp
import numpy as np

from fbpinns.domains import RectangularDomainND
from fbpinns.problems import WaveEquationConstantVelocity3D
from fbpinns.networks import FCN
from fbpinns.constants import Constants
from fbpinns.trainers import PINNTrainer


class WaveEquationInterfaceVelocity3D(WaveEquationConstantVelocity3D):

    @staticmethod
    def init_params(c0=1, c1=2,
                    source=np.array([[-0.5, 0., 0.1, 1.]])):

        static_params = {
            "dims":(1,3),
            "c0":c0,
            "c1":c1,
            "c_fn":WaveEquationInterfaceVelocity3D.c_fn,# velocity function
            "source":jnp.array(source),# location, width and amplitude of initial gaussian sources (k, 4)
            }
        return static_params, {}
    
    @staticmethod
    def c_fn(all_params, x_batch):
        "Computes the velocity model"
        
        p = all_params["static"]["problem"]
        c0, c1 = p["c0"], p["c1"]
        
        x = x_batch[:,0:1]
        c = c0 + (c1-c0)*(1+jnp.tanh(x/0.1))/2
        
        return c

c = Constants(
    domain=RectangularDomainND,
    domain_init_kwargs=dict(
        xmin=np.array([-1,-1,0]),
        xmax=np.array([1,1,1]),
    ),
    problem=WaveEquationInterfaceVelocity3D,
    problem_init_kwargs=dict(
        c0=1, 
        c1=2,
    ),
    network=FCN,
    network_init_kwargs=dict(
        layer_sizes=[3,32,32,32,1],
    ),
    decomposition_init_kwargs=dict(
        unnorm=(0.,0.1),
    ),
    ns=((50,50,50),),
    n_test=(100,100,5),
    n_steps=15000,
    clear_output=True,
)

run = PINNTrainer(c)
_ = run.train()