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

# pylint: disable=invalid-name

"""Cyclotron radiation heat sink for electron heat equation.."""
import dataclasses
from typing import Annotated, ClassVar, Literal

import chex
import jax
from jax import numpy as jnp
import pydantic
from torax._src import array_typing
from torax._src import jax_utils
from torax._src import math_utils
from torax._src import state
from torax._src.config import runtime_params_slice
from torax._src.geometry import geometry
from torax._src.neoclassical.conductivity import base as conductivity_base
from torax._src.sources import base
from torax._src.sources import runtime_params as runtime_params_lib
from torax._src.sources import source
from torax._src.sources import source_profiles
from torax._src.torax_pydantic import torax_pydantic
import typing_extensions

# Default value for the model function to be used for the Cyclotron radiation
# heat sink source. This is also used as an identifier for the model function in
# the default source config for Pydantic to "discriminate" against.
DEFAULT_MODEL_FUNCTION_NAME: str = 'albajar_artaud'


@jax.tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class RuntimeParams(runtime_params_lib.RuntimeParams):
  wall_reflection_coeff: array_typing.FloatScalar
  beta_min: array_typing.FloatScalar
  beta_max: array_typing.FloatScalar
  beta_grid_size: int = dataclasses.field(metadata={'static': True})


def _alpha_closed_form(
    *,
    beta: array_typing.FloatScalar,
    rho_norm: array_typing.FloatVector,
    profile_data: array_typing.FloatVector,
    profile_edge_value: array_typing.FloatScalar,
) -> array_typing.FloatScalar:
  """Returns analytical closed form of alpha for parameterized profiles.

  See albajar_artaud for more details.

  We find alpha for the best fit for either the n_e_data or te_data to the
  parameterized functions:

  n_e = n_e_data[0]*(1 - rhonorm**2)**alpha
  te = (te_data[0] - te_edge)*(1 - rhonorm**beta)**alpha + te_edge

  Where n_e_edge = 0 above. The parameterizations are from the Albajar paper.

  The fit is from the magnetic axis to rhonorm=0.9, to avoid pedestal effects.

  We analytically solve for alpha by taking the derivative of the loss function
  with respect to alpha.

  Defining:
  T = (profile_data - profile_edge_data)/(profile_data[0] - profile_edge_data)

  The loss function is:
  loss = sum((log(T) -  alpha * log(1 - rhonorm**beta)) ** 2) / 2

  The solution is provided by setting d loss / d alpha = 0.

  Args:
    beta: The beta parameter to use in the parameterized functions. For the
      density fit, this is always 2.
    rho_norm: Normalized toroidal flux coordinate.
    profile_data: The profile data to be fit (i.e. density or temperature).
    profile_edge_value: The value of profile data at the edge to use. This can't
      be taken from profile_data since there are different assumptions for the
      temperature and density fits.

  Returns:
    The alpha parameter being fit for either the density or temperature fits.
  """
  # To avoid dealing with slicing of non-concrete values, we do a masking trick
  # where we replace the values of n_e_data and rhonorm above 0.9 with values
  # that will not contribute to the sums in the numerator and denominator.

  mask = rho_norm < 0.9
  profile_data_norm = (profile_data - profile_edge_value) / (
      profile_data[0] - profile_edge_value
  )
  sliced_profile_data_norm = jnp.where(mask, profile_data_norm, 1.0)
  sliced_rhonorm = jnp.where(mask, rho_norm, 0.0)

  num = jnp.sum(
      (jnp.log(sliced_profile_data_norm) * jnp.log(1 - sliced_rhonorm**beta))
  )
  den = jnp.sum(jnp.log(1 - sliced_rhonorm**beta) ** 2)

  alpha_n = num / den
  return alpha_n


def _loss_for_beta_t(
    beta_t: array_typing.FloatScalar,
    rho_norm: array_typing.FloatVector,
    te_data: array_typing.FloatVector,
) -> array_typing.FloatScalar:
  """Returns the loss function for the temperature fit for a given beta_t.

  The fit is from the magnetic axis to rhonorm=0.9, to avoid pedestal effects.
  alpha_t is calculated with an analytical closed form solution.

  Args:
    beta_t: The beta parameter to use in the parameterized functions.
    rho_norm: Normalized toroidal flux coordinate.
    te_data: The temperature data to be fit, assumed to be on the face grid.

  Returns:
    The loss function for the temperature fit for a given beta_t.
  """
  alpha_t = _alpha_closed_form(
      profile_data=te_data,
      profile_edge_value=te_data[-1],
      rho_norm=rho_norm,
      beta=beta_t,
  )
  return _te_loss_fn(
      alpha_t=alpha_t,
      beta_t=beta_t,
      rho_norm=rho_norm,
      te_data=te_data,
  )


def _te_pred_fn(
    *,
    alpha_t: array_typing.FloatScalar,
    beta_t: array_typing.FloatScalar,
    rho_norm: array_typing.FloatVector,
    te_data: array_typing.FloatVector,
) -> array_typing.FloatVector:
  return (te_data[0] - te_data[-1]) * (
      (1 - rho_norm**beta_t)
  ) ** alpha_t + te_data[-1]


