# -*- coding: utf-8 -*-
"""
Created on Wed Mar  4 14:57:56 2026

@author: e6192
"""

#%% GPU 

import jax
from desc import set_device

print(jax.devices())
jax.print_environment_info()

# set_device("gpu")

# jax.config.update("jax_compilation_cache_dir", "../jax-caches")
# jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
# jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)

#%% Imports
import os
import sys

# Adjust import search path
sys.path.insert(0, os.path.abspath("."))
sys.path.append(os.path.abspath("../../"))

# script_name = os.path.splitext(os.path.basename(__file__))[0] + '/'
# figpath = '/home/sergio/Figures/' + script_name

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from desc.compat import flip_theta
from desc.vmec import VMECIO
from desc.equilibrium import Equilibrium
from desc.geometry import FourierRZToroidalSurface
from desc.integrals import Bounce2D
from desc.grid import LinearGrid
from desc.magnetic_fields import OmnigenousField
from desc.objectives import (
    AspectRatio,
    CurrentDensity,
	ForceBalance,
    FixCurrent,
    FixOmniBmax,
    FixPressure,
    FixPsi,
    GenericObjective,
	get_fixed_boundary_constraints,
    #LinearObjectiveFromUser,
    ObjectiveFunction,
    Omnigenity,
    #RotationalTransform,
)
from desc.optimize import Optimizer
from desc.profiles import PowerSeriesProfile
from desc.plotting import (
    plot_boozer_surface, 
    plot_boundaries,
    plot_comparison,
    plot_section,
    plot_surfaces,
    plot_1d,
    plot_3d,
)
plt.rcParams["font.size"] = 14

#%% TCV Parameters

#TCV parameters
R0 = 0.86
B0 = 3.68
aspect = 2.96
z_shift = -0.296

a = R0/aspect
psi = B0 * np.pi * (a)**2

#%% Initial Guess

surf = FourierRZToroidalSurface(
    R_lmn=[R0, 0.289, 0.279],
    Z_lmn=[-0.289, -0.279],
    modes_R=[[0, 0], [1, 0], [0, 1]],
    modes_Z=[[-1, 0], [0, -1]],
    NFP=2,
)

eq = Equilibrium(M=4, N=4, Psi=psi, surface=surf)

eq = flip_theta(eq)

#%% Solve Equilibrium

eq, _ = eq.solve(
    objective="force", 
    ftol=1e-8,
    xtol=1e-10,
    gtol=1e-10, 
    verbose=3)

eq0 = eq.copy()

#%% Define Omnigenous Field

field = OmnigenousField(
    L_B=1,  # radial resolution of B_lm parameters
    M_B=3,  # number of spline knots on each flux surface
    L_x=1,  # radial resolution of x_lmn parameters
    M_x=1,  # eta resolution of x_lmn parameters
    N_x=1,  # alpha resolution of x_lmn parameters
    NFP=eq.NFP,  # number of field periods; should always be equal to Equilibrium.NFP
    helicity=(1, 0),  # helicity for toroidally closed |B| contours
)

#%%% Objective Function

eq_half_grid = LinearGrid(rho=np.sqrt(0.5), M=4 * eq.M, N=4 * eq.N, NFP=eq.NFP, sym=False)
eq_lcfs_grid = LinearGrid(rho=1.0, M=4 * eq.M, N=4 * eq.N, NFP=eq.NFP, sym=False)

field_half_grid = LinearGrid(rho=np.sqrt(0.5), theta=16, zeta=8, NFP=field.NFP, sym=False)
field_lcfs_grid = LinearGrid(rho=1.0, theta=16, zeta=8, NFP=field.NFP, sym=False)

objective = ObjectiveFunction(
    (
        GenericObjective("R0", thing=eq, target=R0, name="major radius"),
        AspectRatio(eq=eq, target=aspect),
        # omnigenity on the rho=0.5 surface
        Omnigenity(
            eq=eq,
            field=field,
            eq_grid=eq_half_grid,
            field_grid=field_half_grid,
            eta_weight=1,
        ),
        # omnigenity on the rho=1.0 surface
        Omnigenity(
            eq=eq,
            field=field,
            eq_grid=eq_lcfs_grid,
            field_grid=field_lcfs_grid,
            eta_weight=2,
        ),
    )
)

#%%% Constraints

constraints = (
    CurrentDensity(eq=eq),  # vacuum equilibrium force balance
    FixPressure(eq=eq),  # fix vacuum pressure profile
    FixCurrent(eq=eq),  # fix vacuum current profile
    FixPsi(eq=eq),  # fix total toroidal magnetic flux
    # ensure the B_max contour is straight in Boozer coordinates
    FixOmniBmax(field=field),
)

#%%% Optimizer

optimizer = Optimizer("lsq-auglag")
#optimizer = Optimizer("lsq-exact")
(eq, field), _ = optimizer.optimize(
    (eq, field), 
    objective, 
    constraints,
    ftol=1e-6,
    xtol=1e-8,
    gtol=1e-8, 
    maxiter=100, 
    verbose=3,
    options={
        "initial_trust_radius": 0.5,
    },
)

#%%% Re-Solve Fixed-Boundary Equilibrium 

eq, _ = eq.solve(
    objective="force",
    ftol=1e-10,
    xtol=1e-12,
    gtol=1e-12,
    maxiter=50,
    verbose=3)