#!/usr/bin/env python3

import argparse
import datetime
from itertools import chain, repeat
import logging
import os
import signal
import subprocess
import sys

import numpy as np
import h5py

import freestream
import frzout


def run_cmd(*args):
    """
    Run and log a subprocess.

    """
    cmd = ' '.join(args)
    logging.info('running command: %s', cmd)

    try:
        proc = subprocess.run(
            cmd.split(), check=True,
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            universal_newlines=True
        )
    except subprocess.CalledProcessError as e:
        logging.error(
            'command failed with status %d:\n%s',
            e.returncode, e.output.strip('\n')
        )
        raise
    else:
        logging.debug(
            'command completed successfully:\n%s',
            proc.stdout.strip('\n')
        )
        return proc


def read_text_file(filename):
    """
    Read a text file into a nested list of bytes objects,
    skipping comment lines (#).

    """
    with open(filename, 'rb') as f:
        return [l.split() for l in f if not l.startswith(b'#')]


class Parser(argparse.ArgumentParser):
    """
    ArgumentParser that parses files with 'key = value' lines.

    """
    def __init__(self, *args, fromfile_prefix_chars='@', **kwargs):
        super().__init__(
            *args, fromfile_prefix_chars=fromfile_prefix_chars, **kwargs
        )

    def convert_arg_line_to_args(self, arg_line):
        # split each line on = and prepend prefix chars to first arg so it is
        # parsed as a long option
        args = [i.strip() for i in arg_line.split('=', maxsplit=1)]
        args[0] = 2*self.prefix_chars[0] + args[0]
        return args


parser = Parser(description='run heavy-ion collision events')

parser.add_argument(
    'results', type=os.path.abspath,
    help='results output file'
)
parser.add_argument(
    '--buffering', type=int, default=0,
    help='results file buffer size'
)
parser.add_argument(
    '--nevents', type=int,
    help='number of events to run (default: run until interrupted)'
)
parser.add_argument(
    '--rankvar',
    help='environment variable containing process rank'
)
parser.add_argument(
    '--rankfmt',
    help='format string for rank integer'
)
parser.add_argument(
    '--workdir', default='.', type=os.path.abspath,
    help='working directory (default: the current directory)'
)
parser.add_argument(
    '--logfile', type=os.path.abspath,
    help='log file (default: stdout)'
)
parser.add_argument(
    '--loglevel', choices=['debug', 'info', 'warning', 'error', 'critical'],
    default='info',
    help='log level (default: %(default)s)'
)
parser.add_argument(
    '--trento-args', default='Pb Pb',
    help="arguments passed to trento (default: '%(default)s')"
)
parser.add_argument(
    '--tau-fs', type=float,
    help='free streaming time [fm] (default: no free streaming)'
)
parser.add_argument(
    '--vishnew-args', default='',
    help='arguments passed to vishnew (default: empty)'
)
parser.add_argument(
    '--Tswitch', type=float, default=.150,
    help='particlization temperature [GeV] (default: %(default).3f GeV)'
)


