import FESTIM as F

my_sim = F.Simulation()

my_sim.mesh = F.MeshFromVertices([0, 1, 2, 3, 4])

tungsten = F.Material(id=1, D_0=4.1e-9, E_D=0.3)
copper = F.Material(id=2, D_0=8e-7, E_D=0.4)

my_sim.materials = F.Materials(
	[tungsten, copper]
)

my_sim.boundary_conditions = [
	F.DirichletBC(surfaces=[1, 2], value=2)
]

my_sim.T = F.Temperature(2*F.x + F.t)

my_sim.sources = F.Source(1, volumes=1, field="solute")

my_sim.settings = F.Settings(
	absolute_tolerance=1e-10,
	relative_tolerance=1e-10,
	maximum_iterations=50,
	transient=True,
	final_time=10
)

my_sim.dt= F.Stepsize(0.1)

my_sim.initialise()
my_sim.run()

class Trap(F.Concentration):
    def __init__(
            self, k_0, E_k, p_0, E_p, materials, density, id=None):

class Trap(F.Concentration):
    def __init__(
            self, k_0, E_k, p_0, E_p, materials, density, id=None):
        """Inits Trap
        Args:
            k_0 (float): trapping pre-exponential factor
            E_k (float): trapping activation energy
            p_0 (float): detrapping pre-exponential factor
            E_p (float): detrapping activation energy
            materials (list or int): the materials ids the trap is living in
            density (sp.Add, float): the trap density
            id (int, optional): The trap id. Defaults to None.
        Raises:
            ValueError: if duplicates are found in materials
        """

class CustomSource(F.Source):
    def __init__(self, A, B, volume):
        value = A*F.t + B*(F.x + F.y**2)

        super().__init__(value, volume, field="solute")
        
my_source = CustomSource(A=2, B=4, volume=1)

class Stepsize:

    def adapt(self, t, nb_it, converged):
        """Changes the stepsize based on convergence.
        Args:
            t (float): current time.
            nb_it (int): number of iterations the solver required to converge.
            converged (bool): True if the solver converged, else False.
        """
        change_ratio = self.adaptive_stepsize["stepsize_change_ratio"]
        dt_min = self.adaptive_stepsize["dt_min"]
        stepsize_stop_max = self.adaptive_stepsize["stepsize_stop_max"]
        t_stop = self.adaptive_stepsize["t_stop"]
        if not converged:
            self.value.assign(float(self.value)/change_ratio)
            if float(self.value) < dt_min:
                raise ValueError('stepsize reached minimal value')
        if nb_it < 5:
            self.value.assign(float(self.value)*change_ratio)
        else:
            self.value.assign(float(self.value)/change_ratio)

        if t_stop is not None:
            if t >= t_stop:
                if float(self.value) > stepsize_stop_max:
                    self.value.assign(stepsize_stop_max)

class Stepsize:

    def adapt(self, t, nb_it, converged):
        """Changes the stepsize based on convergence.
        Args:
            t (float): current time.
            nb_it (int): number of iterations the solver required to converge.
            converged (bool): True if the solver converged, else False.
        """
        change_ratio = self.adaptive_stepsize["stepsize_change_ratio"]
        dt_min = self.adaptive_stepsize["dt_min"]
        stepsize_stop_max = self.adaptive_stepsize["stepsize_stop_max"]
        t_stop = self.adaptive_stepsize["t_stop"]
        if not converged:
            self.value.assign(float(self.value)/change_ratio)
            if float(self.value) < dt_min:
                raise ValueError('stepsize reached minimal value')
        if nb_it < 5:
            self.value.assign(float(self.value)*change_ratio)
        else:
            self.value.assign(float(self.value)/change_ratio)

        if t_stop is not None:
            if t >= t_stop:
                if float(self.value) > stepsize_stop_max:
                    self.value.assign(stepsize_stop_max)

class CustomSource(F.Source):
    def __init__(self, A, B, volume):
        value = A*F.t + B*(F.x + F.y**2)

        super().__init__(value, volume, field="solute")
        
my_source = CustomSource(A=2, B=4, volume=1)

class CustomStepsize(F.Stepsize):

    def adapt(self, t, nb_it, converged):
	 

class CustomStepsize(F.Stepsize):

    def adapt(self, t, nb_it, converged):
      print('Adapting stepsize')
      if converged:
      	self.value *= 1.2 * (1 - nb_it/5)
      else:
        self.value *= 1/2

class RecombinationFlux(F.SurfaceFlux):
    def __init__(self, Kr_0, E_Kr, surface):
      super().__init__(field="solute", surface=surface)

      self.Kr_0 = Kr_0
      self.E_Kr = E_Kr

    def compute(self):
      Kr = self.Kr_0*exp(-self.E_Kr/F.k_B/self.T)
      return Kr*self.function**2

