"""Calculate PF flux contribution and resistive, internal, and external flux consumed over the ramp-up."""

import numpy as np

from ....algorithm_class import Algorithm
from ....unit_handling import Unitfull, ureg, wraps_ufunc


@Algorithm.register_algorithm(return_keys=["internal_flux"])
def calc_internal_flux(plasma_current: Unitfull, internal_inductance: Unitfull) -> Unitfull:
    """Calculate the flux due to the plasma current and internal inductance of the plasma (assuming a circular cross-section).

    From: A power-balance model for local helicity injection startup in a spherical tokamak :cite:`Barr_2018`
    NOTE: This is (plasma current times) equation 25 from Barr but applied to a plasma with a
    circular cross-section and a non-cirular cross-section.

    Args:
        plasma_current: [A] :term:`glossary link<plasma_current>`
        internal_inductance: [henry] :term:`glossary link<internal_inductance>`

    Returns:
        [weber] :term:`internal_flux`
    """
    internal_flux = plasma_current * internal_inductance

    return internal_flux


@Algorithm.register_algorithm(return_keys=["external_flux"])
def calc_external_flux(plasma_current: Unitfull, external_inductance: Unitfull) -> Unitfull:
    """Calculate the surface flux generated by the plasma current.

    From: A power-balance model for local helicity injection startup in a spherical tokamak :cite:`Barr_2018`

    Args:
        plasma_current: [A] :term:`glossary link<plasma_current>`
        external_inductance: [henry] :term:`glossary link<external_inductance>`

    Returns:
        [weber] :term:`external_flux`
    """
    return plasma_current * external_inductance


@Algorithm.register_algorithm(return_keys=["resistive_flux"])
def calc_resistive_flux(plasma_current: Unitfull, major_radius: Unitfull, ejima_coefficient: Unitfull) -> Unitfull:
    """Calculate the resistive flux.

    Chapter 8: Plasma operation and control: Physics cite:`Gribov_2007`
    NOTE: CE, the Ejima coefficient, is chosen by default to be 0.4, See for example...
    Chapter 8: Plasma operation and control: cite:`Gribov_2007`
    Ohmic flux consumption during initial operation of the NSTX spherical torus :cite:`Menard_2001`

    Args:
        plasma_current: [A] :term:`glossary link<plasma_current>`
        major_radius: [m] :term:`glossary link<major_radius>`
        ejima_coefficient: [~] :term:`glossary link<ejima_coefficient>`

    Returns:
        [weber] :term:`resistive_flux`
    """
    return ejima_coefficient * ureg.mu_0 * plasma_current * major_radius


@Algorithm.register_algorithm(return_keys=["poloidal_field_flux"])
def calc_poloidal_field_flux(
    vertical_field_mutual_inductance: Unitfull, vertical_magnetic_field: Unitfull, major_radius: Unitfull
) -> Unitfull:
    """Calculate the surface flux contribution from the vertical magnetic field required for radial force balance (which arises from the poloidal field coils).

    From: A power-balance model for local helicity injection startup in a spherical tokamak :cite:`Barr_2018`

    Args:
        vertical_field_mutual_inductance: [~] :term:`glossary link<vertical_field_mutual_inductance>`
        vertical_magnetic_field: [tesla] :term:`glossary link<vertical_magnetic_field>`
        major_radius: [m] :term:`glossary link<major_radius>`

    Returns:
        [weber] :term:`poloidal_field_flux`
    """
    return np.pi * major_radius**2 * vertical_field_mutual_inductance * vertical_magnetic_field


@Algorithm.register_algorithm(return_keys=["flux_needed_from_CS_over_rampup"])
def calc_flux_needed_from_solenoid_over_rampup(
    internal_flux: Unitfull, external_flux: Unitfull, resistive_flux: Unitfull, poloidal_field_flux: Unitfull
) -> Unitfull:
    """Calculate the total flux needed from the central solenoid over rampup.

    Sums together the required fluxes, subtracting the contribution from the poloidal field coils.

    Args:
        internal_flux: [weber] :term:`glossary link<internal_flux>`
        external_flux: [weber] :term:`glossary link<external_flux>`
        resistive_flux: [weber] :term:`glossary link<resistive_flux>`
        poloidal_field_flux: [weber] :term:`glossary link<poloidal_field_flux>`

    Returns:
        [weber] :term:`flux_needed_from_CS_over_rampup`
    """
    return internal_flux + external_flux + resistive_flux - poloidal_field_flux


@Algorithm.register_algorithm(return_keys=["max_flattop_duration"])
def calc_max_flattop_duration(
    total_flux_available_from_CS: Unitfull, flux_needed_from_CS_over_rampup: Unitfull, loop_voltage: Unitfull
) -> Unitfull:
    """Calculate the maximum-duration flattop time that can be driven by the central solenoid.

    Args:
        total_flux_available_from_CS: [weber] :term:`glossary link<total_flux_available_from_CS>`
        flux_needed_from_CS_over_rampup: [weber] :term:`glossary link<flux_needed_from_CS_over_rampup>`
        loop_voltage: [volt] :term:`glossary link<loop_voltage>`

    Returns:
        [seconds] :term:`max_flattop_duration`
    """
    max_flux_for_flattop = total_flux_available_from_CS - flux_needed_from_CS_over_rampup

    return max_flux_for_flattop / loop_voltage


@Algorithm.register_algorithm(return_keys=["breakdown_flux_consumption"])
@wraps_ufunc(input_units=dict(major_radius=ureg.m), return_units=dict(breakdown_flux_consumption=ureg.weber))
def calc_breakdown_flux_consumption(major_radius: float) -> float:
    """Calculate the resistive flux required for breakdown.

    Plasma Design Considerations of Near Term Tokamak Fusion Experimental Reactor :cite:`Sugihara`
    NOTE: given the way the ejima_coefficient is empirically derived in :cite:`Gribov_2007` (i.e., with an implicit assumption that the ramp is defined from Ip=0)
    there is reason to believe a separate calculation for flux consumed over breakdown is not necessary, but it is included here anyways.

    Args:
        major_radius: [m] :term:`glossary link<major_radius>`

    Returns:
        [weber] :term:`breakdown_flux_consumption`
    """
    return float(0.073 * major_radius - 0.00665)