def _te_loss_fn(
    *,
    alpha_t: array_typing.FloatScalar,
    beta_t: array_typing.FloatScalar,
    rho_norm: array_typing.FloatVector,
    te_data: array_typing.FloatVector,
) -> array_typing.FloatScalar:
  """Returns the loss function for the temperature fit.

  The fit is from the magnetic axis to rhonorm=0.9, to avoid pedestal effects.
  The pedestal will not extend to rhonorm=0.9. The choice of this value is from
  the Artaud paper.

  Args:
    alpha_t: The alpha parameter to use in the parameterized functions.
    beta_t: The beta parameter to use in the parameterized functions.
    rho_norm: Normalized toroidal flux coordinate.
    te_data: The temperature data to be fit, assumed to be on the face grid.

  Returns:
    The loss function for the temperature fit.
  """
  te_pred = _te_pred_fn(
      alpha_t=alpha_t,
      beta_t=beta_t,
      rho_norm=rho_norm,
      te_data=te_data,
  )
  mask = rho_norm < 0.9
  sliced_diff = jnp.where(mask, te_pred - te_data, 0.0)
  return jnp.sum(sliced_diff**2) / 2


def _solve_alpha_t_beta_t_grid_search(
    *,
    rho_norm: array_typing.FloatVector,
    te_data: array_typing.FloatVector,
    beta_scan_parameters: tuple[float, float, int],
) -> tuple[array_typing.FloatScalar, array_typing.FloatScalar]:
  """Returns the alpha and beta parameters that minimize the temperature loss function.

  Grid search is used for computational efficiency.

  Args:
    rho_norm: Normalized toroidal flux coordinate.
    te_data: The temperature data to be fit, assumed to be on the face grid.
    beta_scan_parameters: A tuple of (beta_min, beta_max, beta_grid_size)
      parameters for the grid search.

  Returns:
    The alpha and beta parameters that minimize the temperature loss function.
  """
  beta_t_trials = jnp.linspace(
      beta_scan_parameters[0],
      beta_scan_parameters[1],
      beta_scan_parameters[2],
  )

  losses = jax.vmap(_loss_for_beta_t, in_axes=(0, None, None))(
      beta_t_trials,
      rho_norm,
      te_data,
  )
  min_index = jnp.argmin(jnp.array(losses, dtype=jax_utils.get_dtype()))
  best_beta_t = beta_t_trials[min_index]
  best_alpha_t = _alpha_closed_form(
      beta=best_beta_t,
      rho_norm=rho_norm,
      profile_data=te_data,
      profile_edge_value=te_data[-1],
  )
  return best_alpha_t, best_beta_t


def cyclotron_radiation_albajar(
    runtime_params: runtime_params_slice.RuntimeParams,
    geo: geometry.Geometry,
    source_name: str,
    core_profiles: state.CoreProfiles,
    unused_calculated_source_profiles: source_profiles.SourceProfiles | None,
    unused_conductivity: conductivity_base.Conductivity | None,
) -> tuple[array_typing.FloatVector, ...]:
  """Calculates the cyclotron radiation heat sink contribution to the electron heat equation.

  Total cyclotron radiation is from:
  F. Albajar et al 2001 Nucl. Fusion 41 665
  https://doi.org/10.1088/0029-5515/41/6/301

  Radial profile of the cyclotron radiation is from:
  J.F. Artaud et al 2018 Nucl. Fusion 58 105001
  https://doi.org/10.1088/1741-4326/aad5b1

  The alpha_n, alpha_T and beta_T parameters are calculated from best fits
  of the following electron density and temperature parameterization to the
  actual plasma data.

  n_e = n_e(0)*(1 - rhonorm**2)**alpha_n
  T_e = (T_e(0)-T_e(a))*(1 - rhonorm**beta_T)**alpha_T + T_e(a)

  Where we take a=0.9 in rhonorm space, and perform the best fit between
  0<rhonorm<0.9, to avoid pedestal effects.

  Args:
    runtime_params: A slice of runtime parameters.
    geo: The geometry object.
    source_name: The name of the source.
    core_profiles: The core profiles object.
    unused_calculated_source_profiles: Unused.

  Returns:
    The cyclotron radiation heat sink contribution to the electron heat
    equation.
  """
  source_params = runtime_params.sources[source_name]
  assert isinstance(source_params, RuntimeParams)

  # Notation conventions based on the Albajar and Artaud papers
  # pylint: disable=invalid-name

  n_e20_face = core_profiles.n_e.face_value() / 1e20
  n_e20_cell = core_profiles.n_e.value / 1e20

  # Dimensionless optical thickness parameter, on-axis:
  # Simplified form of omega_pe**2 / (c * omega_ce) where omega_pe is the
  # plasma frequency and omega_ce is the cyclotron frequency.
  p_a_0 = 6.04e3 * geo.a_minor * n_e20_face[0] / geo.B_0

  # Dimensionless correction term for aspect ratio (equation 15 in Albajar)
  G = 0.93 * (1 + 0.85 * jnp.exp(-0.82 * geo.R_major / geo.a_minor))

  # Calculate profile fit parameters
  alpha_n = _alpha_closed_form(
      beta=2.0,
      rho_norm=geo.rho_face_norm,
      profile_data=n_e20_face,
      profile_edge_value=0.0,
  )
  beta_scan_parameters = (
      source_params.beta_min,
      source_params.beta_max,
      source_params.beta_grid_size,
  )
  alpha_t, beta_t = _solve_alpha_t_beta_t_grid_search(
      rho_norm=geo.rho_face_norm,
      te_data=core_profiles.T_e.face_value(),
      beta_scan_parameters=beta_scan_parameters,
  )

  # The "profile factor" (equation 13 in Albajar)
  K = (
      (alpha_n + 3.87 * alpha_t + 1.46) ** -0.79
      * (1.98 + alpha_t) ** 1.36
      * beta_t**2.14
      * (beta_t**1.53 + 1.87 * alpha_t - 0.16) ** -1.33
  )

  # Calculate power loss in [W]
  P_cycl_total = (
      3.84e-2
      * jnp.sqrt(1 - source_params.wall_reflection_coeff)
      * geo.R_major
      * geo.a_minor**1.38
      * geo.elongation_face[-1] ** 0.79
      * geo.B_0**2.62
      * n_e20_face[0] ** 0.38
      * core_profiles.T_e.face_value()[0]
      * (16 + core_profiles.T_e.face_value()[0]) ** 2.61
      * (1 + 0.12 * core_profiles.T_e.face_value()[0] / p_a_0**0.41) ** -1.51
      * K
      * G
  )

  # Calculate the radial profile on the cell grid,
  # according to the Artaud formula (A.45)
  Q_cycl_shape = (
      geo.R_major
      * geo.elongation**0.79
      * (geo.F / geo.R_major) ** 2.62
      * n_e20_cell**0.38
      * core_profiles.T_e.value**3.61
  )

  # Scale the profile shape to match the total integrated power loss
  denom = math_utils.volume_integration(Q_cycl_shape, geo)
  rescaling_factor = P_cycl_total / denom
  Q_cycl = Q_cycl_shape * rescaling_factor

  return (-Q_cycl,)


