import types
import numpy as np
import h5py
from pyspawn.fmsobj import fmsobj
from pyspawn.traj import traj
import general as gen
import os
import shutil
import complexgaussian as cg
import datetime
import time


class simulation(fmsobj):
    """Simulation object contains the current state of the simulation.
    It is analagous to the "bundle" object in the original FMS code"""

    def __init__(self):
        # traj is a dictionary of trajectory basis functions (TBFs)
        self.traj = dict()

        # centroids is a dictionary of TBFs representing the centroids
        # between the basis functions
        self.centroids = dict()

        # queue is a list of tasks to be run
        self.queue = ["END"]
        # tasktimes is a list of the simulation times associated with each task
        self.tasktimes = [1e10]

        # olapmax is the maximum overlap allowed for a spawn.  Above this,
        # the spawn is cancelled
        self.olapmax = 0.8

        # quantum time is the current time of the quantum amplitudes
        self.quantum_time = 0.0
        # quantum time is the current time of the quantum amplitudes
        self.quantum_time_half_step = 0.0
        # timestep for quantum propagation
        self.timestep = 0.0
        # quantum propagator
#         self.qm_propagator = "RK2"
        # quantum hamiltonian
#         self.qm_hamiltonian = "adiabatic"

        # maps trajectories to matrix element indices
        self.traj_map = dict()

        # dictionary with istates for each trajectory, needed for analysis
        self.istates_dict = dict()

        # quantum amplitudes
        self.qm_amplitudes = np.zeros(0, dtype=np.complex128)

        # energy shift for quantum propagation
        self.qm_energy_shift = 0.0

        # variables to be output to hdf5 mapped to the size of each data point
        self.h5_datasets = dict()
        self.h5_types = dict()

        # maximium quantum walltime in seconds
        self.max_quantum_time = -1.0

        # maximium walltime in seconds
        self.max_walltime = -1.0

    def from_dict(self, **tempdict):
        """Convert dict to simulation data structure"""

        for key in tempdict:
            if isinstance(tempdict[key], types.UnicodeType):
                tempdict[key] = str(tempdict[key])
            if isinstance(tempdict[key], types.ListType):
                if isinstance((tempdict[key])[0], types.FloatType):
                    # convert 1d float lists to np arrays
                    tempdict[key] = np.asarray(tempdict[key])
                if isinstance((tempdict[key])[0], types.StringTypes):
                    if (tempdict[key])[0][0] == "^":
                        for i in range(len(tempdict[key])):
                            tempdict[key][i] = eval(tempdict[key][i][1:])
                        tempdict[key] = np.asarray(tempdict[key],
                                                   dtype=np.complex128)
                else:
                    if isinstance((tempdict[key])[0], types.ListType):
                        if isinstance((tempdict[key])[0][0], types.FloatType):
                            # convert 2d float lists to np arrays
                            tempdict[key] = np.asarray(tempdict[key])
                        if isinstance((tempdict[key])[0][0],
                                      types.StringTypes):
                            if (tempdict[key])[0][0][0] == "^":
                                for i in range(len(tempdict[key])):
                                    for j in range(len(tempdict[key][i])):
                                        tempdict[key][i][j] = eval(tempdict[key][i][j][1:])
                                tempdict[key] = np.asarray(tempdict[key],
                                                           dtype=np.complex128)
            if isinstance(tempdict[key], types.DictType) :
                if 'fmsobjlabel' in (tempdict[key]).keys():
                    fmsobjlabel = (tempdict[key]).pop('fmsobjlabel')
                    obj = eval(fmsobjlabel[8:])()
                    obj.from_dict(**(tempdict[key]))
                    tempdict[key] = obj
                else: