#!/usr/bin/env python3
import sys
import numpy as np
import matplotlib.pyplot as plt

import scipy.constants as const

if __name__ == '__main__':
    filename = sys.argv[1]
    savepic = sys.argv[2]

    with open(filename, 'r') as fp:
        line = fp.readline()
    ncols = len(line.split())

    fig = plt.figure()
    ax1 = fig.add_subplot(111)
    plt.ylabel(r'$W_{\alpha}$ / $g_0$ [$\Omega^{-1}$]')
    plt.xlabel('Wavenumber [1/cm]')
    ax2 = ax1.twiny()
    plt.xlabel('Wavenumber [1/cm]')
    plt.ylabel(r'$W_{\alpha}$ / $g_0$ [$\Omega^{-1}$]')

    g0 = (const.e**2)/const.h
    wavenum_1cm2mV = 1e5*const.c*const.h/const.e

    if ncols == 4:
        iets = np.loadtxt(filename, dtype={
            'names': ('mode_idx', 'wavenumber', 'energy', 'intensity'),
            'formats': ('i8', 'f8', 'f8', 'f8')
        })

        ax1.stem(iets['wavenumber'], iets['intensity']/g0)
        ax2.plot([min(iets['wavenumber'])*wavenum_1cm2mV,
                  max(iets['wavenumber'])*wavenum_1cm2mV],
                 2*[0.0], ' ')

    if ncols == 8:
        iets = np.loadtxt(filename, dtype={
            'names': ('mode_idx', 'wavenumber', 'energy', 'intensity_tot',
                      'intensity_uu', 'intensity_du',
                      'intensity_ud', 'intensity_dd'),
            'formats': ('i8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8')
        })
        ax1.stem(iets['wavenumber'], iets['intensity_uu']/g0)
        ax1.stem(iets['wavenumber'], -iets['intensity_dd']/g0)
        ax2.plot([min(iets['wavenumber'])*wavenum_1cm2mV,
                  max(iets['wavenumber'])*wavenum_1cm2mV],
                 2*[0.0], ' ')

    ax1.yaxis.grid()
    plt.xlabel('Bias Voltage [mV]')
    plt.savefig(savepic, bbox_inches='tight')
