#!/usr/bin/env python3

import argparse
import collections
import re

import numpy as np
import h5py
import hic.flow


_cent_bins = [5, 10, 20, 30, 40, 50, 60, 70, 80]
nch_pt_cent_bins = [0] + _cent_bins
# not yet enough stats for small bins
# flow_cent_bins = [0, 1, 2, 3, 4] + _cent_bins
flow_cent_bins = nch_pt_cent_bins


def middles(x):
    """
    Compute the midpoints of an array of edges.

    """
    x = np.asarray(x)
    return (x[1:] + x[:-1])/2


def split_cent_bins(array, bin_edges):
    """
    Split an array into chunks for each centrality bin.  The array must already
    be sorted by centrality along its first axis.

    """
    for a, b in zip(bin_edges[:-1], bin_edges[1:]):
        i, j = (int(array.shape[0]*c/100) for c in (a, b))
        yield array[i:j]


def compute_observables(input, output, num, error=False):
    """
    Read event observables from the input file, postprocess into batch
    observables, and write into the output file as the given batch number.

    """
    # sort by centrality (dNch/deta)
    dNch_deta = np.array(input['dNch_deta'])
    order = dNch_deta.argsort()[::-1]
    dNch_deta = dNch_deta[order]

    # save dNch/deta as a function of centrality
    output['dNch_deta/Y'][num] = [
        i.mean() for i in split_cent_bins(dNch_deta, nch_pt_cent_bins)
    ]
    if error:
        output['dNch_deta/Yerr'][num] = [
            i.std()/np.sqrt(i.size)
            for i in split_cent_bins(dNch_deta, nch_pt_cent_bins)
        ]

    # identified particles dN/dy and <pT>
    dN_dy, mean_pT = (
        collections.OrderedDict(
            (k, np.array(v)[order]) for k, v in input[i].items()
        )
        for i in ['dN_dy', 'mean_pT']
    )

    for name, data in dN_dy.items():
        output['dN_dy'][name]['Y'][num] = [
            i.mean() for i in split_cent_bins(data, nch_pt_cent_bins)
        ]
        if error:
            output['dN_dy'][name]['Yerr'][num] = [
                i.std()/np.sqrt(i.size)
                for i in split_cent_bins(data, nch_pt_cent_bins)
            ]

    for name, data in mean_pT.items():
        output['mean_pT'][name]['Y'][num] = [
            np.average(i, weights=j) for i, j in
            zip(
                split_cent_bins(data, nch_pt_cent_bins),
                split_cent_bins(dN_dy[name], nch_pt_cent_bins)
            )
        ]
        if error:
            output['mean_pT'][name]['Yerr'][num] = [
                i.std()/np.sqrt(i.size)
                for i in split_cent_bins(data, nch_pt_cent_bins)
            ]

    # flow cumulants
    M, Qn = (np.array(input['flow'][i])[order] for i in ['M', 'Qn'])

    cumulants = [
        hic.flow.Cumulant(m, *qn.T) for m, qn in
        zip(
            split_cent_bins(M, flow_cent_bins),
            split_cent_bins(Qn, flow_cent_bins)
        )
    ]

    for n in [2, 3, 4]:
        output['vn'][str(n)]['Y'][num] = [
            c.flow(n, 2, imaginary='zero') for c in cumulants
        ]
        if error:
            output['vn'][str(n)]['Yerr'][num] = [
                c.flow(n, 2, error=True, imaginary='zero')[1]
                for c in cumulants
            ]


def main():
    parser = argparse.ArgumentParser(
        description='compute batch observables')

    parser.add_argument('--overwrite', action='store_const', const='w',
                        default='w-', dest='mode',
                        help='overwrite existing output file')
    parser.add_argument('--pattern', default=r'\b[0-9]{3}\b',
                        help='regular expression pattern for extracting \
                              batch names from filenames')
    parser.add_argument('--error', action='store_true',
                        help='whether to compute statistical errors')

    parser.add_argument('input_files', nargs='+',
                        help='HDF5 event observable files')
    parser.add_argument('output_file', help='output HDF5 file')

    args = parser.parse_args()

    input_files = sorted(args.input_files)
    nbatches = len(input_files)

    with h5py.File(args.output_file, args.mode, libver='latest') as f:
        kwargs = dict(compression='lzf')

        f.create_dataset('batches',
                         data=np.array([re.search(args.pattern, i).group()
                                        for i in input_files], dtype=np.int32),
                         **kwargs)

        f.create_dataset('dNch_deta/x', data=middles(nch_pt_cent_bins),
                         **kwargs)

        empty_kwargs = dict(kwargs, dtype=np.float64,
                            shape=(nbatches, len(nch_pt_cent_bins) - 1))

        f.create_dataset('dNch_deta/Y', **empty_kwargs)

        if args.error:
            f.create_dataset('dNch_deta/Yerr', **empty_kwargs)

        for obs in ['dN_dy', 'mean_pT']:
            g = f.create_group(obs)
            for name in ['pion', 'kaon', 'proton']:
                h = g.create_group(name)
                h['x'] = f['dNch_deta/x']
                h.create_dataset('Y', **empty_kwargs)
                if args.error:
                    h.create_dataset('Yerr', **empty_kwargs)

        empty_kwargs['shape'] = (nbatches, len(flow_cent_bins) - 1)

        for n in [2, 3, 4]:
            g = f.create_group('vn/{}'.format(n))

            if n == 2:
                g.create_dataset('x', data=middles(flow_cent_bins), **kwargs)
            else:
                g['x'] = f['vn/2/x']

            g.create_dataset('Y', **empty_kwargs)
            if args.error:
                g.create_dataset('Yerr', **empty_kwargs)

        for n, infile in enumerate(input_files):
            print(n, infile)
            with h5py.File(infile, 'r') as i:
                compute_observables(i, f, n, error=args.error)


if __name__ == "__main__":
    main()
