from westpa.tools import WESTMasterCommand, WESTParallelTool
from warnings import warn

from westpa.cli.tools.w_reweight import RWMatrix

# Just a shim to make sure everything works and is backwards compatible.
# We're making sure it has the appropriate functions so that it can be called
# as a regular tool, and not a subcommand.


class PAMatrix(RWMatrix):
    subcommand = 'init'
    help_text = 'averages and CIs for path-tracing kinetics analysis'
    default_output_file = 'flux_matrices.h5'
    # This isn't strictly necessary, but for the moment, here it is.
    # We really need to modify the underlying class so that we don't pull this sort of stuff if it isn't necessary.
    # That'll take some case handling, which is fine.
    # default_kinetics_file = 'assign.h5'


class WReweight(WESTMasterCommand, WESTParallelTool):
    prog = 'w_postanalysis_matrix'
    subcommands = [PAMatrix]
    subparsers_title = 'calculate state-to-state kinetics by tracing trajectories'
    description = '''\
Generate a colored transition matrix from a WE assignment file. The subsequent
analysis requires that the assignments are calculated using only the initial and
final time points of each trajectory segment. This may require downsampling the
h5file generated by a WE simulation. In the future w_assign may be enhanced to optionally
generate the necessary assignment file from a h5file with intermediate time points.
Additionally, this analysis is currently only valid on simulations performed under
either equilibrium or steady-state conditions without recycling target states.

-----------------------------------------------------------------------------
Output format
-----------------------------------------------------------------------------

The output file (-o/--output, by default "reweight.h5") contains the
following datasets:

  ``/bin_populations`` [window, bin]
     The reweighted populations of each bin based on windows. Bins contain
     one color each, so to recover the original un-colored spatial bins,
     one must sum over all states.

  ``/iterations`` [iteration]
    *(Structured -- see below)*  Sparse matrix data from each
    iteration.  They are reconstructed and averaged within the
    w_reweight {kinetics/probs} routines so that observables may
    be calculated.  Each group contains 4 vectors of data:

      flux
        *(Floating-point)* The weight of a series of flux events
      cols
        *(Integer)* The bin from which a flux event began.
      cols
        *(Integer)* The bin into which the walker fluxed.
      obs
        *(Integer)* How many flux events were observed during this
        iteration.

-----------------------------------------------------------------------------
Command-line options
-----------------------------------------------------------------------------
'''


def entry_point():
    warn('{} is being deprecated.  Please use w_reweight instead.'.format(WReweight.prog))
    # If we're not really supporting subcommands...
    import sys

    try:
        if sys.argv[1] != 'init':
            sys.argv.insert(1, 'init')
    except Exception:
        sys.argv.insert(1, 'init')
    WReweight().main()


if __name__ == '__main__':
    entry_point()
