# modified from stellopt.pySTEL

import numpy as np
import xarray
import matplotlib.pyplot as plt
from .misc import trig2real
from .surface import FourSurf

__all__ = ["VMECout"]


class VMECout(object):
    """
    Class to parse a VMEC wout file

    :param filename: filename passed to xarray.open_dataset()

    The entire dataset is stored in `self.wout` and you can access
    to variables via `self.wout['iotaf'].values`.
    The flux surfaces are all parsed as `FourSurf` classes and stored
    in the list of `self.surface`.
    Magnetic fields are Fourier transformed to real space
    (only the stellarator symmetric part, `bmnc`) and store in `self.data['b]`.

    `self.plot` has several options to plot the profiles, etc.

    """

    def __init__(self, filename, **kwargs):
        self.wout = xarray.open_dataset(filename)
        self.data = {}
        self.data["ns"] = int(self.wout["ns"].values)
        self.data["nfp"] = int(self.wout["nfp"].values)
        self.data["nu"] = int(self.wout["mpol"].values * 4)
        self.data["nv"] = int(self.wout["ntor"].values * 4)
        self.data["nflux"] = np.linspace(
            0, 1, self.data["ns"]
        )  # np.ndarray((self.data['ns'],1))
        self.data["theta"] = np.linspace(
            0, 2 * np.pi, self.data["nu"]
        )  # np.ndarray((self.data['nu'],1))
        self.data["zeta"] = np.linspace(
            0, 2 * np.pi, self.data["nv"]
        )  # np.ndarray((self.data['nv'],1))
        self.surface = []
        self.data["b"] = []
        for i in range(self.data["ns"]):
            self.surface.append(
                FourSurf(
                    xm=self.wout["xm"].values,
                    xn=self.wout["xn"].values,
                    rbc=self.wout["rmnc"][i].values,
                    rbs=np.zeros_like(self.wout["rmnc"][i].values),
                    zbs=self.wout["zmns"][i].values,
                    zbc=np.zeros_like(self.wout["zmns"][i].values),
                )
            )
            self.data["b"].append(
                trig2real(
                    self.data["theta"],
                    self.data["zeta"],
                    self.wout["xm_nyq"].values,
                    self.wout["xn_nyq"].values / self.data["nfp"],
                    self.wout["bmnc"][i].values,
                )
            )
        return

    def plot(self, plot_name="none", ax=None, **kwargs):
        """Plot various VMEC quantities

        Args:
            plot_name (str, optional): The quantity to be plotted, should be one of
                                       iota, q, pressue, <Buco>, <Bvco>, <jcuru>, <jcurv>,
                                       <j.B>, LPK, none. Defaults to 'none'.
            ax (Matplotlib axis, optional): The Matplotlib axis to be plotted on. Defaults to None.
        """
        if ax is None:
            fig, ax = plt.subplots()
        else:
            fig = ax.get_figure()
        if plot_name == "none":
            print("You can plot: iota, q, pressue, <Buco>, <Bvco>, <jcuru>, <jcurv>, ")
            print("               <j.B>, LPK")
        elif plot_name == "iota":
            ax.plot(self.data["nflux"], self.wout["iotaf"].values, **kwargs)
            ax.set_xlabel("Normalized Flux")
            ax.set_ylabel("iota")
            ax.set_title("Rotational Transform")
            # ax.set(xlabel='s',ylabel='iota',aspect='square')
        elif plot_name == "q":
            ax.plot(self.data["nflux"], 1.0 / self.wout["iotaf"].values, **kwargs)
            ax.set_xlabel("Normalized Flux")
            ax.set_ylabel("q")
            ax.set_title("Safety Factor")
        elif plot_name == "pressure":
            ax.plot(self.data["nflux"], self.wout["presf"].values / 1000, **kwargs)
            ax.set_xlabel("Normalized Flux")
            ax.set_ylabel("Pressure [kPa]")
            ax.set_title("Pressure Profile")
        elif plot_name == "<Buco>":
            ax.plot(self.data["nflux"], self.wout["buco"].values, **kwargs)
            ax.set_xlabel("Normalized Flux")
            ax.set_ylabel("<B^u> [T]")
            ax.set_title("Flux surface Averaged B^u")
        elif plot_name == "<Bvco>":
            ax.plot(self.data["nflux"], self.wout["bvco"].values, **kwargs)
            ax.set_xlabel("Normalized Flux")
            ax.set_ylabel("<B^v> [T]")
            ax.set_title("Flux surface Averaged B^v")
        elif plot_name == "<jcuru>":
            ax.plot(self.data["nflux"], self.wout["jcuru"].values / 1000, **kwargs)
            ax.set_xlabel("Normalized Flux")
            ax.set_ylabel("<j^u> [kA/m^2]")
            ax.set_title("Flux surface Averaged j^u")
        elif plot_name == "<jcurv>":
            ax.plot(self.data["nflux"], self.wout["jcurv"].values / 1000, **kwargs)
            ax.set_xlabel("Normalized Flux")
            ax.set_ylabel("<j^v> [kA/m^2]")
            ax.set_title("Flux surface Averaged j^v")
        elif plot_name == "<j.B>":
            ax.plot(self.data["nflux"], self.wout["jdotb"].values / 1000, **kwargs)
            ax.set_xlabel("Normalized Flux")
            ax.set_ylabel("<j.B> [T*kA/m^2]")
            ax.set_title("Flux surface Averaged j.B")
        elif plot_name == "LPK":
            self.surface[-1].plot(zeta=0, color="red", label=r"$\phi=0$", **kwargs)
            self.surface[-1].plot(
                zeta=0.5 * np.pi / self.data["nfp"], color="green", label=r"$\phi=0.25$", **kwargs
            )
            self.surface[-1].plot(
                zeta=np.pi / self.data["nfp"], color="blue", label=r"$\phi=0.5$", **kwargs
            )
            ax.set_title("LPK Plot")
        elif plot_name[0] == "-":
            print(plot_name)
        else:
            return ax
