# Source code for spacepy.pybats.rim

    #!/usr/bin/env python
    '''
    Classes, functions, and methods for reading, writing, and plotting output
    from the Ridley Ionosphere Model (RIM) and the similar legacy code,
    Ridley Serial.

    Copyright 2010 Los Alamos National Security, LLC.

    .. rubric:: Classes

    .. autosummary::
        :toctree:
        :template: clean_class.rst

        Iono
        OvalDebugFile

    .. rubric:: Functions

    .. autosummary::
        :toctree:

        get_iono_cb
        tex_label
    '''

    import os
    import numpy as np
    from spacepy import deprecated
    import spacepy.plot.apionly
    from spacepy.plot import set_target
    from spacepy.pybats import PbData, dmarray

    def get_iono_cb(ct_name='bwr'):
        '''
        Several custom colorbars used by RIM and AMIE have become standard when
        visualizing data from these models.  These are 'blue_white_red' and
        'white_red', used for data that have positive and negative values and
        for data that have only positive values, respectively.  This function
        builds and returns these colorbars when called with the initials of the
        color table name as the only argument.

        Other Parameters
        ================
        ct_name : str
           Select the color table.  Can be 'bwr' for blue-white-red or 'wr' for
           white-red.  Defaults to 'bwr'.

        Examples
        ========
        >>> bwr_map = get_iono_cb('bwr')
        >>> wr_map  = get_iono_cb('wr')

        '''

        from matplotlib.colors import LinearSegmentedColormap as lsc

        if ct_name=='bwr':
            table = {
                'red':  [(0.,0.,.0),(.34,0.,0.),(.5,1.,1.),(1.,1.,1.)],
                'green':[(0.,0.,0.),(.35,1.,1.),(.66,1.,1.),(1.,0.,0.)],
                'blue' :[(0.,1.,1.),(.5,1.,1.),(.66,0.,0.),(.85,0.,0.),(1.,.1,.1)]
                }
            cmap = lsc('blue_white_red',table)
        elif ct_name=='wr':
            table = {
                'red':  [(0.,1.,1.),(1.,1.,1.)],
                'green':[(0.,1.,1.),(1.,0.,0.)],
                'blue' :[(0.,1.,1.),(1.,.0,.0)]
                }
            cmap = lsc('white_red',table)

        return cmap

    def tex_label(varname):
        '''
        Many variable names used in the Ridley Ionosphere Model look much better
        in LaTeX format with their proper Greek letters.  This function takes
        a variable name, and if it is recognized, returns a properly formatted
        string that uses MatPlotLib's MathText functionality to display the
        proper characters.  If it is not recognized, the varname is returned.

        Parameters
        ==========
        varname : string
           The variable to convert to a LaTeX label.

        Examples
        ========
        >>>tex_label('n_phi')
        '\\Phi_{Ionosphere}'
        >>>tex_label('Not Recognized')
        'Not Recognized'

        '''

        if varname[:2]=='n_' or varname[:2]=='s_':
            varname=varname[2:]

        known = {
            'phi': r'$\Phi_{Ionosphere}$',
            'sigmah':r'$\sigma_{Hall}$',
            'sigmap':r'$\sigma_{Peder}$',
            'jr':r'$J_{radial}$',
            'mA/m^2':r'$\mu A/m^{2}$',
            'W/m2':r'$W/m^2$',
            'eV':'$eV$'
            }

        if varname in known:
            label = known[varname]
        else:
            label = varname

        return label

    class Iono(PbData):
        '''
        A class for handling 2D output from the Ridley Ionosphere Model.
        Instantiate an object as follows:

        >>> iono = rim.Iono('filename.idl')

        ...where filename.idl is the name of a RIM 2D output file.
        '''

        def __init__(self, infile, *args, **kwargs):
            super(Iono, self).__init__(*args, **kwargs)
            self.attrs['file']=infile
            self.readascii()

        def readascii(self):
            '''
            Read an ascii ".idl" output file and load the contents into the object.
            '''

            from sys import version_info
            import re
            import datetime as dt
            from numpy import zeros, reshape
            import gzip

            # slurp entire file.
            if self.attrs['file'][-3:]=='.gz':
                try:
                    infile = gzip.open(self.attrs['file'],'rt')
                except ValueError: #Python2.7 (Windows) compatibility
                    infile = gzip.open(self.attrs['file'])
            else:
                infile = open(self.attrs['file'], 'r')
            raw = infile.readlines()
            infile.close()

            # Parse header
            title = raw[raw.index('TITLE\n')+1]
            self.attrs['title'] = title[title.index('"')+1:title.rindex('"')]

            i = raw.index('NUMERICAL VALUES\n')
            self.attrs['nvars'] = int(raw[i+1].split()[0])
            self.attrs['ntheta']= int(raw[i+2].split()[0])
            self.attrs['nphi']  = int(raw[i+3].split()[0])

            # Convenience:
            nphi, ntheta = self.attrs['nphi'], self.attrs['ntheta']

            i = raw.index('TIME\n')
            self.attrs['time'] = dt.datetime(
                int(raw[i+1].split()[0]),      #year
                int(raw[i+2].split()[0]),      #month
                int(raw[i+3].split()[0]),      #day
                int(raw[i+4].split()[0]),      #hour
                int(raw[i+5].split()[0]),      #min
                int(raw[i+6].split()[0]),      #sec
                int(raw[i+7].split()[0])*1000  #microsec
                )

            i = raw.index('SIMULATION\n')
            self.attrs['iter']    =   int(raw[i+1].split()[0])
            self.attrs['simtime'] = float(raw[i+2].split()[0])

            i = raw.index('DIPOLE TILT\n')
            self.tilt = zeros(2)
            self.tilt[0] = float(raw[i+1].split()[0])
            self.tilt[1] = float(raw[i+2].split()[0])

            i = raw.index('VARIABLE LIST\n')
            namevar = []
            units   = {}
            for j in range(i+1,i+self.attrs['nvars']+1):
                match = re.match(r'\s*\d+\s+([\w\s\W]+)\[([\w\s\W]+)\]',raw[j])
                if match:
                    name = (match.group(1).strip()).lower()
                    namevar.append(name)
                    units[name] = match.group(2).strip()
                else:
                    raise ValueError('Could not parse %s' % raw[j])

            ### Read all data ###

            # Create data arrays
            nPts = self.attrs['ntheta']*self.attrs['nphi']
            for key in namevar:
                self['n_'+key] = dmarray(zeros(nPts), {'units':units[key]})
                self['s_'+key] = dmarray(zeros(nPts), {'units':units[key]})
            i = raw.index('BEGIN NORTHERN HEMISPHERE\n')+1

            # Some compilers insert line breaks automatically when fortran format
            # string is not adequately specified.  Let's see if that's the
            # case here: how many lines does it take to cover all variables?
            nvars, nvarline, nwrap = len(namevar), 0, 0
            while nvarline<nvars:
                nvarline += len(raw[i+nwrap].split())
                nwrap += 1

            # Fill data arrays:
            for j in range(nPts):
                # Build list of variables; accounting for line-wrapping:
                parts = []
                iLine = i + j*nwrap
                for iwrap in range(nwrap):
                    parts += raw[iLine+iwrap].split()
                for k in range(self.attrs['nvars']):
                    self['n_'+namevar[k]][j] = parts[k]
            i = raw.index('BEGIN SOUTHERN HEMISPHERE\n')+1
            for j in range(nPts):
                parts = []
                iLine = i + j*nwrap
                for iwrap in range(nwrap):
                    parts += raw[iLine+iwrap].split()
                for k in range(self.attrs['nvars']):
                    self['s_'+namevar[k]][j] = parts[k]

            # Create 2-D arrays.
            for key in namevar:
                nkey, skey = 'n_'+key, 's_'+key
                self[nkey] = reshape(self[nkey], (ntheta, nphi), 'F')
                self[skey] = reshape(self[skey], (ntheta, nphi), 'F')

            # Some extra grid info:
            self.dlon = self['n_psi'  ][0,3]-self['n_psi'  ][0,2]
            self.dlat = self['n_theta'][3,0]-self['n_theta'][2,0]

        def calc_I(self):
            '''
            Integrate radial current to get the following values:
            I:     The total radial current over one hemisphere.
            Iup:   The total upward current over one hemisphere.
            Idown: The total downward current over one hemisphere.

            Values are stored in the object with the prefix 'n_' or 's_'
            to indicate the hemisphere and may be accessed via self['n_I'], etc.

            Values are stored in units of mega Amps.

            Parameters
            ==========

            Returns
            =======
            True

            Examples
            ========
            >>> a = rim.Iono('spacepy/tests/data/pybats_test/it000321_104510_000.idl.gz')
            >>> a.calc_I()
            >>> print(a['n_Iup'])

            '''

            # Calculate some physically meaningful values/units
            units = 1E-6*1E-6  # micro amps to amps, amps to MegaAmps
            R = (6371.0+110.0)*1000.0 # Radius of Earth + iono altitude
            dTheta = np.pi*self.dlat/180.
            dPhi   = np.pi*self.dlon/180.

            # -----NORTHERN HEMISPHERE-----
            # Get relevant values:
            colat     = self['n_theta']*np.pi/180.
            integrand = self['n_jr']*np.sin(colat)*dTheta*dPhi
            # Get locations of "up" and "down"
            loc_up = self['n_jr']>0
            loc_do = self['n_jr']<0
            self['n_I']     = units*R**2 * np.sum(integrand)
            self['n_Iup']   = units*R**2 * np.sum(integrand[loc_up])
            self['n_Idown'] = units*R**2 * np.sum(integrand[loc_do])

            # -----SOUTHERN HEMISPHERE-----
            # Get relevant values:
            colat     = self['s_theta']*np.pi/180.
            integrand = self['s_jr']*np.sin(colat)*dTheta*dPhi
            # Get locations of "up" and "down"
            loc_up = self['s_jr']>0
            loc_do = self['s_jr']<0
            self['s_I']     = units*R**2 * np.sum(integrand)
            self['s_Iup']   = units*R**2 * np.sum(integrand[loc_up])
            self['s_Idown'] = units*R**2 * np.sum(integrand[loc_do])

        def add_cont(self, var, target=None, n=24, maxz=False, lines=True,
                     cmap=False, add_cbar=False, label=None, loc=111,
                     xticksize=12, yticksize=12, max_colat=40, **kwargs):
            '''
            Create a polar contour of variable *var*.  Plot will be either drawn
            on a new matplotlib figure and axes, or you can specify a plot target
            using the *target* kwarg.

            Parameters
            ==========
            var : str
               The name of the variable to plot.

            Returns
            =======
            fig : matplotlib figure object
            ax  : matplotlib axes object
            cnt : matplotlib contour object
            cb  : matplotlib colorbar object

            Other Parameters
            ================
            target : Figure or Axes
                If None (default), a new figure is generated from scratch.
                If a matplotlib Figure object, a new axis is created
                to fill that figure.
                If a matplotlib Axes object, the plot is placed
                into that axis.
            loc : int
                Use to specify the subplot placement of the axis
                (e.g. loc=212, etc.) Used if target is a Figure or None.
                Default 111 (single plot).
            n : int
                Set number of levels.  Should be a multiple of 3 for best match
                between filled and traced contours.  Default is 21.
            lines : bool
                Add unfilled black solid/dashed contours to plot for additional
                contrast.  Default is **True**.
            maxz : real
                Set the max/min value for the color bar.  Default is set by data.
            max_colat : real
                Set the co-latitude range of the plot, in degrees.  Defaults to 40.
            cmap : str
                Set the colormap.  Default is to autoselect using classic IE maps.
                Can be 'bwr', 'wr', or any name of a matplotlib colar map.
            add_cbar : bool
                Add colorbar to plot.  Default is **False** which will
                not add one to the plot.
            label : str or None
                Label to place at top of plot.  Defaults to variable name.
            xticksize : int
                Size of longitude markers in text points.  Defaults to 12.
            yticksize : int
                Size of latitude markers in text points.  Defaults to 12.

            Extra keywords are passed to the contourf routine.

            '''
            # Get only what we need to decrease runtime.
            from math import pi
            from numpy import linspace
            from matplotlib.colors import Normalize
            from matplotlib.ticker import MaxNLocator, MultipleLocator
            from matplotlib.pyplot import clabel, colorbar

            fig, ax = set_target(target, polar=True, loc=loc)

            hemi = var[:2]

            # user defined variables may not have hemisphere marking.
            # Assume nothern hemi in those cases.
            if hemi!='n_' and hemi!='s_': hemi = 'n_'

            # Set levels and ticks:
            if label==None:
                label=tex_label(var)
            lt_labels = ['06',    label, '18',   '00']
            xticks    = [   0,   pi/2,   pi, 3*pi/2]
            lct = MultipleLocator(10)
            minz = self[var].min()
            if minz < 0.0:
                if not maxz:
                    maxz = max([abs(minz),self[var].max()])
                crange = Normalize(vmin=-1.*maxz, vmax=maxz)
                levs = linspace(-1.*maxz, maxz, n)
            else:
                if not maxz:
                    maxz = self[var].max()
                crange = Normalize(vmin=0., vmax=maxz)
                levs = linspace(0., maxz, n)

            # Get color map if not given:
            if not cmap:
                if self[var].min() >= 0.0:
                    cmap=get_iono_cb('wr')
                else:
                    cmap=get_iono_cb('bwr')

            # Set the latitude based on hemisphere:
            theta = self[hemi+'theta']
            if 's_' in hemi: theta = 180-self[hemi+'theta']

            # Create contour:
            cnt1 = ax.contourf(self[hemi+'psi']*pi/180.0+pi/2., theta,
                               np.array(self[var]), levs, norm=crange, cmap=cmap,
                               **kwargs)
            # Set xtick label size, increase font of top label.
            labels = ax.get_xticklabels()
            for l in labels: l.set_size(xticksize)
            labels[1].set_size(xticksize*1.25)

            if lines:
                nk = int(round(n/3.0))
                cnt2 = ax.contour(self[hemi+'psi']*pi/180.0+pi/2., theta,
                                  np.array(self[var]), nk, colors='k')
                #clabel(cnt2,fmt='%3i',fontsize=10)

            if add_cbar:
                cbarticks = MaxNLocator(7)
                cbar = colorbar(cnt1, ticks=cbarticks, shrink=0.75, pad=0.08, ax=ax)
                cbar.set_label(tex_label(self[var].attrs['units']))
            else:
                cbar=False
            ax.set_xticks(xticks)
            ax.set_xticklabels(lt_labels)
            ax.yaxis.set_major_locator(lct)
            ax.set_ylim([0,max_colat])

            # Use text function to manually add pretty ticks.
            ax.set_yticklabels('') # old ticks off.
            opts = {'size':yticksize, 'rotation':-45, 'ha':'center', 'va':'center'}
            for theta in [80.,70.,60.]:
                txt = '{:02.0f}'.format(theta)+r'$^{\circ}$'
                ax.text(pi/4., 90.-theta, txt, color='w', weight='heavy', **opts)
                ax.text(pi/4., 90.-theta, txt, color='k', weight='light', **opts)

            return fig, ax, cnt1, cbar

    class OvalDebugFile(PbData):
        '''
        The auroral oval calculations in RIM may spit out special debug files that
        are extremely useful.  This class handles reading and plotting the data
        contained within those files.
        '''

        def __init__(self, infile, *args, **kwargs):
            super(OvalDebugFile, self).__init__(*args, **kwargs)
            self.attrs['file']=infile
            self._readascii()

        def _readascii(self):
            '''
            Loads data & populates object upon instantiation.
            '''

            import datetime as dt

            # Slurp in lines:
            f = open(self.attrs['file'])
            lines = f.readlines()
            f.close()

            # Parse header to get longitude in radians:
            l = lines.pop(0)
            self['lon'] = np.array(l.split('=')[-1].split(), dtype=float)* np.pi/180.
            l = lines.pop(0)

            # Some helper vars:
            nLons = self['lon'].size
            nLine = len(lines)

            # Create container arrays:
            self['time'] = np.zeros(nLine, dtype=object)
            self['oval'] = np.zeros( (nLine, nLons) )

            # Parse rest of file:
            for j,l in enumerate(lines):
                self['time'][j]   = dt.datetime.strptime(l[:19], '%Y %m %d %H %M %S')
                self['oval'][j,:] = l.split()[7:]

        def get_oval(self, time, interp=True):
            '''
            Given a datetime, *time*, interpolate the oval position to that time.
            If *time* is outside the range of self['time'], an exception will be
            raised.  Linear interpolation is used unless kwarg *interp* is False.
            In that case, the value of the oval nearest to *time* will be used
            without interpolation.

             Parameters
            ==========
            time : datetime.datetime
               The time at which the oval is requested.

            Returns
            =======
            oval : numpy.ndarray
               Oval location at time *time* in radians.  Zero corresponds to local
               noon.

            Other Parameters
            ================
            interp : bool
                Set interpolation behavior.  If **True**, linear interpolation is
                used.  If **False**, the oval location nearest to *time* is used
                without interpolation.

            Examples
            ========
            >>>
            '''

            from matplotlib.dates import date2num

            # If *time* is outside the bounds of self['time'], raise exception
            # to avoid extrapolation.
            if time<self['time'][0] or time>self['time'][-1]:
                raise ValueError('Given time outside object range ' +
                                 'and requires extrapolation')

            # Turn datetimes into numbers.
            time     = date2num(time)
            ovaltime = date2num(self['time'])

            # Start by obtaining the indices of the time array that bound
            index = np.arange(self['time'].size)
            i1 = index[time>=ovaltime][-1] # Value before time
            i2 = index[time<=ovaltime][ 0] # Value after time

            # Check for exact matches:
            dT1 = np.abs(ovaltime[i1]-time)
            dT2 = np.abs(ovaltime[i2]-time)
            if dT1==0: return self['oval'][i1]
            if dT2==0: return self['oval'][i2]

            # If no interpolation, just send back nearest neighbor:
            if not interp:
                if dT1<dT2:
                    return self['oval'][i1]
                else:
                    return self['oval'][i2]

            # If interpolation, get slope and y-intercept vectors:
            dT = ovaltime[i2]-ovaltime[i1]
            m  = (self['oval'][i2]-self['oval'][i1])/dT
            b  = self['oval'][i2] - m*ovaltime[i2]

            # Interpolate and return:
            return m*time + b

        def add_oval_line(self, time, *args, **kwargs):
            '''
            Adds the location of the auroral oval at time *time* as a line plot on to
            *target*, which must be a Matplotlib figure/axes.

            If *target* not given, a new axes object is created.  If *target* is
            not an existing axes object, a new axes object is created and customized
            to be an ionosphere polar plot.

            Extra arguments/kwargs are sent to :meth:`matplotlib.pyplot.plot`.

            Parameters
            ==========
            time : integer or datetime
               Sets the time at which to plot the line.  If an integer is used,
               the integer is used to index the internal array naively.  If a
               datetime object is given and is within the bounds of self['time'],
               the oval location will be interpolated to *time*.  Control of the
               interpolation happens via the *interp* keyword.

            Returns
            =======
            fig  : matplotlib figure object
            ax   : matplotlib axes object
            line : matplotlib line object

            Other Parameters
            ================
            target : Figure or Axes
                If None (default), a new figure is generated from scratch.
                If a matplotlib Figure object, a new axis is created
                to fill that figure.
                If a matplotlib Axes object, the plot is placed
                into that axis.
            loc : int
                Use to specify the subplot placement of the axis
                (e.g. loc=212, etc.) Used if target is a Figure or None.
                Default 111 (single plot).
            interp : bool
                Control the behavior of time interpolation when the *time* argument
                is given as a datetime object.  Defaults to **True**, meaning that
                oval location is interpolated in time.

            Examples
            ========
            >>>
            '''

            import datetime as dt

            # Handle kwargs not being handed to plot:
            target=None; loc=111; interp=True
            if 'target' in kwargs: target=kwargs.pop('target')
            if 'interp' in kwargs: interp=kwargs.pop('interp')
            if 'loc'    in kwargs: loc=kwargs.pop('loc')

            # Set plot targets:
            fig, ax = set_target(target, polar=True, loc=loc)

            # Get oval interpolated in time:
            if type(time) == dt.datetime:
                oval = self.get_oval(time, interp=interp)
            elif type(time) == int:
                oval = self['oval'][time]
            else:
                raise TypeError('Unrecognized type for *time*:'+type(time))

            # Plot oval:
            line = ax.plot(self['lon']+np.pi/2., oval, *args, **kwargs)

            return fig, ax, line