"""

"""

import sys
import glob
import argparse
from crrlpy import crrls
from astropy.io import fits


def add_axis(head, axis, naxis, crpix, crval, cdelt, ctype, cunit):
    """
    Adds an axis to the fits cube
    """
    
    head['NAXIS'] = head['NAXIS'] + 1
    head['NAXIS{0}'.format(axis)] = naxis
    head['CRPIX{0}'.format(axis)] = crpix
    head['CRVAL{0}'.format(axis)] = crval
    head['CDELT{0}'.format(axis)] = cdelt
    head['CTYPE{0}'.format(axis)] = ctype
    head['CUNIT{0}'.format(axis)] = cunit
   

def cube2vel(cube, transition='RRL_CIalpha', z=0, f_col=3, v_col=3):
    """
    Read the frequency axis of a fits cube and creates and axis 
    with velocity with respect to the specified transition.
    
    """
    
    hdu = fits.open(cube, mode='update')
    head = hdu[0].header
    freq = crrls.get_axis(head, f_col)
    
    # Invert frequency axis when searching for lines if necessary
    fi = 1
    if freq[0] > freq[-1]: fi = -1
#    print(freq[::fi][0], freq[::fi][-1])

    # Find the lines in the cube
    qns, reff = crrls.find_lines_sb(freq[::fi]*1e-6, transition, z) # returns in MHz
    #print reff

    if not reff:
        print('No {0} line found.'.format(transition))
        print('Will now exit.')
        sys.exit(0)
    
    # Get a velocity axis
    vel = crrls.freq2vel(reff*1e6, freq)
    
    add_axis(head, v_col, len(vel), 1, vel[0], vel[1]-vel[0], 'VELO', 'm/s')
    
    hdu[0].header = head
    
    hdu.flush()
    hdu.close()

if __name__ == '__main__':
    
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument('cube', type=str,
                        help="Cube to process.\n" \
                             "E.g., \"lba_hgh_*.fits\" (string).\n" \
                             "Wildcards and [] accepted.")
    parser.add_argument('-t', '--transition', type=str, default='RRL_CIalpha',
                        help="Transition to convert in the spectra.\n" \
                             "E.g., RRL_CI13beta" \
                             "Default: RRL_CIalpha")
    parser.add_argument('--z', type=float, default=0.0, dest='z',
                        help="Redshift to apply to the transition rest frequency.\n" \
                             "Default: 0")
    parser.add_argument('--f_col', type=int, default=3,
                        help="Header axis with frequency values.\n" \
                             "This will be converted to velocity. Default: 3")
    parser.add_argument('--v_col', type=int, default=3,
                        help="Header axis where the velocity information will be stored.\n" \
                             "Default: 3")
    parser.add_argument('--v_rng', type=str, default=None,
                        help="Velocity range to keep in the cube.\n" \
                             "E.g., -100e3,100e3 will keep from -100 to 100 km/s.\n" \
                             "Default: keep all \n" \
                             "Not implemented.")
    args = parser.parse_args()
    
    cube2vel(args.cube, args.transition, args.z, args.f_col, args.v_col)

@@ -581,16 +581,40 @@
  v_col=3
+, unit='m/s', qnidx=None
 ):%0A    %22
@@ -984,50 +984,8 @@
 = -1
-%0A#    print(freq%5B::fi%5D%5B0%5D, freq%5B::fi%5D%5B-1%5D)
 %0A%0A  
@@ -1103,24 +1103,8 @@
  MHz
-%0A    #print reff
 %0A%0A  
@@ -1226,24 +1226,348 @@
 xit(0)%0A    %0A
+    # Handle multiple lines in one cube.%0A    if len(reff) %3E 1:%0A        print('More than one line in the cube.')%0A        if qnidx is None:%0A            print('Split the cube or set %60qnidx%60.')%0A            print('Will now exit.')%0A            sys.exit(0)%0A        else:%0A            qns = qns%5Bqnidx%5D%0A            reff = qns%5Bqnidx%5D%0A%0A
     # Get a 
@@ -1622,24 +1622,71 @@
  freq)%0A    %0A
+    if unit == 'km/s':%0A        vel = vel*1e-3%0A%0A
     add_axis
@@ -1747,19 +1747,112 @@
 O', 
-'m/s'
+unit
 )%0A    
+# Add principal quantum number to header.%0A    head%5B'PQN'%5D = (qns, 'Principal quantum number')%0A
 %0A   
@@ -1911,16 +1911,17 @@
 lose()%0A%0A
+%0A
 if __nam

Adds principal quantum number to cube