# ---
# jupyter:
#   jupytext:
#     cell_metadata_filter: tags,-all
#     notebook_metadata_filter: -jupytext.text_representation.jupytext_version
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# %% tags=["remove-cell"]
# SPDX-FileCopyrightText: 2021-present M. Coleman, J. Cook, F. Franza
# SPDX-FileCopyrightText: 2021-present I.A. Maione, S. McIntosh
# SPDX-FileCopyrightText: 2021-present J. Morris, D. Short
#
# SPDX-License-Identifier: LGPL-2.1-or-later

"""
Run PROCESS using the PROCESSTemplateBuilder
"""

# %% [markdown]
# # Running PROCESS from "scratch"
# PROCESS is one of the codes bluemira can use to compliment a reactor design.
# As with any of the external codes bluemira uses, a solver object is created.
# The solver object abstracts away most of the complexities of running different
# programs within bluemira.
#
# This example shows how to build a PROCESS template IN.DAT file

# %%
from bluemira.base.look_and_feel import bluemira_error
from bluemira.codes import systems_code_solver
from bluemira.codes.error import CodesError
from bluemira.codes.process.api import Impurities
from bluemira.codes.process.equation_variable_mapping import Constraint, Objective
from bluemira.codes.process.model_mapping import (
    AlphaJModel,
    AlphaPressureModel,
    AvailabilityModel,
    BetaLimitModel,
    BetaNormMaxModel,
    BootstrapCurrentScalingLaw,
    CSSuperconductorModel,
    ConfinementTimeScalingLaw,
    CostModel,
    CurrentDriveEfficiencyModel,
    DensityLimitModel,
    OperationModel,
    OutputCostsSwitch,
    PFSuperconductorModel,
    PROCESSOptimisationAlgorithm,
    PlasmaCurrentScalingLaw,
    PlasmaGeometryModel,
    PlasmaNullConfigurationModel,
    PlasmaPedestalModel,
    PlasmaProfileModel,
    PowerFlowModel,
    PrimaryPumpingModel,
    SecondaryCycleModel,
    ShieldThermalHeatUse,
    SolenoidSwitchModel,
    TFNuclearHeatingModel,
    TFSuperconductorModel,
)
from bluemira.codes.process.template_builder import PROCESSTemplateBuilder

# %% [markdown]
# First we are going to build a template using the :py:class:`PROCESSTemplateBuilder`,
# without interacting with any of PROCESS' integers.

# %%

template_builder = PROCESSTemplateBuilder()


# %% [markdown]
# Now we're going to specify which optimisation algorithm we want to use, and the
# number of iterations and tolerance.

# %%
template_builder.set_run_title("Example that won't converge")
template_builder.set_optimisation_algorithm(PROCESSOptimisationAlgorithm.VMCON)
template_builder.set_optimisation_numerics(maxiter=200, tolerance=1e-8)


# %% [markdown]
# Let's select the optimisation objective as the major radius:

# %%
template_builder.set_minimisation_objective(Objective.MAJOR_RADIUS)

# %% [markdown]
# You can inspect what options are available by taking a look at the
# :py:class:`Objective` Enum. The options are hopefully self-explanatory.
# The values of the options correspond to the PROCESS integers.

# %%
print("\n".join(str(o) for o in list(Objective)))


# %% [markdown]
# Now we will add a series of constraint equations to the PROCESS problem
# we wish to solve. State True if you wish the constraint to be an equality,
# or False for an inequality.
# You can read more about these constraints an what
# they mean in the PROCESS documentation

# %%
for constraint, equality in (
    (Constraint.BETA_CONSISTENCY, True),
    (Constraint.GLOBAL_POWER_CONSISTENCY, True),
    (Constraint.DENSITY_UPPER_LIMIT, False),
    (Constraint.RADIAL_BUILD_CONSISTENCY, True),
    (Constraint.BURN_TIME_LOWER_LIMIT, False),
    (Constraint.LH_THRESHHOLD_LIMIT, False),
    (Constraint.NET_ELEC_LOWER_LIMIT, False),
    (Constraint.TF_CASE_STRESS_UPPER_LIMIT, False),
    (Constraint.TF_JACKET_STRESS_UPPER_LIMIT, False),
    (Constraint.TF_JCRIT_RATIO_UPPER_LIMIT, False),
    (Constraint.TF_DUMP_VOLTAGE_UPPER_LIMIT, False),
    (Constraint.TF_CURRENT_DENSITY_UPPER_LIMIT, False),
    (Constraint.TF_T_MARGIN_LOWER_LIMIT, False),
    (Constraint.CS_T_MARGIN_LOWER_LIMIT, False),
    (Constraint.CONFINEMENT_RATIO_LOWER_LIMIT, False),
    (Constraint.DUMP_TIME_LOWER_LIMIT, False),
    (Constraint.CS_STRESS_UPPER_LIMIT, False),
    (Constraint.DENSITY_PROFILE_CONSISTENCY, False),
    (Constraint.PSEPB_QAR_UPPER_LIMIT, False),
):
    template_builder.add_constraint(constraint, equality=equality)