@dataclasses.dataclass(kw_only=True, frozen=True, eq=True)
class CyclotronRadiationHeatSink(source.Source):
  """Cyclotron radiation heat sink for electron heat equation."""

  SOURCE_NAME: ClassVar[str] = 'cyclotron_radiation'
  model_func: source.SourceProfileFunction = cyclotron_radiation_albajar

  @property
  def source_name(self) -> str:
    return self.SOURCE_NAME

  @property
  def affected_core_profiles(self) -> tuple[source.AffectedCoreProfile, ...]:
    return (source.AffectedCoreProfile.TEMP_EL,)


class CyclotronRadiationHeatSinkConfig(base.SourceModelBase):
  """Cyclotron radiation heat sink for electron heat equation.

  Attributes:
    wall_reflection_coeff: The wall reflection coefficient is a
      machine-dependent dimensionless parameter corresponding to the fraction of
      cyclotron radiation reflected off the wall and back into the plasma where
      it is re-absorbed. The default value is a typical value.
    beta_min: The minimum value of the beta parameter used in the parameterized
      function for the temperature fit.
    beta_max: The maximum value of the beta parameter used in the parameterized
      function for the temperature fit.
    beta_grid_size: The number of points to use in the grid search for the best
      fit of the temperature function.
  """

  model_name: Annotated[
      Literal['albajar_artaud'], torax_pydantic.JAX_STATIC
  ] = 'albajar_artaud'
  mode: Annotated[runtime_params_lib.Mode, torax_pydantic.JAX_STATIC] = (
      runtime_params_lib.Mode.MODEL_BASED
  )
  wall_reflection_coeff: float = 0.9
  beta_min: Annotated[float, torax_pydantic.JAX_STATIC] = 0.5
  beta_max: Annotated[float, torax_pydantic.JAX_STATIC] = 8.0
  beta_grid_size: Annotated[pydantic.PositiveInt, torax_pydantic.JAX_STATIC] = (
      32
  )

  @pydantic.model_validator(mode='after')
  def _check_fields(self) -> typing_extensions.Self:
    if not self.beta_min < self.beta_max:
      raise ValueError('beta_min must be less than beta_max.')
    return self

  @property
  def model_func(self) -> source.SourceProfileFunction:
    return cyclotron_radiation_albajar

  def build_runtime_params(
      self,
      t: chex.Numeric,
  ) -> 'RuntimeParams':
    return RuntimeParams(
        prescribed_values=tuple(
            [v.get_value(t) for v in self.prescribed_values]
        ),
        wall_reflection_coeff=self.wall_reflection_coeff,
        mode=self.mode,
        is_explicit=self.is_explicit,
        beta_min=self.beta_min,
        beta_max=self.beta_max,
        beta_grid_size=self.beta_grid_size,
    )

  def build_source(self) -> CyclotronRadiationHeatSink:
    return CyclotronRadiationHeatSink(model_func=self.model_func)
