import numpy as np
import re, os, sys

def get_dyson(
    out_dir,
    traj,  
    njobs
    ):

    # TODO: verify frame corresponds to molpro run based on ID

    frames_per_job = len(traj.frames) / njobs
    nframes = len(traj.frames)
    for job in range(njobs+1):
        ind1 = frames_per_job * job
        ind2 = frames_per_job * (job + 1)
        if job == njobs: ind2 = nframes
        for ind, frame in enumerate(traj.frames[ind1 : ind2]):
            data = np.loadtxt('%s%02djob/%04d/norms.out' %  (out_dir, job, ind + ind1))
            dyson_norms = {}
            if len(data) == 3:
                dyson_norms = np.array([[int(data[1])+1, data[2]],])
                frame.properties['dyson_norms'] = dyson_norms
            else:
                dyson_norms = np.array([(int(d[1])+1, d[2]) for d in data])
            frame.properties['dyson_norms'] = dyson_norms

    return traj