class MobileTrap(F.Trap):
  def create_form(self, **kwargs):
    super().create_form(**kwargs)

    self.F += 2 * dot(
      grad(self.solution), grad(self.test_function))*kwargs["dx"]

class LinearMesh(F.MeshFromVertices):
    def __init__(self, N, start_dx, ratio):
        dx = np.empty((N,))
        dx[0] = start_dx
        dx[1:] = ratio
        dx = np.cumprod(dx)
        vertices = np.concatenate([0], np.cumsum(dx))
        super().create_form(vertices)

class CustomStepsize(F.Stepsize):

    def adapt(self, t, nb_it, converged):
      print('Adapting stepsize')
      if converged:
      	self.value *= 1.2 * (1 - nb_it/5)
      else:
        self.value *= 1/2

class RecombinationFlux(F.SurfaceFlux):
    def __init__(self, Kr_0, E_Kr, surface):
      super().__init__(field="solute", surface=surface)

      self.Kr_0 = Kr_0
      self.E_Kr = E_Kr

    def compute(self):
      Kr = self.Kr_0*exp(-self.E_Kr/F.k_B/self.T)
      return Kr*self.function**2

class MobileTrap(F.Trap):
  def create_form(self, **kwargs):
    super().create_form(**kwargs)

    self.F += 2 * dot(
      grad(self.solution), grad(self.test_function))*kwargs["dx"]

class LinearMesh(F.MeshFromVertices):
    def __init__(self, N, start_dx, ratio):
        dx = np.empty((N,))
        dx[0] = start_dx
        dx[1:] = ratio
        dx = np.cumprod(dx)
        vertices = np.concatenate([0], np.cumsum(dx))
        super().create_form(vertices)

my_sim.mesh = my_mesh

my_mesh = MeshFromVertices([0, 1, 2, 3, 4])

my_sim.mesh = my_mesh

my_mesh = MeshFromRefinements(
    initial_number_of_cells=10,
    size=1,
    refinements=[{
        "x": 0.5, "cells": 10
    }]
)

my_sim.mesh = my_mesh

my_mesh = MeshFromXDMF(
    volume_file="mesh.xdmf",
    boundary_file="boundaries.xdmf"
)


my_sim.mesh = my_mesh

my_sim.materials = my_materials

tungsten = F.Material(id=1, D_0=1, E_D=0.1)
copper = F.Material(id=2, D_0=2, E_D=0.2)

my_sim.materials = my_materials

tungsten = F.Material(id=1, D_0=1, E_D=0.1)
copper = F.Material(id=2, D_0=2, E_D=0.2)

my_materials = F.Materials([
  tungsten,
  copper,
])

my_sim.materials = my_materials

my_sim.traps = ...

trap1 = F.Trap(k_0=1, E_k=0.1, p_0=2, E_p=0.8, materials=[1, 2], density=10)
trap2 = F.Trap(k_0=1, E_k=0.1, p_0=3, E_p=1, materials=[1], density=1)


my_sim.traps = [trap1, trap2]

trap1 = F.Trap(k_0=1, E_k=0.1, p_0=2, E_p=0.8, materials=[1, 2], density=10)
trap2 = F.Trap(k_0=1, E_k=0.1, p_0=3, E_p=1, materials=[1], density=1)
trap3 = F.ExtrinsicTrap(
  k_0=1, E_k=0.1, p_0=3, E_p=1, materials=[1],
  form_parameters=...
)


my_sim.traps = [trap1, trap2, trap3]

trap1 = F.Trap(k_0=1, E_k=0.1, p_0=2, E_p=0.8, materials=[1, 2], density=10)
trap2 = F.Trap(k_0=1, E_k=0.1, p_0=3, E_p=1, materials=[1], density=1)
trap3 = F.ExtrinsicTrap(
  k_0=1, E_k=0.1, p_0=3, E_p=1, materials=[1],
  form_parameters=...
)


my_sim.traps = [trap1, trap2, trap3]

my_sim.exports = F.Exports(list_of_exports)

list_of_exports = [
  F.TXTExport("solute", times=[10, 100], label="mobile", folder="out"),
]


my_sim.exports = F.Exports(list_of_exports)

list_of_exports = [
  F.TXTExport("solute", times=[10, 100], label="mobile", folder="out"),
  F.XDMFExport("retention", label="retention", folder="out"),
]


my_sim.exports = F.Exports(list_of_exports)

list_of_exports = [
  F.TXTExport("solute", times=[10, 100], label="mobile", folder="out"),
  F.XDMFExport("retention", label="retention", folder="out"),
  derived_quantities
]


