# Copyright 2024 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Update routines for core_profiles.

Set of routines that relate to updating/creating core profiles from existing
core profiles.

Includes:
- get_prescribed_core_profile_values: Updates core profiles which are not being
  evolved by PDE.
- update_core_profiles_during_step: Intra-step updates of the evolving variables
  in core profiles during solver iterations.
- update_all_core_profiles_after_step: Updates all core_profiles after a step.
  Includes the evolved variables and derived variables like q_face, psidot, etc.
- compute_boundary_conditions_for_t_plus_dt: Computes boundary conditions for
  the next timestep and returns updates to State.
"""
import dataclasses
import functools

import jax
from jax import numpy as jnp
from torax._src import array_typing
from torax._src import jax_utils
from torax._src import state
from torax._src.config import runtime_params_slice
from torax._src.core_profiles import convertors
from torax._src.core_profiles import getters
from torax._src.fvm import cell_variable
from torax._src.geometry import geometry
from torax._src.neoclassical import neoclassical_models as neoclassical_models_lib
from torax._src.physics import formulas
from torax._src.physics import psi_calculations
from torax._src.sources import source_models as source_models_lib
from torax._src.sources import source_profile_builders
from torax._src.sources import source_profiles as source_profiles_lib

_trapz = jax.scipy.integrate.trapezoid

# pylint: disable=invalid-name


def _calculate_psi_value_constraint_from_v_loop(
    dt: array_typing.FloatScalar,
    theta: array_typing.FloatScalar,
    v_loop_lcfs_t: array_typing.FloatScalar,
    v_loop_lcfs_t_plus_dt: array_typing.FloatScalar,
    psi_lcfs_t: array_typing.FloatScalar,
) -> jax.Array:
  """Calculates the value constraint on the poloidal flux for the next time step from loop voltage."""
  theta_weighted_v_loop_lcfs = (
      1 - theta
  ) * v_loop_lcfs_t + theta * v_loop_lcfs_t_plus_dt
  return psi_lcfs_t + theta_weighted_v_loop_lcfs * dt


@jax_utils.jit
def get_prescribed_core_profile_values(
    runtime_params: runtime_params_slice.RuntimeParams,
    geo: geometry.Geometry,
    core_profiles: state.CoreProfiles,
) -> dict[str, array_typing.FloatVector]:
  """Updates core profiles which are not being evolved by PDE.

  Uses same functions as for profile initialization.

  Args:
    runtime_params: Runtime parameters at t=t_initial.
    geo: Torus geometry.
    core_profiles: Core profiles dataclass to be updated

  Returns:
    Updated core profiles values on the cell grid.
  """
  # If profiles are not evolved, they can still potential be time-evolving,
  # depending on the runtime params. If so, they are updated below.
  if not runtime_params.numerics.evolve_ion_heat:
    T_i = getters.get_updated_ion_temperature(
        runtime_params.profile_conditions, geo
    ).value
  else:
    T_i = core_profiles.T_i.value
  if not runtime_params.numerics.evolve_electron_heat:
    T_e_cell_variable = getters.get_updated_electron_temperature(
        runtime_params.profile_conditions, geo
    )
    T_e = T_e_cell_variable.value
  else:
    T_e_cell_variable = core_profiles.T_e
    T_e = T_e_cell_variable.value
  if not runtime_params.numerics.evolve_density:
    n_e_cell_variable = getters.get_updated_electron_density(
        runtime_params.profile_conditions, geo
    )
  else:
    n_e_cell_variable = core_profiles.n_e
  ions = getters.get_updated_ions(
      runtime_params,
      geo,
      n_e_cell_variable,
      T_e_cell_variable,
  )
  n_e = n_e_cell_variable.value
  n_i = ions.n_i.value
  n_impurity = ions.n_impurity.value
  impurity_fractions = ions.impurity_fractions

  return {
      'T_i': T_i,
      'T_e': T_e,
      'n_e': n_e,
      'n_i': n_i,
      'n_impurity': n_impurity,
      'impurity_fractions': impurity_fractions,
      'Z_i': ions.Z_i,
      'Z_i_face': ions.Z_i_face,
      'Z_impurity': ions.Z_impurity,
      'Z_impurity_face': ions.Z_impurity_face,
      'A_i': ions.A_i,
      'A_impurity': ions.A_impurity,
      'A_impurity_face': ions.A_impurity_face,
      'Z_eff': ions.Z_eff,
      'Z_eff_face': ions.Z_eff_face,
  }


@functools.partial(jax_utils.jit, static_argnames=['evolving_names'])
def update_core_profiles_during_step(
    x_new: tuple[cell_variable.CellVariable, ...],
    runtime_params: runtime_params_slice.RuntimeParams,
    geo: geometry.Geometry,
    core_profiles: state.CoreProfiles,
    evolving_names: tuple[str, ...],
) -> state.CoreProfiles:
  """Returns the new core profiles after updating the evolving variables.

  Intended for use during iterative solves in the step function. Only updates
  the core profiles which are being evolved by the PDE and directly derivable
  quantities like q_face, s_face. core_profile calculations which require
  sources are not updated.

  Args:
    x_new: The new values of the evolving variables.
    runtime_params: The runtime params slice.
    geo: Magnetic geometry.
    core_profiles: The old set of core plasma profiles.
    evolving_names: The names of the evolving variables.
  """

  updated_core_profiles = convertors.solver_x_tuple_to_core_profiles(
      x_new, evolving_names, core_profiles
  )

  ions = getters.get_updated_ions(
      runtime_params,
      geo,
      updated_core_profiles.n_e,
      updated_core_profiles.T_e,
  )

  return dataclasses.replace(
      updated_core_profiles,
      n_i=ions.n_i,
      n_impurity=ions.n_impurity,
      impurity_fractions=ions.impurity_fractions,
      Z_i=ions.Z_i,
      Z_i_face=ions.Z_i_face,
      Z_impurity=ions.Z_impurity,
      Z_impurity_face=ions.Z_impurity_face,
      A_i=ions.A_i,
      A_impurity=ions.A_impurity,
      A_impurity_face=ions.A_impurity_face,
      Z_eff=ions.Z_eff,
      Z_eff_face=ions.Z_eff_face,
      q_face=psi_calculations.calc_q_face(geo, updated_core_profiles.psi),
      s_face=psi_calculations.calc_s_face(geo, updated_core_profiles.psi),
  )


def update_core_and_source_profiles_after_step(
    dt: array_typing.FloatScalar,
    x_new: tuple[cell_variable.CellVariable, ...],
    runtime_params_t_plus_dt: runtime_params_slice.RuntimeParams,
    geo: geometry.Geometry,
    core_profiles_t: state.CoreProfiles,
    core_profiles_t_plus_dt: state.CoreProfiles,
    explicit_source_profiles: source_profiles_lib.SourceProfiles,
    source_models: source_models_lib.SourceModels,
    neoclassical_models: neoclassical_models_lib.NeoclassicalModels,
    evolving_names: tuple[str, ...],
) -> tuple[state.CoreProfiles, source_profiles_lib.SourceProfiles]:
  """Returns a core profiles and source profiles after the solver has finished.

  Updates the evolved variables and derived variables like q_face, psidot, etc.

  Args:
    dt: The size of the last timestep.
    x_new: The new values of the evolving variables.
    runtime_params_t_plus_dt: The runtime params slice at t=t_plus_dt.
    geo: Magnetic geometry.
    core_profiles_t: The old set of core plasma profiles.
    core_profiles_t_plus_dt: The partially new set of core plasma profiles. On
      input into this function, all prescribed profiles and used boundary
      conditions are already set. But evolving values are not.
    explicit_source_profiles: The explicit source profiles.
    source_models: The source models.
    neoclassical_models: The neoclassical models.
    evolving_names: The names of the evolving variables.

  Returns:
    A tuple of the new core profiles and the source profiles.
  """

  updated_core_profiles_t_plus_dt = convertors.solver_x_tuple_to_core_profiles(
      x_new, evolving_names, core_profiles_t_plus_dt
  )

  ions = getters.get_updated_ions(
      runtime_params_t_plus_dt,
      geo,
      updated_core_profiles_t_plus_dt.n_e,
      updated_core_profiles_t_plus_dt.T_e,
  )

  v_loop_lcfs = (
      runtime_params_t_plus_dt.profile_conditions.v_loop_lcfs  # pylint: disable=g-long-ternary
      if runtime_params_t_plus_dt.profile_conditions.use_v_loop_lcfs_boundary_condition
      else _update_v_loop_lcfs_from_psi(
          core_profiles_t.psi,
          updated_core_profiles_t_plus_dt.psi,
          dt,
      )
  )

  j_total, j_total_face, Ip_profile_face = psi_calculations.calc_j_total(
      geo,
      updated_core_profiles_t_plus_dt.psi,
  )

  # A wholly new core profiles object is defined as a guard against neglecting
  # to update one of the attributes if doing dataclasses.replace
  intermediate_core_profiles = state.CoreProfiles(
      T_i=updated_core_profiles_t_plus_dt.T_i,
      T_e=updated_core_profiles_t_plus_dt.T_e,
      psi=updated_core_profiles_t_plus_dt.psi,
      n_e=updated_core_profiles_t_plus_dt.n_e,
      n_i=ions.n_i,
      n_impurity=ions.n_impurity,
      impurity_fractions=ions.impurity_fractions,
      Z_i=ions.Z_i,
      Z_i_face=ions.Z_i_face,
      Z_impurity=ions.Z_impurity,
      Z_impurity_face=ions.Z_impurity_face,
      psidot=core_profiles_t_plus_dt.psidot,
      q_face=psi_calculations.calc_q_face(
          geo, updated_core_profiles_t_plus_dt.psi
      ),
      s_face=psi_calculations.calc_s_face(
          geo, updated_core_profiles_t_plus_dt.psi
      ),
      A_i=ions.A_i,
      A_impurity=ions.A_impurity,
      A_impurity_face=ions.A_impurity_face,
      Z_eff=ions.Z_eff,
      Z_eff_face=ions.Z_eff_face,
      v_loop_lcfs=v_loop_lcfs,
      sigma=core_profiles_t_plus_dt.sigma,  # Not yet updated
      sigma_face=core_profiles_t_plus_dt.sigma_face,  # Not yet updated
      j_total=j_total,
      j_total_face=j_total_face,
      Ip_profile_face=Ip_profile_face,
  )

  conductivity = neoclassical_models.conductivity.calculate_conductivity(
      geo, intermediate_core_profiles
  )

  intermediate_core_profiles = dataclasses.replace(
      intermediate_core_profiles,
      sigma=conductivity.sigma,
      sigma_face=conductivity.sigma_face,
  )

  # build_source_profiles calculates the union with explicit + implicit
  total_source_profiles = source_profile_builders.build_source_profiles(
      runtime_params=runtime_params_t_plus_dt,
      geo=geo,
      source_models=source_models,
      neoclassical_models=neoclassical_models,
      core_profiles=intermediate_core_profiles,
      explicit=False,
      explicit_source_profiles=explicit_source_profiles,
      conductivity=conductivity,
  )

  if (
      not runtime_params_t_plus_dt.numerics.evolve_current
      and runtime_params_t_plus_dt.profile_conditions.psidot
      is not None
  ):
    # If psidot is prescribed and current does not evolve, use prescribed value
    psidot_value = (
        runtime_params_t_plus_dt.profile_conditions.psidot
    )
  else:
    # Otherwise, calculate psidot from psi sources.
    psi_sources = total_source_profiles.total_psi_sources(geo)
    psidot_value = psi_calculations.calculate_psidot_from_psi_sources(
        psi_sources=psi_sources,
        sigma=intermediate_core_profiles.sigma,
        resistivity_multiplier=runtime_params_t_plus_dt.numerics.resistivity_multiplier,
        psi=intermediate_core_profiles.psi,
        geo=geo,
    )
  psidot = dataclasses.replace(
      core_profiles_t_plus_dt.psidot,
      value=psidot_value,
      right_face_constraint=v_loop_lcfs,
      right_face_grad_constraint=None,
  )

  core_profiles_t_plus_dt = dataclasses.replace(
      intermediate_core_profiles,
      psidot=psidot,
  )
  return core_profiles_t_plus_dt, total_source_profiles


def compute_boundary_conditions_for_t_plus_dt(
    dt: array_typing.FloatScalar,
    runtime_params_t: runtime_params_slice.RuntimeParams,
    runtime_params_t_plus_dt: runtime_params_slice.RuntimeParams,
    geo_t_plus_dt: geometry.Geometry,
    core_profiles_t: state.CoreProfiles,
) -> dict[str, dict[str, jax.Array | None]]:
  """Computes boundary conditions for the next timestep and returns updates to State.

  Args:
    dt: Size of the next timestep
    runtime_params_t: Runtime parameters for the current timestep. Will not be
      used if runtime_params_t.profile_conditions.v_loop_lcfs is
      None, i.e. if the dirichlet psi boundary condition based on Ip is used
    runtime_params_t_plus_dt: Runtime parameters for the next timestep
    geo_t_plus_dt: Geometry object for the next timestep
    core_profiles_t: Core profiles at the current timestep. Will not be used if
      runtime_params_t_plus_dt.profile_conditions.v_loop_lcfs is
      None, i.e. if the dirichlet psi boundary condition based on Ip is used

  Returns:
    Mapping from State attribute names to dictionaries updating attributes of
    each CellVariable in the state. This dict can in theory recursively replace
    values in a State object.
  """
  profile_conditions_t_plus_dt = (
      runtime_params_t_plus_dt.profile_conditions
  )
  # TODO(b/390143606): Separate out the boundary condition calculation from the
  # core profile calculation.
  n_e = getters.get_updated_electron_density(
      profile_conditions_t_plus_dt, geo_t_plus_dt
  )
  n_e_right_bc = n_e.right_face_constraint

  # Used for edge calculations and input arguments have correct edge BCs.
  ions_edge = getters.get_updated_ions(
      runtime_params_t_plus_dt,
      geo_t_plus_dt,
      dataclasses.replace(
          core_profiles_t.n_e,
          right_face_constraint=profile_conditions_t_plus_dt.n_e_right_bc,
      ),
      dataclasses.replace(
          core_profiles_t.T_e,
          right_face_constraint=profile_conditions_t_plus_dt.T_e_right_bc,
      ),
  )

  Z_i_edge = ions_edge.Z_i_face[-1]
  Z_impurity_edge = ions_edge.Z_impurity_face[-1]

  dilution_factor_edge = formulas.calculate_main_ion_dilution_factor(
      Z_i_edge,
      Z_impurity_edge,
      runtime_params_t_plus_dt.plasma_composition.Z_eff_face[-1],
  )

  n_i_bound_right = n_e_right_bc * dilution_factor_edge
  n_impurity_bound_right = (
      n_e_right_bc - n_i_bound_right * Z_i_edge
  ) / Z_impurity_edge

  return {
      'T_i': dict(
          left_face_grad_constraint=jnp.zeros(()),
          right_face_grad_constraint=None,
          right_face_constraint=profile_conditions_t_plus_dt.T_i_right_bc,
      ),
      'T_e': dict(
          left_face_grad_constraint=jnp.zeros(()),
          right_face_grad_constraint=None,
          right_face_constraint=profile_conditions_t_plus_dt.T_e_right_bc,
      ),
      'n_e': dict(
          left_face_grad_constraint=jnp.zeros(()),
          right_face_grad_constraint=None,
          right_face_constraint=jnp.array(n_e_right_bc),
      ),
      'n_i': dict(
          left_face_grad_constraint=jnp.zeros(()),
          right_face_grad_constraint=None,
          right_face_constraint=jnp.array(n_i_bound_right),
      ),
      'n_impurity': dict(
          left_face_grad_constraint=jnp.zeros(()),
          right_face_grad_constraint=None,
          right_face_constraint=jnp.array(n_impurity_bound_right),
      ),
      'psi': dict(
          right_face_grad_constraint=(
              psi_calculations.calculate_psi_grad_constraint_from_Ip(  # pylint: disable=g-long-ternary
                  Ip=profile_conditions_t_plus_dt.Ip,
                  geo=geo_t_plus_dt,
              )
              if not runtime_params_t.profile_conditions.use_v_loop_lcfs_boundary_condition
              else None
          ),
          right_face_constraint=(
              _calculate_psi_value_constraint_from_v_loop(  # pylint: disable=g-long-ternary
                  dt=dt,
                  v_loop_lcfs_t=runtime_params_t.profile_conditions.v_loop_lcfs,
                  v_loop_lcfs_t_plus_dt=profile_conditions_t_plus_dt.v_loop_lcfs,
                  psi_lcfs_t=core_profiles_t.psi.right_face_constraint,
                  theta=runtime_params_t.solver.theta_implicit,
              )
              if runtime_params_t.profile_conditions.use_v_loop_lcfs_boundary_condition
              else None
          ),
      ),
      'Z_i_edge': Z_i_edge,
      'Z_impurity_edge': Z_impurity_edge,
  }


def provide_core_profiles_t_plus_dt(
    dt: jax.Array,
    runtime_params_t: runtime_params_slice.RuntimeParams,
    runtime_params_t_plus_dt: runtime_params_slice.RuntimeParams,
    geo_t_plus_dt: geometry.Geometry,
    core_profiles_t: state.CoreProfiles,
) -> state.CoreProfiles:
  """Provides state at t_plus_dt with new boundary conditions and prescribed profiles."""
  updated_boundary_conditions = compute_boundary_conditions_for_t_plus_dt(
      dt=dt,
      runtime_params_t=runtime_params_t,
      runtime_params_t_plus_dt=runtime_params_t_plus_dt,
      geo_t_plus_dt=geo_t_plus_dt,
      core_profiles_t=core_profiles_t,
  )
  updated_values = get_prescribed_core_profile_values(
      runtime_params=runtime_params_t_plus_dt,
      geo=geo_t_plus_dt,
      core_profiles=core_profiles_t,
  )
  T_i = dataclasses.replace(
      core_profiles_t.T_i,
      value=updated_values['T_i'],
      **updated_boundary_conditions['T_i'],
  )
  T_e = dataclasses.replace(
      core_profiles_t.T_e,
      value=updated_values['T_e'],
      **updated_boundary_conditions['T_e'],
  )
  psi = dataclasses.replace(
      core_profiles_t.psi, **updated_boundary_conditions['psi']
  )
  n_e = dataclasses.replace(
      core_profiles_t.n_e,
      value=updated_values['n_e'],
      **updated_boundary_conditions['n_e'],
  )
  n_i = dataclasses.replace(
      core_profiles_t.n_i,
      value=updated_values['n_i'],
      **updated_boundary_conditions['n_i'],
  )
  n_impurity = dataclasses.replace(
      core_profiles_t.n_impurity,
      value=updated_values['n_impurity'],
      **updated_boundary_conditions['n_impurity'],
  )

  # pylint: disable=invalid-name
  # Update Z_face with boundary condition Z, needed for cases where T_e
  # is evolving and updated_prescribed_core_profiles is a no-op.
  Z_i_face = jnp.concatenate(
      [
          updated_values['Z_i_face'][:-1],
          jnp.array([updated_boundary_conditions['Z_i_edge']]),
      ],
  )
  Z_impurity_face = jnp.concatenate(
      [
          updated_values['Z_impurity_face'][:-1],
          jnp.array([updated_boundary_conditions['Z_impurity_edge']]),
      ],
  )
  # pylint: enable=invalid-name
  core_profiles_t_plus_dt = dataclasses.replace(
      core_profiles_t,
      T_i=T_i,
      T_e=T_e,
      psi=psi,
      n_e=n_e,
      n_i=n_i,
      n_impurity=n_impurity,
      impurity_fractions=updated_values['impurity_fractions'],
      Z_i=updated_values['Z_i'],
      Z_i_face=Z_i_face,
      Z_impurity=updated_values['Z_impurity'],
      Z_impurity_face=Z_impurity_face,
      A_i=updated_values['A_i'],
      A_impurity=updated_values['A_impurity'],
      A_impurity_face=updated_values['A_impurity_face'],
      Z_eff=updated_values['Z_eff'],
      Z_eff_face=updated_values['Z_eff_face'],
  )
  return core_profiles_t_plus_dt


# TODO(b/406173731): Find robust solution for underdetermination and solve this
# for general theta_implicit values.
def _update_v_loop_lcfs_from_psi(
    psi_t: cell_variable.CellVariable,
    psi_t_plus_dt: cell_variable.CellVariable,
    dt: array_typing.FloatScalar,
) -> jax.Array:
  """Updates the v_loop_lcfs for the next timestep.

  For the Ip boundary condition case, the v_loop_lcfs formula is in principle
  calculated from:

  (psi_lcfs_t_plus_dt - psi_lcfs_t) / dt =
    v_loop_lcfs_t_plus_dt*theta_implicit - v_loop_lcfs_t*(1-theta_implicit)

  However this set of equations is underdetermined. We thus restrict this
  calculation assuming a fully implicit system, i.e. theta_implicit=1, which is
  the TORAX default. Be cautious when interpreting the results of this function
  when theta_implicit != 1 (non-standard usage).

  Args:
    psi_t: The psi CellVariable at the beginning of the timestep interval.
    psi_t_plus_dt: The updated psi CellVariable for the next timestep.
    dt: The size of the last timestep.

  Returns:
    The updated v_loop_lcfs for the next timestep.
  """
  psi_lcfs_t = psi_t.face_value()[-1]
  psi_lcfs_t_plus_dt = psi_t_plus_dt.face_value()[-1]
  v_loop_lcfs_t_plus_dt = (psi_lcfs_t_plus_dt - psi_lcfs_t) / dt
  return v_loop_lcfs_t_plus_dt