# %% [markdown]
# Many of these constraints require certain iteration variables to have been
# specified, or certain input values. The novice user can easily not be
# aware that this is the case, or simply forget to specify these.
#
# The :py:class:`PROCESSTemplateBuilder` will warn the user if certain
# values have not been specified. For example, if we try to make a set of
# inputs for an IN.DAT now, we will get many warning messages:

# %%
inputs = template_builder.make_inputs()


# %% [markdown]
# So let's go ahead and add the iteration variables we want to the problem:

# %%
template_builder.add_variable("b_plasma_toroidal_on_axis", 5.7, upper_bound=20.0)
template_builder.add_variable("rmajor", 8.0, lower_bound=8.0, upper_bound=9.0)
template_builder.add_variable(
    "temp_plasma_electron_vol_avg_kev", 27.33, upper_bound=100.0
)
template_builder.add_variable("beta_total_vol_avg", 3.0e-2)
template_builder.add_variable("nd_plasma_electrons_vol_avg", 7.5e19)
template_builder.add_variable("q95", 3.5, lower_bound=3.0)
template_builder.add_variable("p_hcd_primary_extra_heat_mw", 75.0)
template_builder.add_variable("f_nd_alpha_electron", 6.8940e-02, upper_bound=0.1)
template_builder.add_variable("dr_bore", 2.0, lower_bound=0.1)
template_builder.add_variable("dr_cs", 0.5, lower_bound=0.3)
template_builder.add_variable("dx_tf_turn_steel", 8.0e-3, lower_bound=8.0e-3)
template_builder.add_variable("dr_tf_nose_case", 0.5)
template_builder.add_variable("c_tf_turn", 6.5e4, lower_bound=6.50e4, upper_bound=9.0e4)
template_builder.add_variable("t_tf_superconductor_quench", 2.5e01)
template_builder.add_variable(
    "f_a_tf_turn_cable_copper", 0.80, lower_bound=0.5, upper_bound=0.94
)
template_builder.add_variable("f_c_plasma_non_inductive", 0.8)
template_builder.add_variable("dr_tf_wp_with_insulation", 0.5, lower_bound=0.4)
template_builder.add_variable("f_a_cs_turn_steel", 0.8)

# %% [markdown]
# Often one wants to specify certain impurity concentrations, and even use
# one of these as an iteration variable.

# %%
template_builder.add_impurity(Impurities.H, 1.0)
template_builder.add_impurity(Impurities.He, 0.1)
template_builder.add_impurity(Impurities.W, 5.0e-5)
template_builder.add_variable(Impurities.Xe.id(), 3.8e-04)


# %% [markdown]
# We also want to specify some input values that are not variables:

# %%
template_builder.add_input_values({
    # Profile parameterisation inputs
    "alphan": 1.0,
    "alphat": 1.45,
    "radius_plasma_pedestal_density_norm": 0.94,
    "radius_plasma_pedestal_temp_norm": 0.94,
    "tbeta": 2.0,
    "temp_plasma_pedestal_kev": 5.5,
    "temp_plasma_separatrix_kev": 0.1,
    "nd_plasma_pedestal_electron": 0.5e20,
    "nd_plasma_separatrix_electron": 0.2e20,
    "beta_norm_max": 3.0,
    # Plasma impurity stuff
    "radius_plasma_core_norm": 0.75,
    "f_p_plasma_core_rad_reduction": 0.6,
    # Important stuff
    "p_plant_electric_net_required_mw": 800.0,
    "t_burn_min": 7.2e3,
    "sig_tf_case_max": 7.5e8,
    "sig_tf_wp_max": 7.5e8,
    "alstroh": 7.5e8,
    "psepbqarmax": 10.0,
    "aspect": 3.0,
    "triang": 0.5,
    "q0": 1.0,
    "f_sync_reflect": 0.6,
    "plasma_res_factor": 0.7,
    "ejima_coeff": 0.3,
    "hfact": 1.1,
    # Radial build inputs
    "dr_vv_inboard": 0.3,
    "dr_shld_inboard": 0.3,
    "dr_shld_blkt_gap": 0.02,
    "dr_blkt_inboard": 0.7,
    "dr_fw_plasma_gap_inboard": 0.25,
    "dr_fw_plasma_gap_outboard": 0.25,
    "dr_blkt_outboard": 1.0,
    "dr_vv_outboard": 0.3,
    "dr_shld_outboard": 0.8,
    "dr_cryostat": 0.15,
    # Vertical build inputs
    "dz_vv_upper": 0.3,
    "dz_divertor": 0.62,
    "dz_vv_lower": 0.3,
    # HCD inputs
    "p_hcd_injected_max": 200.0,
    "eta_cd_norm_ecrh": 0.3,
    "eta_ecrh_injector_wall_plug": 0.5,
    "f_c_plasma_bootstrap_max": 0.95,
    # BOP inputs
    "eta_turbine": 0.4,
    "eta_coolant_pump_electric": 0.87,
    "etaiso": 0.9,
    "vfshld": 0.6,
    "t_plant_pulse_dwell": 1800.0,
    "t_plant_pulse_coil_precharge": 500.0,
    # CS / PF coil inputs
    "fcuohsu": 0.7,
    "f_z_cs_tf_internal": 0.9,
    "rpf2": -1.825,
    "c_pf_coil_turn_peak_input": [
        4.0e4,
        4.0e4,
        4.0e4,
        4.0e4,
        4.0e4,
        4.0e4,
        4.0e4,
        4.0e4,
    ],
    "i_pf_location": [2, 2, 3, 3],
    "n_pf_coils_in_group": [1, 1, 2, 2],
    "n_pf_coil_groups": 4,
    "j_pf_coil_wp_peak": [1.1e7, 1.1e7, 6.0e6, 6.0e6, 8.0e6, 8.0e6, 8.0e6, 8.0e6],
    # TF coil inputs
    "n_tf_coils": 16,
    "dr_tf_plasma_case": 0.06,
    "dx_tf_side_case_min": 0.05,
    "ripple_b_tf_plasma_edge_max": 0.6,
    "dia_tf_turn_coolant_channel": 0.01,
    "tftmp": 4.75,
    "dx_tf_wp_insulation": 0.008,
    "tmargmin": 1.5,
    "f_a_tf_turn_cable_space_extra_void": 0.3,
})