my_sim.exports = F.Exports(list_of_exports)

derived_quantities = F.DerivedQuantities(
  [
    F.AverageVolume("retention", volume=1)
  ]
)

list_of_exports = [
  F.TXTExport("solute", times=[10, 100], label="mobile", folder="out"),
  F.XDMFExport("retention", label="retention", folder="out"),
  derived_quantities
]


my_sim.exports = F.Exports(list_of_exports)

derived_quantities = F.DerivedQuantities(
  [
    F.AverageVolume("retention", volume=1),
    F.HydrogenFlux(surface=2)
  ]
)

list_of_exports = [
  F.TXTExport("solute", times=[10, 100], label="mobile", folder="out"),
  F.XDMFExport("retention", label="retention", folder="out"),
  derived_quantities
]


my_sim.exports = F.Exports(list_of_exports)

derived_quantities = F.DerivedQuantities(
  [
    F.AverageVolume("retention", volume=1),
    F.HydrogenFlux(surface=2),
    F.ThermalFlux(surface=[1, 2])
  ]
)

list_of_exports = [
  F.TXTExport("solute", times=[10, 100], label="mobile", folder="out"),
  F.XDMFExport("retention", label="retention", folder="out"),
  derived_quantities
]


my_sim.exports = F.Exports(list_of_exports)

my_sim.boundary_conditions = list_of_bcs

list_of_bcs = [
  F.DirichletBC(surfaces=1, value=10, component="T"),
]

my_sim.boundary_conditions = list_of_bcs

list_of_bcs = [
  F.DirichletBC(surfaces=1, value=10, component="T"),
  F.SievertsBC(surfaces=[2, 3], S_0=100, E_S=0.1, pressure=1e5)
]

my_sim.boundary_conditions = list_of_bcs

list_of_bcs = [
  F.DirichletBC(surfaces=1, value=10, component="T"),
  F.SievertsBC(surfaces=[2, 3], S_0=100, E_S=0.1, pressure=1e5),
  F.ImplantationDirichlet(surfaces=4, phi=10, R_p=5e-9, D_0=1, E_D=0.1)
]

my_sim.boundary_conditions = list_of_bcs

def bc(T, cm, K):
  return K*T*cm

list_of_bcs = [
  F.DirichletBC(surfaces=1, value=10, component="T"),
  F.SievertsBC(surfaces=[2, 3], S_0=100, E_S=0.1, pressure=1e5),
  F.ImplantationDirichlet(surfaces=4, phi=10, R_p=5e-9, D_0=1, E_D=0.1)
]

my_sim.boundary_conditions = list_of_bcs

def bc(T, cm, K):
  return K*T*cm

list_of_bcs = [
  F.DirichletBC(surfaces=1, value=10, component="T"),
  F.SievertsBC(surfaces=[2, 3], S_0=100, E_S=0.1, pressure=1e5),
  F.ImplantationDirichlet(surfaces=4, phi=10, R_p=5e-9, D_0=1, E_D=0.1),
  F.CustomDirichlet(surfaces=5, function=bc, component="solute", K=2*F.t)
]

my_sim.boundary_conditions = list_of_bcs

list_of_bcs = [
  F.DirichletBC(surfaces=1, value=10, component="T"),
  F.SievertsBC(surfaces=[2, 3], S_0=100, E_S=0.1, pressure=1e5),
  F.ImplantationDirichlet(surfaces=4, phi=10, R_p=5e-9, D_0=1, E_D=0.1),
  F.CustomDirichlet(surfaces=5, function=bc, component="solute", K=2*F.t),
  F.ConvectiveFlux(surfaces=6, h_coeff=10, T_ext=350)
]

my_sim.boundary_conditions = list_of_bcs

list_of_bcs = [
  F.DirichletBC(surfaces=1, value=10, component="T"),
  F.SievertsBC(surfaces=[2, 3], S_0=100, E_S=0.1, pressure=1e5),
  F.ImplantationDirichlet(surfaces=4, phi=10, R_p=5e-9, D_0=1, E_D=0.1),
  F.CustomDirichlet(surfaces=5, function=bc, component="solute", K=2*F.t),
  F.ConvectiveFlux(surfaces=6, h_coeff=10, T_ext=350),
  F.RecombinationFlux(surfaces=7, Kr_0=10, E_Kr=0.1, order=2)
]

my_sim.boundary_conditions = list_of_bcs

