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_ascii import *
from omfit_classes.omfit_gato import OMFITdskgato
import numpy as np

__all__ = ['OMFITtoqProfiles', 'OMFITdskeqdata']


class OMFITtoqProfiles(SortedDict, OMFITascii):
    """TOQ profiles data files"""

    def __init__(self, filename, **kw):
        OMFITascii.__init__(self, filename, **kw)
        SortedDict.__init__(self)
        self.dynaLoad = True

    @dynaLoad
    def load(self):
        with open(self.filename, 'r') as f:
            lines = f.readlines()
        lines = lines[2:]
        for k, line in enumerate(lines):
            if re.findall('[a-z]', lines[k]):
                lines[k] = ''
                continue
            lines[k] = lines[k].rstrip('\n')
        lines = [_f for _f in lines if _f]
        #                       0, 1, 2, 3, 4, 5, 6, 7, 9,10,11,12,13,14,15,16,17,18,19,20,21
        for nn, n in enumerate(
            [
                [0, 8, 8, 8, 8, 8, 8, 8, 7, 7, 7, 8, 7, 8, 8, 8, 8, 8, 7, 7, 8],
                [0, 8, 8, 8, 8, 8, 9, 8, 7, 7, 8, 7, 8, 8, 8, 8, 8, 7, 8, 7, 8],
            ]
        ):
            try:
                tmp = [x[0] + x[1] for x in zip(lines[0::2], lines[1::2])]
                for k, line in enumerate(tmp):
                    line = '  ' + line
                    line = [line[np.sum(n[: i + 1]) : np.sum(n[: i + 2])] for i in range(len(n) - 1)]
                    for k1 in range(len(line)):
                        # print(k1,'='+line[k1]+'=')
                        if '*' in line[k1]:
                            line[k1] = 'nan'
                    line = list(map(float, line))
                    tmp[k] = line
                break
            except Exception:
                if nn == 1:
                    raise
                else:
                    pass
        tmp = np.array(tmp).T

        names = ['ped', 'center', 'top', 'rho_0.95', 'rho_0.90']

        self.clear()
        for k1 in range(5):
            self[names[k1]] = SortedDict()
        for k, item in enumerate(
            [
                'PsiN',
                'a-r(cm)',
                'p(MPa)',
                'Te(keV)',
                'n(13cm-3)',
                'pp(MPa/W)',
                'j.B/B0',
                'beta%',
                'betan',
                'nu*',
                'rho',
                'width/a',
                'q',
                's',
                'alpha',
                'al_cyl',
                'eps',
                'betap',
                'rhot_dt',
                'Ti(kev)',
            ]
        ):
            self[item] = tmp[k, 5:]
            for k1 in range(5):
                self[names[k1]][item] = tmp[k, k1]

        for item in self:
            if isinstance(self[item], np.ndarray):
                self[item] = self[item][::-1]

    @dynaLoad
    def plot(self):
        ls = ['o', '.', 's', 'd', 'x']
        for k, item in enumerate([k for k in list(self.keys()) if k != 'PsiN' and isinstance(self[k], np.ndarray)]):
            if k == 0:
                ax = pyplot.subplot(4, 5, k + 1)
            else:
                pyplot.subplot(4, 5, k + 1, sharex=ax)
            pyplot.plot(self['PsiN'], self[item])
            for k1, point_item in enumerate([k2 for k2 in list(self.keys()) if isinstance(self[k2], dict)]):
                pyplot.plot(
                    self[point_item]['PsiN'], self[point_item][item], ls[k1], color=pyplot.gca().lines[-1].get_color(), label=point_item
                )
            title_inside(item, y=0.8)
        autofmt_sharex()
        pyplot.xlim([0, 1])


