#!/usr/bin/env python3
import argparse
# import math
import os
import glob

import matplotlib.pyplot as plt
# import matplotlib.ticker as ticker
import pandas as pd

import artistools as at

DEFAULTSPECPATH = '../example_run/spec.out'


def main(argsraw=None):
    """
        Plot the electron energy distribution
    """
    defaultoutputfile = 'plotnonthermal_cell{0:03d}_timestep{1:03d}.pdf'
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description='Plot ARTIS radiation field.')
    parser.add_argument('modelpath', nargs='?', default='',
                        help='Path to ARTIS folder')
    parser.add_argument('-listtimesteps', action='store_true', default=False,
                        help='Show the times at each timestep')
    parser.add_argument('-timestep', '-ts', type=int, default=-1,
                        help='Timestep number to plot')
    parser.add_argument('-timestepmax', type=int, default=-1,
                        help='Make plots for all timesteps up to this timestep')
    parser.add_argument('-modelgridindex', '-cell', type=int, default=0,
                        help='Modelgridindex to plot')
    parser.add_argument('-xmin', type=int, default=40,
                        help='Plot range: minimum energy in eV')
    parser.add_argument('-xmax', type=int, default=10000,
                        help='Plot range: maximum energy in eV')
    parser.add_argument('-o', action='store', dest='outputfile', default=defaultoutputfile,
                        help='Filename for PDF file')
    args = parser.parse_args(argsraw)

    if os.path.isdir(args.outputfile):
        args.outputfile = os.path.join(args.outputfile, defaultoutputfile)

    if args.listtimesteps:
        at.showtimesteptimes('spec.out')
    else:
        nonthermaldata = None
        nonthermal_files = (
            glob.glob(os.path.join(args.modelpath, 'nonthermalspec_????.out'), recursive=True) +
            glob.glob(os.path.join(args.modelpath, 'nonthermalspec-*.out'), recursive=True) +
            glob.glob(os.path.join(args.modelpath, 'nonthermalspec.out'), recursive=True))

        for nonthermal_file in nonthermal_files:
            print(f'Loading {nonthermal_file}...')

            nonthermaldata_thisfile = pd.read_csv(nonthermal_file, delim_whitespace=True)
            nonthermaldata_thisfile.query('modelgridindex==@args.modelgridindex', inplace=True)
            if len(nonthermaldata_thisfile) > 0:
                if nonthermaldata is None:
                    nonthermaldata = nonthermaldata_thisfile.copy()
                else:
                    nonthermaldata.append(nonthermaldata_thisfile, ignore_index=True)

        if args.timestep < 0:
            timestepmin = max(nonthermaldata['timestep'])
        else:
            timestepmin = args.timestep

        if not args.timestepmax or args.timestepmax < 0:
            timestepmax = timestepmin + 1
        else:
            timestepmax = args.timestepmax

        list_timesteps = range(timestepmin, timestepmax)

        for timestep in list_timesteps:
            nonthermaldata_currenttimestep = nonthermaldata.query('timestep==@timestep')

            if len(nonthermaldata_currenttimestep) > 0:
                print(f'Plotting timestep {timestep:d}')
                outputfile = args.outputfile.format(args.modelgridindex, timestep)
                make_plot(nonthermaldata_currenttimestep, timestep, outputfile, args)
            else:
                print(f'No data for timestep {timestep:d}')


def make_plot(nonthermaldata, timestep, outputfile, args):
    """
        Draw the bin edges, fitted field, and emergent spectrum
    """
    import numpy as np
    fig, axis = plt.subplots(1, 1, sharex=True, figsize=(6, 4),
                             tight_layout={"pad": 0.2, "w_pad": 0.0, "h_pad": 0.0})

    # ymax = max(nonthermaldata['y'])

    # nonthermaldata.plot(x='energy_ev', y='y', linewidth=1.5, ax=axis, color='blue', legend=False)
    axis.plot(nonthermaldata['energy_ev'], np.log10(nonthermaldata['y']), linewidth=2.0, color='blue')

    axis.annotate(f'Timestep {timestep:d}\nCell {args.modelgridindex:d}',
                  xy=(0.02, 0.96), xycoords='axes fraction',
                  horizontalalignment='left', verticalalignment='top', fontsize=8)

    axis.set_xlabel(r'Energy (eV)')
    axis.set_ylabel(r'log [y (e$^-$ / cm$^2$ / s / eV)]')
    # axis.yaxis.set_minor_locator(ticker.MultipleLocator(base=0.1))
    # axis.set_yscale("log", nonposy='clip')
    # axis.set_xlim(xmin=args.xmin, xmax=args.xmax)
    # axis.set_ylim(ymin=0.0, ymax=ymax)

    # axis.legend(loc='upper center', handlelength=2,
    #             frameon=False, numpoints=1, prop={'size': 13})

    print(f'Saving to {outputfile:s}')
    fig.savefig(outputfile, format='pdf')
    plt.close()


if __name__ == "__main__":
    main()

@@ -158,16 +158,34 @@
 s as pd%0A
+import numpy as np
 %0Aimport 
@@ -248,16 +248,246 @@
 .out'%0A%0A%0A
+def ar_xs(energy_ev, ionpot_ev, A, B, C, D):%0A    u = energy_ev / ionpot_ev%0A    if u %3C= 1:%0A        return 0%0A    return 1e-14 * (A * (1 - 1/u) + B * pow((1 - 1/u), 2) + C * np.log(u) + D * np.log(u) / u) / (u * pow(ionpot_ev, 2))%0A%0A%0A
 def main
@@ -4357,24 +4357,359 @@
 or='blue')%0A%0A
+    arr_xs = %5Bxs_fe1(en) for en in nonthermaldata%5B'energy_ev'%5D%5D%0A    # arr_xs_old = %5Bxs_fe1_old(en) for en in nonthermaldata%5B'energy_ev'%5D%5D%0A    arr_xs_times_y = %5Bxs_fe1(en) * y for en, y in zip(nonthermaldata%5B'energy_ev'%5D, nonthermaldata%5B'y'%5D)%5D%0A%0A    # axis.plot(nonthermaldata%5B'energy_ev'%5D, arr_xs_times_y, linewidth=2.0, color='blue')%0A%0A
     axis.ann

Add impact ionisation cross section function