try:
    # framework is running
    from .startup_choice import *
except ImportError as _excp:
    # class is imported by itself
    if (
        'attempted relative import with no known parent package' in str(_excp)
        or 'No module named \'omfit_classes\'' in str(_excp)
        or "No module named '__main__.startup_choice'" in str(_excp)
    ):
        from startup_choice import *
    else:
        raise


from omfit_classes.omfit_hdf5 import OMFIThdf5
from omfit_classes.utils_fusion import available_EFITs
import numpy as np

__all__ = ['OMFITefitCDF', 'available_EFITs']


class OMFITefitCDF(OMFIThdf5):
    r"""
    class used to interface CDF files generated by EFIT++

    :param filename: filename passed to OMFITascii class

    :param \**kw: keyword dictionary passed to OMFITascii class
    """

    @property
    def times(self):
        return self['time']['data']

    def plot(self, time=None):
        """
        Function used to plot input constraints and outputs

        This method is called by .plot() when the object is a CDF file

        :param time: time in seconds

        :return: None
        """

        if time is None:
            time = self.times[0]

        fig = pyplot.gcf()
        fig.set_size_inches(11, 8)

        self.ax = pyplot.subplot(221)
        self.plot_trace('boundaries', title='Boundary', yname='[m]', xname='Time [s]', tgtName='rCoords', sigName='rSigmas')

        self.ax = pyplot.subplot(222)
        self.plot_trace('plasmaCurrent', title='Current', yname='[A]', xname='Time [s]', sigName='sigma')

        self.ax = pyplot.subplot(223)
        self.plot_profile('pressures', time, title='Pressure', yname='[Pa]', rName="rCoords")

        self.ax = pyplot.subplot(224)
        self.plot_profile('mse', time, title='MSE', yname='[Radians]', ylim_min=-5, ylim_max=5)

    def plot_trace(
        self,
        grpName,
        yname=None,
        xname='Radius [m]',
        tgtName="target",
        cmpName="computed",
        sigName="sigmas",
        title=None,
        ylim_min=None,
        ylim_max=None,
    ):
        """

        :param grpName: EFIT++ constraint dictionary name e.g. [input][constraints][grpName]

        :param yname: y-axis figure title

        :param xname: x-axis figure title

        :param tgtName: EFIT++ target profile dictionary name e.g. [input][constraints][grpName][tgtName]

        :param cmpName: EFIT++ computed profile dictionary name e.g. [input][constraints][grpName][cmpName]

        :param sigName: EFIT++ sigma profile dictionary name e.g. [input][constraints][grpName][sigName]

        :param title: figure title

        :param ylim_min: y-axis figure minimum value

        :param ylim_max: y-axis figure maximum value

        :return:
        """

        if grpName not in self['input']['constraints']:
            print(grpName + " constraint does not exist")
            self.ax.set_visible(False)
            return

        # target profiles
        target = np.squeeze(self['input']['constraints'][grpName][tgtName]['data'])
        sigmas = np.squeeze(self['input']['constraints'][grpName][sigName]['data'])

        # computed profiles
        if grpName == 'boundaries':
            computed = np.squeeze(self['output']['separatrixGeometry']['rmidplaneOut']['data'])
        else:
            computed = np.squeeze(self['input']['constraints'][grpName][cmpName]['data'])

        self.ax.ticklabel_format(style='sci', scilimits=(0, 0), axis='y')
        self.ax.errorbar(self.times, target, marker='o', markersize=5, yerr=sigmas, label='Target', color='b')
        self.ax.plot(self.times, computed, marker='x', markersize=5, label='Computed', linewidth=2, color='r')
        self.ax.set_ylim(ylim_min, ylim_max)
        self.ax.legend(title=title, prop={'size': 10}, bbox_to_anchor=[0.98, 0.98], fancybox=True)
        self.ax.set_ylabel(yname)
        self.ax.set_xlabel(xname)
        self.ax.grid(True, linestyle='-', color='0.75', linewidth=0.2)

    def plot_profile(
        self,
        grpName,
        time=None,
        yname=None,
        xname='Radius [m]',
        tgtName="target",
        cmpName="computed",
        sigName="sigmas",
        rName="r",
        title=None,
        ylim_min=None,
        ylim_max=None,
    ):
        '''

        :param grpName: EFIT++ constraint dictionary name e.g. [input][constraints][grpName]

        :param time: single time slice in seconds to plot profile data

        :param yname: y-axis figure title

        :param xname: x-axis figure title

        :param tgtName: EFIT++ target profile dictionary name e.g. [input][constraints][grpName][tgtName]

        :param cmpName: EFIT++ computed profile dictionary name e.g. [input][constraints][grpName][cmpName]

        :param sigName: EFIT++ sigma profile dictionary name e.g. [input][constraints][grpName][sigName]

        :param rName: EFIT++ radius profile dictionary name e.g. [input][constraints][grpName][rName]

        :param title: figure title

        :param ylim_min: y-axis figure minimum value

        :param ylim_max: y-axis figure maximum value

        :return:
        '''

        if time is None:
            time = self.times[0]

        if grpName not in self['input']['constraints']:
            print(grpName + " constraint does not exist")
            self.ax.set_visible(False)
            return

        # profile time slice
        tind = np.abs(time - self.times).argmin()

        # target profiles
        target = self['input']['constraints'][grpName][tgtName]['data']
        sigmas = self['input']['constraints'][grpName][sigName]['data']

        # computed profiles
        computed = self['input']['constraints'][grpName][cmpName]['data']

        # radius
        radius = self['input']['constraints'][grpName][rName]['data']

        # separatrix
        rSepOut = self['output']['separatrixGeometry']['rmidplaneOut']['data']

        if len(np.shape(radius)) > 1:
            radius = radius[tind, :]

        self.ax.ticklabel_format(style='sci', scilimits=(0, 0), axis='y')
        self.ax.errorbar(radius, target[tind, :], marker='o', markersize=5, yerr=sigmas[tind, :], label='Target', color='b')
        self.ax.plot(radius, computed[tind, :], marker='x', markersize=5, label='Computed', linewidth=2, color='r')
        self.ax.set_ylim(ylim_min, ylim_max)
        self.ax.plot([rSepOut[tind], rSepOut[tind]], self.ax.get_ylim(), linestyle='--', color='r', label='Separatrix')
        self.ax.legend(title=title + '; t= ' + str(time) + ' s', prop={'size': 10}, bbox_to_anchor=[0.98, 0.98], fancybox=True)
        self.ax.set_ylabel(yname)
        self.ax.set_xlabel(xname)
        self.ax.grid(True, linestyle='-', color='0.75', linewidth=0.2)


############################################
if '__main__' == __name__:
    test_classes_main_header()

    tmp = OMFITefitCDF(OMFITsrc + '/../samples/efitOut.nc')
    print(tmp)