def run_events(args, results_file):
    """
    Run events as determined by user input:

        - Read options from `args`, as returned by `parser.parse_args()`.
        - Write results to binary file object `results_file`.

    """
    grid_step = 0.1
    grid_max = 13.05

    def _initial_conditions(nevents=10, initial_file='initial.hdf'):
        """
        Run trento and yield HDF5 datasets.

        """
        try:
            os.remove(initial_file)
        except FileNotFoundError:
            pass

        run_cmd(
            'trento',
            '--number-events {}'.format(nevents),
            '--grid-step {} --grid-max {}'.format(grid_step, grid_max),
            '--output', initial_file,
            args.trento_args
        )

        with h5py.File(initial_file, 'r') as f:
            yield from f.values()

    # if nevents was specified, generate that number of initial conditions
    # otherwise generate indefinitely
    initial_conditions = (
        chain.from_iterable(_initial_conditions() for _ in repeat(None))
        if args.nevents is None else
        _initial_conditions(args.nevents)
    )

    # enable free streaming if time is set
    enable_fs = args.tau_fs is not None

    # create sampler HRG object (to be reused for all events)
    hrg = frzout.HRG(args.Tswitch, species='urqmd', res_width=True)
    eswitch = hrg.energy_density()

    # append switching energy density to vishnew arguments
    vishnew_args = [
        args.vishnew_args,
        'initialuread=1 iein=0 t0={}'.format(args.tau_fs) if enable_fs else
        'initialuread=0 iein=1',
        'edec={}'.format(eswitch)
    ]

    # species (name, ID) for identified particle observables
    species = [
        ('pion', 211),
        ('kaon', 321),
        ('proton', 2212),
    ]

    # fully specify numeric data types, including endianness and size, to
    # ensure consistency across all machines
    float_t = '<f8'
    int_t = '<i8'
    complex_t = '<c16'

    # results "array" (one element)
    # to be overwritten for each event
    results = np.empty((), dtype=[
        ('initial_entropy', float_t),
        ('mult_factor', float_t),
        ('nsamples', int_t),
        ('dNch_deta', float_t),
        ('dN_dy', [(s, float_t) for (s, _) in species]),
        ('mean_pT', [(s, float_t) for (s, _) in species]),
        ('M', int_t),
        ('Qn', complex_t, 6),
    ])

    def run_single_event(ic):
        """
        Run the initial condition event contained in HDF5 dataset object `ic`
        and save observables to `results`.

        """
        results.fill(0)
        results['initial_entropy'] = ic.attrs['mult']

        # read initial density array from dataset
        density = np.array(ic)

        # write initial file[s] for vishnew
        if enable_fs:
            logging.info(
                'free streaming initial condition for %.3f fm',
                args.tau_fs
            )
            fs = freestream.FreeStreamer(density, grid_max, args.tau_fs)

            e = fs.energy_density()
            e_above = e[e > eswitch].sum()
            results['mult_factor'] = e.sum()/e_above if e_above > 0 else 1

            np.savetxt('ed.dat', e)
            for i in [1, 2]:
                np.savetxt('u{}.dat'.format(i), fs.flow_velocity(i))
            for ij in [(1, 1), (1, 2), (2, 2)]:
                np.savetxt('pi{}{}.dat'.format(*ij), fs.shear_tensor(*ij))
        else:
            logging.info('free streaming disabled')
            np.savetxt('sd.dat', density)

        # hydro
        run_cmd('vishnew', *vishnew_args)

        # read freeze-out surface data
        surface_data = np.array(read_text_file('surface.dat'), dtype=float)

        # end event if the surface is empty -- this occurs in ultra-peripheral
        # events where the initial condition doesn't exceed Tswitch
        if surface_data.size == 0:
            logging.info('empty surface, ending event')
            return

        # unpack surface_data columns:
        #   0    1  2  3         4         5         6    7
        #   tau  x  y  dsigma^t  dsigma^x  dsigma^y  v_x  v_y
        #   8     9     10    11    12    13    14    15
        #   pitt  pitx  pity  pixx  pixy  piyy  pizz  Pi
        x, sigma, v, _ = np.hsplit(surface_data, [3, 6, 8])
        pi = dict(zip(['xx', 'xy', 'yy'], surface_data.T[11:14]))
        Pi = surface_data.T[15]

        # create sampler surface object
        surface = frzout.Surface(x, sigma, v, pi=pi, Pi=Pi, ymax=2)

        minsamples, maxsamples = 10, 1000  # reasonable range for nsamples
        minparts = 10**5  # min number of particles to sample
        nparts = 0  # for tracking total number of sampled particles

        logging.info('sampling surface with frzout')

        # sample particles and write to file
        with open('particles_in.dat', 'w') as f:
            for nsamples in range(1, maxsamples + 1):
                parts = frzout.sample(surface, hrg)
                if parts.size == 0:
                    continue
                nparts += parts.size
                print('#', parts.size, file=f)
                for p in parts:
                    print(p['ID'], *chain(p['x'], p['p']), file=f)
                if nparts >= minparts and nsamples >= minsamples:
                    break

        logging.info('produced %d particles in %d samples', nparts, nsamples)

        if nparts == 0:
            logging.info('no particles produced, ending event')
            return

        results['nsamples'] = nsamples

        # hadronic afterburner
        run_cmd('afterburner particles_in.dat particles_out.dat')

        # read final particle data
        ID, charge, pT, phi, y, eta = (
            np.array(col, dtype=dtype) for (col, dtype) in
            zip(
                zip(*read_text_file('particles_out.dat')),
                (2*[int] + 4*[float])
            )
        )

        logging.info('computing observables')
        charged = (charge != 0)
        abs_eta = np.fabs(eta)

        results['dNch_deta'] = \
            np.count_nonzero(charged & (abs_eta < .5)) / nsamples

        abs_ID = np.abs(ID)
        midrapidity = (np.fabs(y) < .5)

        for name, i in species:
            cut = (abs_ID == i) & midrapidity
            N = np.count_nonzero(cut)
            results['dN_dy'][name] = N / nsamples
            results['mean_pT'][name] = (0. if N == 0 else pT[cut].mean())

        phi_alice = phi[charged & (abs_eta < .8) & (.2 < pT) & (pT < 5.)]
        results['M'] = phi_alice.size
        results['Qn'] = [np.exp(1j*n*phi_alice).sum() for n in range(1, 7)]

    nfail = 0

    # run each initial condition event and save results to file
    for n, ic in enumerate(initial_conditions):
        logging.info('starting event %d', n)

        try:
            run_single_event(ic)
        except Exception:
            logging.exception('event %d failed', n)
            nfail += 1
            if nfail > 3 and nfail/n > .5:
                logging.critical('too many failures, stopping events')
                break
            logging.warning('continuing to next event')
            continue

        results_file.write(results.tobytes())
        logging.info('event %d completed successfully', n)