list_of_bcs = [
  F.DirichletBC(surfaces=1, value=10, component="T"),
  F.SievertsBC(surfaces=[2, 3], S_0=100, E_S=0.1, pressure=1e5),
  F.ImplantationDirichlet(surfaces=4, phi=10, R_p=5e-9, D_0=1, E_D=0.1),
  F.CustomDirichlet(surfaces=5, function=bc, component="solute", K=2*F.t),
  F.ConvectiveFlux(surfaces=6, h_coeff=10, T_ext=350),
  F.RecombinationFlux(surfaces=7, Kr_0=10, E_Kr=0.1, order=2),
  F.FluxBC(surfaces=8, value=10 + sp.sin(F.t), component="solute")
]

my_sim.boundary_conditions = list_of_bcs

my_sim.T = my_temperature

my_temperature = F.Temperature(350)

my_sim.T = my_temperature

my_temperature = F.HeatTransferProblem(transient=True, initial_value=350)

my_sim.T = my_temperature

my_sim.settings = my_settings

my_settings = F.Settings(
  absolute_tolerance=1e10,
  relative_tolerance=1e-10,
)

my_sim.settings = my_settings

my_settings = F.Settings(
  absolute_tolerance=1e10,
  relative_tolerance=1e-10,
  maximum_iterations=30,
  transient=True,
  final_time=10,
  chemical_pot=False,
  soret=False,
  traps_element_type="CG",
  update_jacobian=True,
  completion_tone=False
)

my_sim.settings = my_settings

def __init__(
  self,
  absolute_tolerance,
  relative_tolerance,
  maximum_iterations=30,
  transient=True,
  final_time=None,
  chemical_pot=False,
  soret=False,
  traps_element_type="CG",
  update_jacobian=True,
  completion_tone=False
):
"""Inits Settings
Args:
  absolute_tolerance (float): the absolute tolerance of the newton
      solver
  relative_tolerance (float): the relative tolerance of the newton
      solver
  maximum_iterations (int, optional): maximum iterations allowed for
      the solver to converge. Defaults to 30.
  transient (bool, optional): If set to True, the simulation will be
      transient. Defaults to True.
  final_time (float, optional): final time of the simulation.
      Defaults to None.
  chemical_pot (bool, optional): if True, conservation of chemical
      potential will be assumed. Defaults to False.
  soret (bool, optional): if True, Soret effect will be assumed.
      Defaults to False.
  traps_element_type (str, optional): Finite element used for traps.
      If traps densities are discontinuous (eg. different materials)
      "DG" is recommended. Defaults to "CG".
  update_jacobian (bool, optional): If set to False, the Jacobian of
      the formulation will be computed only once at the beggining.
      Else it will be computed at each time step. Defaults to True.
  completion_tone (bool, optional): If True, a native os alert
      tone will alert user upon completion of current run. Defaults
      to False.
        """

my_settings = F.Settings(
  absolute_tolerance=1e10,
  relative_tolerance=1e-10,
  maximum_iterations=30,
  transient=True,
  final_time=10,
  chemical_pot=False,
  soret=False,
  traps_element_type="CG",
  update_jacobian=True,
  completion_tone=False
)

my_sim.settings = my_settings

def __init__(
  self,
  absolute_tolerance,
  relative_tolerance,
  maximum_iterations=30,
  transient=True,
  final_time=None,
  chemical_pot=False,
  soret=False,
  traps_element_type="CG",
  update_jacobian=True,
  completion_tone=False
):
"""Inits Settings
Args:
  absolute_tolerance (float): the absolute tolerance of the newton
      solver
  relative_tolerance (float): the relative tolerance of the newton
      solver
  maximum_iterations (int, optional): maximum iterations allowed for
      the solver to converge. Defaults to 30.
  transient (bool, optional): If set to True, the simulation will be
      transient. Defaults to True.
  final_time (float, optional): final time of the simulation.
      Defaults to None.
  chemical_pot (bool, optional): if True, conservation of chemical
      potential will be assumed. Defaults to False.
  soret (bool, optional): if True, Soret effect will be assumed.
      Defaults to False.
  traps_element_type (str, optional): Finite element used for traps.
      If traps densities are discontinuous (eg. different materials)
      "DG" is recommended. Defaults to "CG".
  update_jacobian (bool, optional): If set to False, the Jacobian of
      the formulation will be computed only once at the beggining.
      Else it will be computed at each time step. Defaults to True.
  completion_tone (bool, optional): If True, a native os alert
      tone will alert user upon completion of current run. Defaults
      to False.
        """

my_sim.dt = my_stepsize

my_stepsize = F.Stepsize(
  initial_value=1.2,
  stepsize_change_change_ratio=1.2,
  dt_min=1e-8
)

my_sim.dt = my_stepsize

my_stepsize = F.Stepsize(
  initial_value=1.2,
  stepsize_change_change_ratio=1.2,
  dt_min=1e-8,
  t_stop=100,
  stepsize_stop_max=2
)

my_sim.dt = my_stepsize