# %% [markdown]
# PROCESS has many different models with integer-value 'switches'. We can specify
# these choices as follows:

# %%
for model_choice in (
    BootstrapCurrentScalingLaw.SAUTER,
    ConfinementTimeScalingLaw.IPB98_Y2_H_MODE,
    PlasmaCurrentScalingLaw.ITER_REVISED,
    BetaNormMaxModel.WESSON,
    PlasmaProfileModel.WESSON,
    AlphaJModel.WESSON,
    PlasmaPedestalModel.PEDESTAL_GW,
    PlasmaNullConfigurationModel.SINGLE_NULL,
    BetaLimitModel.THERMAL,
    DensityLimitModel.GREENWALD,
    AlphaPressureModel.WARD,
    PlasmaGeometryModel.CREATE_A_M_S,
    PowerFlowModel.SIMPLE,
    ShieldThermalHeatUse.LOW_GRADE_HEAT,
    SecondaryCycleModel.INPUT,
    CurrentDriveEfficiencyModel.ECRH_UI_GAM,
    OperationModel.PULSED,
    PFSuperconductorModel.NBTI,
    SolenoidSwitchModel.SOLENOID,
    CSSuperconductorModel.NB3SN_WST,
    TFSuperconductorModel.NB3SN_WST,
    PrimaryPumpingModel.PRESSURE_DROP_INPUT,
    TFNuclearHeatingModel.INPUT,
    CostModel.TETRA_1990,
    AvailabilityModel.INPUT,
    OutputCostsSwitch.NO,
):
    template_builder.set_model(model_choice)

# %% [markdown]
# Some of these model choices also require certain input values
# to be specified. If these are not specified by the user, default
# values are used, which may not be desirable. Let us see what
# we're still missing:

# %%
inputs = template_builder.make_inputs()

# %% [markdown]
# And now let's add those missing inputs:

# %%
template_builder.add_input_value("qnuc", 1.3e4)
template_builder.add_input_values({
    "nflutfmax": 0,
    "rrr_tf_cu": 30.0,
    "t_tf_quench_detection": 0.0,
})


# %% [markdown]
# Finally, let us run PROCESS with our inputs. In this case, we're just running
# PROCESS as an external code (see e.g.
# [External code example](../codes/external_code.ex.py))
# So we are not interesed in passing any parameters into it. In future, once the
# input template has been refined to something desirable, one can pass in parameters
# in mapped names to PROCESS, and not need to explicitly know all the PROCESS
# parameter names.

# %%
solver = systems_code_solver(
    params={}, build_config={"template_in_dat": template_builder.make_inputs()}
)

try:
    result = solver.execute("run")
except CodesError as ce:
    bluemira_error(ce)

# %%
# PROCESS runs but no feasbile solution can be found.
# We can adjust our design problem, perhaps relaxing some of the
# requirements we have put on the fusion powerplant.

# %%
# Let's reduce the requirement for net electric power to a lower value.
template_builder.set_run_title("Example that should converge")
template_builder.add_input_value("p_plant_electric_net_required_mw", 400.0)

# %%

solver = systems_code_solver(
    params={}, build_config={"template_in_dat": template_builder.make_inputs()}
)

result = solver.execute("run")
# %%
# Now PROCESS has found a feasible solution!