def main():
    args = parser.parse_args()

    # must handle rank before setting up logging since it affects the log file
    if args.rankvar:
        rank = os.getenv(args.rankvar)
        if rank is None:
            print(
                'rank variable {} is not set'.format(args.rankvar),
                file=sys.stderr
            )
            sys.exit(1)

        if args.rankfmt:
            rank = args.rankfmt.format(int(rank))

        argdict = vars(args)

        # append rank to path arguments, e.g.:
        #   /path/to/output.log  ->  /path/to/output/<rank>.log
        for arg in ['workdir', 'results', 'logfile']:
            if argdict[arg] is not None:
                root, ext = os.path.splitext(argdict[arg])
                argdict[arg] = os.path.join(root, rank) + ext

    os.makedirs(args.workdir, exist_ok=True)
    os.chdir(args.workdir)

    os.makedirs(os.path.dirname(args.results), exist_ok=True)

    if args.logfile is None:
        logfile_kwargs = dict(stream=sys.stdout)
    else:
        logfile_kwargs = dict(filename=args.logfile, filemode='w')
        os.makedirs(os.path.dirname(args.logfile), exist_ok=True)

    logging.basicConfig(
        level=getattr(logging, args.loglevel.upper()),
        format='[%(levelname)s@%(relativeCreated)d] %(message)s',
        **logfile_kwargs
    )
    logging.captureWarnings(True)

    start = datetime.datetime.now()
    logging.info('started at %s', start)
    logging.info('arguments: %r', args)

    # translate SIGTERM to KeyboardInterrupt
    signal.signal(signal.SIGTERM, signal.default_int_handler)
    logging.debug('set SIGTERM handler')

    with open(args.results, 'wb', buffering=args.buffering) as f:
        try:
            run_events(args, f)
        except KeyboardInterrupt:
            # after catching the initial SIGTERM or interrupt, ignore them
            # during shutdown -- this ensures everything will exit gracefully
            # in case of additional signals (short of SIGKILL)
            signal.signal(signal.SIGTERM, signal.SIG_IGN)
            signal.signal(signal.SIGINT, signal.SIG_IGN)
            logging.info(
                'interrupt or signal at %s, cleaning up...',
                datetime.datetime.now()
            )

    end = datetime.datetime.now()
    logging.info('finished at %s, %s elapsed', end, end - start)


if __name__ == "__main__":
    main()
