import numpy as np
import matplotlib.pyplot as plt
from .sortedDict import SortedDict
from .misc import trig2real

__all__ = ["BOOZ_XFORM"]


class BOOZ_XFORM(SortedDict):
    def __init__(self, filename, **kwargs):
        import xarray

        SortedDict.__init__(self)
        self.output = xarray.open_dataset(filename)
        self.ns = int(self.output["ns_b"].values)
        self.nfp = int(self.output["nfp_b"].values)
        self.xm = self.output["ixm_b"].values
        self.xn = self.output["ixn_b"].values / self.nfp
        self.jlist = self.output["jlist"].values
        self.bmnc = self.output["bmnc_b"].values
        return

    @staticmethod
    def write_input(extension, mbooz, nbooz, surfaces):
        """Write a BOO_XFORM input file

        Args:
            extension (str): VMEC output extension.
            mbooz (int): Maximum poloidal mode.
            nbooz (int): Maximum toroidal mode
            surfaces (list): List of flux surfaces.
        """
        template = """{mbooz:d}  {nbooz:d}
'{extension:s}'
{surfaces:s}
        """
        assert mbooz > 0, "mbooz should be >0."
        assert nbooz > 0, "nbooz should be >0."
        surfaces = ["{:} ".format(i) for i in surfaces]
        surfaces = "".join(surfaces)
        with open("inbooz." + extension, "w") as f:
            f.write(
                template.format(
                    mbooz=mbooz, nbooz=nbooz, extension=extension, surfaces=surfaces
                )
            )
        return

    @staticmethod
    def from_vmec(wout, mbooz=0, nbooz=0, surfaces=[]):
        """Prepare BOOZ_XFORM input file from VMEC wout file

        Args:
            wout (str): VMEC wout filename.
            mbooz (int, optional): Maximum poloildal mode number. Defaults to 0 (4*Mpol).
            nbooz (int, optional): Maximum toroidal mode number. Defaults to 0 (4*Ntor).
            surfaces (list, optional): Flux surfaces list. Defaults to [] (1:NS).

        """
        import xarray

        vmec = xarray.open_dataset(wout)
        ind = wout.index("wout_") + 5
        end = wout.index(".nc")
        extension = wout[ind:end]
        if mbooz == 0:
            mbooz = 4 * int(vmec["mpol"].values) + 1
        if nbooz == 0:
            nbooz = 4 * int(vmec["ntor"].values)
        if len(surfaces) == 0:
            ns = int(vmec["ns"].values)
            surfaces = 1 + np.arange(ns)
        vmec.close()
        return BOOZ_XFORM.write_input(extension, mbooz, nbooz, surfaces)

    def plot(
        self,
        ordering=0,
        mn=(None, None),
        ax=None,
        log=True,
        normalize=True,
        logical_not=False,
        **kwargs
    ):
        """Plot |B| components 1D profile

        Args:
            ordering (integer, optional): Plot the leading Nordering asymmetric modes. Defaults to 0.
            mn (tuple, optional): Plot the particular (m,n) mode. Defaults to (None, None).
            ax (Matplotlib axis, optional): Matplotlib axis to be plotted on. Defaults to None.
            log (bool, optional): Plot in log scale. Defaults to True.
            normalize(bool, optionl): Normalized to B_00. Defaults to True.
            logical_not (bool, optional): Unselect mn modes. Defaults to False.
            kwargs (dict): Keyword arguments for matplotlib.pyplot.plot. Defaults to {}.

        Returns:
            ax (Matplotlib axis): Matplotlib axis plotted on.
        """
        xx = self.jlist / self.ns
        return self.plot_helicity(
            self.bmnc,
            self.xm,
            self.xn,
            xx,
            ordering,
            mn,
            ax,
            log,
            normalize,
            logical_not,
            **kwargs
        )

    @staticmethod
    def plot_helicity(
        vals,
        xm,
        xn,
        xx,
        ordering=0,
        mn=(None, None),
        ax=None,
        log=False,
        normalize=False,
        logical_not=False,
        **kwargs
    ):
        # get figure and ax data
        if ax is None:
            fig, ax = plt.subplots()
        plt.sca(ax)
        # check if normalizing to B_00
        if normalize:
            try:
                vals /= vals[:, np.logical_and(xm == 0, xn == 0)]
            except ValueError:
                print("Something wrong with the normalization to B_00")
        # determine filter condition
        if mn[0] is not None:
            mfilter = xm == mn[0]
            m = "m={:}".format(mn[0])
        else:
            mfilter = np.full(np.shape(xm), True)
            m = "m"
        if mn[1] is not None:
            nfilter = xn == mn[1]
            n = "n={:}".format(mn[1])
        else:
            nfilter = np.full(np.shape(xn), True)
            n = "n"
        cond = np.logical_and(mfilter, nfilter)
        if logical_not:
            cond = np.logical_not(cond)
        # data = np.reshape(vals[:, cond], (ns, -1))
        vals = vals[:, cond]
        xm = xm[cond]
        xn = xn[cond]
        # select the top ordering asymmetric terms
        if ordering:
            assert ordering >= 1
            data = np.linalg.norm(vals, axis=0)
            ind_arg = np.argsort(data)
            for i in range(ordering):
                ind = ind_arg[-1 - i]  # index of the i-th largest term
                m = xm[ind]
                n = xn[ind]
                if m == 0 and n == 0:
                    continue
                kwargs["label"] = "m={:}, n={:}".format(int(m), int(n))
                if log:
                    ax.semilogy(xx, np.abs(vals[:, ind]), **kwargs)
                else:
                    ax.plot(xx, np.abs(vals[:, ind]), **kwargs)
            ylabel = r"$\frac{B_{m,n}}{ B_{00} }$"
        else:
            if log:
                line = ax.semilogy(xx, np.linalg.norm(vals, axis=1), **kwargs)
            else:
                line = ax.plot(xx, np.linalg.norm(vals, axis=1), **kwargs)
            ylabel = r"$ \frac{{ \Vert B_{{ {:},{:} }} \Vert }}{{ B_{{00}} }} $".format(
                m, n
            )
        plt.xlabel("normalized flux (s)", fontsize=16)
        plt.ylabel(ylabel, fontsize=16)
        plt.xticks(fontsize=15)
        plt.yticks(fontsize=15)
        return ax

    def plot2d(self, ns=-1, npol=128, ntor=128, ax=None, contour=True, **kwargs):
        """Plot |B| contour on the flux surface

        Args:
            ns (int, optional): Flux surface index in self.jlist. Defaults to -1.
            npol (int, optional): Poloidal resolution. Defaults to 128.
            ntor (int, optional): Toroidal resolution. Defaults to 128.
            ax (Matplotlib axis, optional): Matplotlib axis to be plotted on. Defaults to None.
            contour (bool, optional): Whether plotting the contour lines. Defaults to True.

        Returns:
            ax (Matplotlib axis): Matplotlib axis plotted on.
        """
        _theta = np.linspace(0, np.pi * 2, npol, endpoint=True)
        _zeta = np.linspace(0, np.pi * 2, ntor, endpoint=True)
        modB = trig2real(_theta, _zeta, self.xm, self.xn, self.bmnc[ns, :])
        # get figure and ax data
        if ax is None:
            fig, ax = plt.subplots()
        plt.sca(ax)
        bl = ax.imshow(modB, origin="lower", extent=[0, 2 * np.pi, 0, 2 * np.pi])
        plt.xticks(
            np.linspace(0, 2 * np.pi, 3),
            ["0", r"$\pi/{:d}$".format(self.nfp), r"$2\pi/{:d}$".format(self.nfp)],
            fontsize=15,
        )
        plt.yticks(
            np.linspace(0, 2 * np.pi, 3), ["0", r"$\pi$", r"$2\pi$"], fontsize=15
        )
        if contour:
            ax.contour(_zeta, _theta, modB, colors="grey")
        clb = plt.colorbar(bl)
        clb.ax.tick_params(labelsize=15)
        clb.ax.set_title("|B|(T)", fontsize=15)
        plt.xlabel(r"$\zeta_{Boozer}$", fontsize=16)
        plt.ylabel(r"$\theta_{Boozer}$", fontsize=16)
        import matplotlib.ticker as tck

        ax.xaxis.set_minor_locator(tck.AutoMinorLocator())
        ax.yaxis.set_minor_locator(tck.AutoMinorLocator())
        return ax

    def to_FOCUS(self, ns=-1, focus_file="plasma.boundary", tol=1.0e-8, iota=None):
        """Write a FOCUS plasma boundary file in Boozer coordinates

        Args:
            ns (int, optional): Flux surface label to be exported. Defaults to -1.
            focus_file (str, optional): FOCUS plasma boundary file name. Defaults to 'plasma.boundary'.
            tol (float, optional): Truncated tolerance. Defaults to 1.0E-8.
            iota (float, optional): The target iota to be extracted. Defaults to None.
        """
        mn = int(self.output["mnboz_b"].values)
        # find the closest ns if specifying an iota
        if iota is not None:
            ns = np.argmin(np.abs(self.output["iota_b"].values - iota))
            print(
                "Target iota: ",
                iota,
                ", closest value: ",
                self.output["iota_b"].values[ns],
            )
        rbc = np.array(self.output["rmnc_b"][ns, :])
        # rbs = np.zeros(mn)
        zbs = np.array(self.output["zmns_b"][ns, :])
        # zbc = np.zeros(mn)
        pmns = np.array(self.output["pmns_b"][ns, :])
        # pmnc = np.zeros(mn)
        Nbnf = 0
        # count non-zero terms
        amn = 0
        for imn in range(mn):
            if (abs(rbc[imn]) + abs(zbs[imn] + abs(pmns[imn]))) > tol:
                amn += 1  # number of nonzero coef.
        # write FOCUS plasma boundary file
        with open(focus_file, "w") as fofile:
            fofile.write("# bmn   bNfp   nbf " + "\n")
            fofile.write("{:d} \t {:d} \t {:d} \n".format(amn, self.nfp, Nbnf))
            fofile.write("#plasma boundary" + "\n")
            fofile.write("# n m Rbc Rbs Zbc Zbs Pmnc Pmns" + "\n")
            for imn in range(mn):
                if (abs(rbc[imn]) + abs(zbs[imn] + abs(pmns[imn]))) > tol:
                    fofile.write(
                        "{:4d}  {:4d} \t {:23.15E}  {:12.5E}  {:12.5E} {:23.15E} {:12.5E}  {:23.15E} \n".format(
                            int(self.xn[imn]),
                            self.xm[imn],
                            rbc[imn],
                            0.0,
                            0.0,
                            zbs[imn],
                            0.0,
                            pmns[imn],
                        )
                    )
            fofile.write("#Bn harmonics \n")
            fofile.write("# n m bnc bns" + "\n")
        return