class OMFITdskeqdata(OMFITdskgato):
    r"""
    OMFIT class used to interface to equilibria files generated by TOQ (dskeqdata files)

    :param filename: filename passed to OMFITobject class

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

    def __init__(self, filename='', **kw):
        OMFITascii.__init__(self, filename, **kw)
        SortedDict.__init__(self)
        self.dynaLoad = True

    @dynaLoad
    def load(self):
        # Parser based on readeqdata.f in TOQ
        with open(self.filename, 'r') as f:
            lines = f.read()
        header, lines = lines.split('&end')
        self['__header__'] = header + '&end'
        lines = re.sub('\n', ' ', lines).split()
        self['NSURF'] = int(lines.pop(0))
        self['NTHT'] = int(lines.pop(0))
        self['PSI'] = np.array([float(lines.pop(0)) for k in range(self['NSURF'])]) / 1e8
        self['PRES'] = np.array([float(lines.pop(0)) for k in range(self['NSURF'])]) / 10.0
        self['PPRIME'] = np.array([float(lines.pop(0)) for k in range(self['NSURF'])]) / 1e-7
        self['FPOL'] = np.array([float(lines.pop(0)) for k in range(self['NSURF'])]) / 1e6
        self['FFPRIM'] = np.array([float(lines.pop(0)) for k in range(self['NSURF'])]) / 1e4
        self['QPSI'] = np.array([float(lines.pop(0)) for k in range(self['NSURF'])])
        self['GROT'] = np.array([float(lines.pop(0)) for k in range(self['NSURF'])])
        self['GROTP'] = np.array([float(lines.pop(0)) for k in range(self['NSURF'])])
        self['RBBBS'] = np.array([float(lines.pop(0)) for k in range(self['NTHT'])]) / 100.0
        self['ZBBBS'] = np.array([float(lines.pop(0)) for k in range(self['NTHT'])]) / 100.0
        # detect symmetric equilibria
        self['SYMMETRIC'] = False
        if self['ZBBBS'][0] == self['ZBBBS'][-1] and self['RBBBS'][0] != self['RBBBS'][-1]:
            self['RBBBS'] = np.hstack((self['RBBBS'], self['RBBBS'][self['NTHT'] - 1 :: -1]))
            self['ZBBBS'] = np.hstack((self['ZBBBS'], -self['ZBBBS'][self['NTHT'] - 1 :: -1]))
            self['SYMMETRIC'] = True
        self['CD'] = np.array([float(lines.pop(0)) for k in range(self['NTHT'])])

        self['BCENTR'] = float(re.findall(r'b vac\(rzero\)\(g\)\s+=\s+[\w\.\+\-]+', self['__header__'])[0].split('=')[-1]) / 4000.0
        self['RCENTR'] = float(re.findall(r'rzero\s+=\s+[\w\.\+\-]+', self['__header__'])[0].split('=')[-1]) / 100.0
        self['RMAXIS'] = float(re.findall(r'r axis \(cm\)\s+=\s+[\w\.\+\-]+', self['__header__'])[0].split('=')[-1]) / 100.0
        self['ZMAXIS'] = 0.0

        self.add_derived()

    @dynaSave
    def save(self):
        lines = [self['__header__']]
        lines.append('%d' % self['NSURF'])
        lines.append('%d' % self['NTHT'])
        for item in [
            ('PSI', 1e8),
            ('PRES', 10.0),
            ('PPRIME', 1e-7),
            ('FPOL', 1e6),
            ('FFPRIM', 1e4),
            'QPSI',
            'GROT',
            'GROTP',
            ('RBBBS', 100.0),
            ('ZBBBS', 100.0),
            'CD',
        ]:
            mul = 1.0
            if isinstance(item, tuple):
                item, mul = item
            if self['SYMMETRIC'] and item in ['RBBBS', 'ZBBBS']:
                data = copy.deepcopy(self[item][: int(len(self[item]) // 2)]).tolist()
            else:
                data = copy.deepcopy(self[item]).tolist()
            while len(data):
                lines.append(' '.join(['%16.16f' % (data.pop(0) * mul) for k in range(min([4, len(data)]))]))
        with open(self.filename, 'w') as f:
            f.write('\n'.join(lines))
