import argparse
import logging

import numpy as np

import westpa
from westpa.core.segment import Segment
from westpa.core.states import InitialState
from westpa.core.data_manager import n_iter_dtype, seg_id_dtype

log = logging.getLogger('w_fork')


def entry_point():
    parser = argparse.ArgumentParser(
        'w_fork',
        description='''\
    Prepare a new weighted ensemble simulation from an existing one at a particular
    point. A new HDF5 file is generated. In the case of executable propagation,
    it is the user's responsibility to prepare the new simulation directory
    appropriately, particularly making the old simulation's restart data from the
    appropriate iteration available as the new simulations initial state data; a
    mapping of old simulation segment to new simulation initial states is
    created, both in the new HDF5 file and as a flat text file, to aid in this.
    Target states and basis states for the new simulation are taken from those
    in the original simulation.
    ''',
    )

    westpa.rc.add_args(parser)
    parser.add_argument(
        '-i',
        '--input',
        dest='input_h5file',
        help='''Create simulation from the given INPUT_H5FILE (default: read from
                                configuration file.''',
    )
    parser.add_argument(
        '-I',
        '--iteration',
        dest='n_iter',
        type=int,
        help='''Take initial distribution for new simulation from iteration N_ITER
                                (default: last complete iteration).''',
    )
    parser.add_argument(
        '-o',
        '--output',
        dest='output_h5file',
        default='forked.h5',
        help='''Save new simulation HDF5 file as OUTPUT (default: %(default)s).''',
    )
    parser.add_argument(
        '--istate-map',
        default='istate_map.txt',
        help='''Write text file describing mapping of existing segments to new initial
                                states in ISTATE_MAP (default: %(default)s).''',
    )
    parser.add_argument('--no-headers', action='store_true', help='''Do not write header to ISTATE_MAP''')
    args = parser.parse_args()
    westpa.rc.process_args(args)

    # Open old HDF5 file
    dm_old = westpa.rc.new_data_manager()
    if args.input_h5file:
        dm_old.we_h5filename = args.input_h5file
    dm_old.open_backing(mode='r')

    # Get iteration if necessary
    n_iter = args.n_iter or dm_old.current_iteration - 1

    # Create and open new HDF5 file
    dm_new = westpa.rc.new_data_manager()
    dm_new.we_h5filename = args.output_h5file
    dm_new.prepare_backing()
    dm_new.open_backing()

    # Copy target states
    target_states = dm_old.get_target_states(n_iter)
    dm_new.save_target_states(target_states, n_iter)

    # Copy basis states
    basis_states = dm_old.get_basis_states(n_iter)
    dm_new.create_ibstate_group(basis_states, n_iter=1)

    # Transform old segments into initial states and new segments
    # We produce one initial state and one corresponding
    # new segment for each old segment. Further adjustment
    # can be accomplished by using w_binning.
    old_iter_group = dm_old.get_iter_group(n_iter)
    old_index = old_iter_group['seg_index'][...]
    old_pcoord_ds = old_iter_group['pcoord']
    n_segments = old_pcoord_ds.shape[0]
    pcoord_len = old_pcoord_ds.shape[1]
    pcoord_ndim = old_pcoord_ds.shape[2]
    old_final_pcoords = old_pcoord_ds[:, pcoord_len - 1, :]

    istates = dm_new.create_initial_states(n_segments, n_iter=1)
    segments = []
    state_map_dtype = np.dtype([('old_n_iter', n_iter_dtype), ('old_seg_id', seg_id_dtype), ('new_istate_id', seg_id_dtype)])
    state_map = np.empty((n_segments,), dtype=state_map_dtype)
    state_map['old_n_iter'] = n_iter

    for iseg, (index_row, pcoord) in enumerate(zip(old_index, old_final_pcoords)):
        istate = istates[iseg]
        istate.iter_created = 0
        istate.iter_used = 1
        istate.istate_type = InitialState.ISTATE_TYPE_RESTART
        istate.istate_status = InitialState.ISTATE_STATUS_PREPARED
        istate.pcoord = pcoord

        segment = Segment(
            n_iter=1,
            seg_id=iseg,
            weight=index_row['weight'],
            parent_id=-(istate.state_id + 1),
            wtg_parent_ids=[-(istate.state_id + 1)],
            status=Segment.SEG_STATUS_PREPARED,
        )
        segment.pcoord = np.zeros((pcoord_len, pcoord_ndim), dtype=pcoord.dtype)
        segment.pcoord[0] = pcoord
        segments.append(segment)
        state_map[iseg]['old_seg_id'] = iseg
        state_map[iseg]['new_istate_id'] = istate.state_id

    dm_new.update_initial_states(istates, n_iter=0)
    dm_new.prepare_iteration(n_iter=1, segments=segments)

    # Update current iteration and close both files
    dm_new.current_iteration = 1
    dm_new.close_backing()
    dm_old.close_backing()

    # Write state map
    istate_map_file = open(args.istate_map, 'wt')
    if not args.no_headers:
        istate_map_file.write('# mapping from previous segment IDs to new initial states\n')
        istate_map_file.write('# generated by w_fork\n')
        istate_map_file.write('# column 0: old simulation n_iter\n')
        istate_map_file.write('# column 1: old simulation seg_id\n')
        istate_map_file.write('# column 2: new simulation initial state ID\n')

    for row in state_map:
        istate_map_file.write(
            '{old_n_iter:20d}    {old_seg_id:20d}    {new_istate_id:20d}\n'.format(
                old_n_iter=int(row['old_n_iter']), old_seg_id=int(row['old_seg_id']), new_istate_id=int(row['new_istate_id'])
            )
        )


if __name__ == '__main__':
    entry_point()
