"""Driver to perform AFQMC calculation"""
import sys
import json
import time
import numpy
import warnings
import uuid
from math import exp
import copy
import h5py
from pauxy.estimators.handler import Estimators
from pauxy.propagation.utils import get_propagator_driver
from pauxy.qmc.options import QMCOpts
from pauxy.qmc.utils import set_rng_seed
from pauxy.systems.utils import get_system
from pauxy.trial_wavefunction.utils import get_trial_wavefunction
from pauxy.utils.misc import get_git_revision_hash, get_sys_info
from pauxy.utils.io import  to_json, serialise, get_input_value
from pauxy.walkers.handler import Walkers


class AFQMC(object):
    """AFQMC driver.

    Zero temperature AFQMC using open ended random walk.

    This object contains all the instances of the classes which parse input
    options.

    Parameters
    ----------
    model : dict
        Input parameters for model system.
    qmc_opts : dict
        Input options relating to qmc parameters.
    estimates : dict
        Input options relating to what estimator to calculate.
    trial : dict
        Input options relating to trial wavefunction.
    propagator : dict
        Input options relating to propagator.
    parallel : bool
        If true we are running in parallel.
    verbose : bool
        If true we print out additional setup information.

    Attributes
    ----------
    uuid : string
        Simulation state uuid.
    sha1 : string
        Git hash.
    seed : int
        RNG seed. This is set during initialisation in calc.
    root : bool
        If true we are on the root / master processor.
    nprocs : int
        Number of processors.
    rank : int
        Processor id.
    cplx : bool
        If true then most numpy arrays are complex valued.
    system : system object.
        Container for model input options.
    qmc : :class:`pauxy.state.QMCOpts` object.
        Container for qmc input options.
    trial : :class:`pauxy.trial_wavefunction.X' object
        Trial wavefunction class.
    propagators : :class:`pauxy.propagation.Projectors` object
        Container for system specific propagation routines.
    estimators : :class:`pauxy.estimators.Estimators` object
        Estimator handler.
    psi : :class:`pauxy.walkers.Walkers` object
        Walker handler. Stores the AFQMC wavefunction.
    """

    def __init__(self, comm, options=None, system=None, trial=None,
                 parallel=False, verbose=False):
        if verbose is not None:
            self.verbosity = verbose
            if comm.rank != 0:
                self.verbosity = 0
            verbose = verbose > 0 and comm.rank == 0
        # 1. Environment attributes
        if comm.rank == 0:
            self.uuid = str(uuid.uuid1())
            get_sha1 = options.get('get_sha1', True)
            if get_sha1:
                self.sha1, self.branch = get_git_revision_hash()
            else:
                self.sha1 = 'None'
            if verbose:
                self.sys_info = get_sys_info(self.sha1, self.branch, self.uuid, comm.size)
        # Hack - this is modified later if running in parallel on
        # initialisation.
        self.root = comm.rank == 0
        self.rank = comm.rank
        self._init_time = time.time()
        self.run_time = time.asctime()
        # 2. Calculation objects.
        # if comm.rank == 0:
            # system = get_system(sys_opts=options.get('model', {}),
                                     # mf=mf, verbose=verbose)
        # else:
            # system = None
        # self.system = comm.bcast(system, root=0)
        if system is not None:
            self.system = system
        else:
            sys_opts = get_input_value(options, 'model',
                                       default={},
                                       alias=['system'],
                                       verbose=self.verbosity>1)
            self.system = get_system(sys_opts, verbose=verbose)
        qmc_opt = get_input_value(options, 'qmc', default={},
                                  alias=['qmc_options'],
                                  verbose=self.verbosity>1)
        self.qmc = QMCOpts(qmc_opt, self.system,
                           verbose=self.verbosity>1)
        self.qmc.rng_seed = set_rng_seed(self.qmc.rng_seed, comm)
        self.cplx = self.determine_dtype(options.get('propagator', {}),
                                         self.system)
        twf_opt = get_input_value(options, 'trial', default={},
                                  alias=['trial_wavefunction'],
                                  verbose=self.verbosity>1)
        if trial is not None:
            self.trial = trial
        else:
            if comm.rank == 0:
                self.trial = (
                    get_trial_wavefunction(self.system, options=twf_opt,
                                           parallel=parallel, verbose=verbose)
                )
            else:
                self.trial = None
            self.trial = comm.bcast(self.trial, root=0)
        if self.system.name == "Generic":
            if self.trial.ndets == 1:
                if self.system.cplx_chol:
                    self.system.construct_integral_tensors_cplx(self.trial)
                else:
                    self.system.construct_integral_tensors_real(self.trial)
        if comm.rank == 0:
            self.trial.calculate_energy(self.system)
        prop_opt = options.get('propagator', {})
        self.propagators = get_propagator_driver(self.system, self.trial,
                                                 self.qmc, options=prop_opt,
                                                 verbose=verbose)
        self.tsetup = time.time() - self._init_time
        wlk_opts = get_input_value(options, 'walkers', default={},
                                   alias=['walker', 'walker_opts'],
                                   verbose=self.verbosity>1)
        est_opts = get_input_value(options, 'estimators', default={},
                                   alias=['estimates','estimator'],
                                   verbose=self.verbosity>1)
        est_opts['stack_size'] = wlk_opts.get('stack_size', 1)
        self.estimators = (
            Estimators(est_opts, self.root, self.qmc, self.system,
                       self.trial, self.propagators.BT_BP, verbose)
        )
        # Reset number of walkers so they are evenly distributed across
        # cores/ranks.
        # Number of walkers per core/rank.
        self.qmc.nwalkers = int(self.qmc.nwalkers/comm.size)
        # Total number of walkers.
        if self.qmc.nwalkers == 0:
            if comm.rank == 0:
                print("# WARNING: Not enough walkers for selected core count.")
                print("# There must be at least one walker per core set in the "
                      "input file.")
                print("# Setting one walker per core.")
            self.qmc.nwalkers = 1
        self.qmc.ntot_walkers = self.qmc.nwalkers * comm.size
        self.psi = Walkers(wlk_opts, self.system, self.trial,
                           self.qmc, verbose,
                           nprop_tot=self.estimators.nprop_tot,
                           nbp=self.estimators.nbp,
                           comm=comm)
        if comm.rank == 0:
            json.encoder.FLOAT_REPR = lambda o: format(o, '.6f')
            json_string = to_json(self)
            self.estimators.json_string = json_string
            self.estimators.dump_metadata()
            if verbose:
                self.estimators.estimators['mixed'].print_key()
                self.estimators.estimators['mixed'].print_header()

    def run(self, psi=None, comm=None, verbose=True):
        """Perform AFQMC simulation on state object using open-ended random walk.

        Parameters
        ----------
        psi : :class:`pauxy.walker.Walkers` object
            Initial wavefunction / distribution of walkers.
        comm : MPI communicator
        """
        if psi is not None:
            self.psi = psi
        self.setup_timers()
        (etot, e1b, e2b) = self.psi.walkers[0].local_energy(self.system)
        eshift = 0
        self.propagators.mean_local_energy = eshift.real
        # Calculate estimates for initial distribution of walkers.
        self.estimators.estimators['mixed'].update(self.system, self.qmc,
                                                   self.trial, self.psi, 0,
                                                   self.propagators.free_projection)
        # Print out zeroth step for convenience.
        if verbose:
            self.estimators.estimators['mixed'].print_step(comm, comm.size, 0, 1)

        for step in range(1, self.qmc.total_steps + 1):
            start_step = time.time()
            if step % self.qmc.nstblz == 0:
                start = time.time()
                self.psi.orthogonalise(self.trial,
                                       self.propagators.free_projection)
                self.tortho += time.time() - start
            start = time.time()
            for w in self.psi.walkers:
                if abs(w.weight) > 1e-8:
                    self.propagators.propagate_walker(w, self.system,
                                                      self.trial, eshift)
                if (abs(w.weight) > w.total_weight * 0.10) and step > 1:
                    w.weight = w.total_weight * 0.10
            self.tprop += time.time() - start
            if step % self.qmc.npop_control == 0:
                start = time.time()
                self.psi.pop_control(comm)
                self.tpopc += time.time() - start
            # calculate estimators
            start = time.time()
            self.estimators.update(self.system, self.qmc,
                                   self.trial, self.psi, step,
                                   self.propagators.free_projection)
            self.testim += time.time() - start
            self.estimators.print_step(comm, comm.size, step)
            if self.psi.write_restart and step % self.psi.write_freq == 0:
                self.psi.write_walkers(comm)
            if step < self.qmc.neqlb:
                eshift = self.estimators.estimators['mixed'].get_shift()
            else:
                eshift += (self.estimators.estimators['mixed'].get_shift()-eshift)
            self.tstep += time.time() - start_step

    def finalise(self, verbose=False):
        """Tidy up.

        Parameters
        ----------
        verbose : bool
            If true print out some information to stdout.
        """
        if self.root:
            if verbose:
                print("# End Time: {:s}".format(time.asctime()))
                print("# Running time : {:.6f} seconds"
                      .format((time.time() - self._init_time)))
                print("# Timing breakdown (per processor, per block/step):")
                print("# - Setup: {:.6f} s".format(self.tsetup))
                nsteps = max(self.qmc.nsteps, 1)
                nstblz = max(nsteps // self.qmc.nstblz, 1)
                npcon = max(nsteps // self.qmc.npop_control, 1)
                print("# - Step: {:.6f} s".format((self.tstep/nsteps)))
                print("# - Orthogonalisation: {:.6f} s".format(self.tortho/nstblz))
                print("# - Propagation: {:.6f} s".format(self.tprop/nsteps))
                print("# - Estimators: {:.6f} s".format(self.testim/nsteps))
                print("# - Population control: {:.6f} s".format(self.tpopc/npcon))


    def determine_dtype(self, propagator, system):
        """Determine dtype for trial wavefunction and walkers.

        Parameters
        ----------
        propagator : dict
            Propagator input options.
        system : object
            System object.
        """
        hs_type = propagator.get('hubbard_stratonovich', 'discrete')
        continuous = 'continuous' in hs_type
        twist = system.ktwist.all() is not None
        return continuous or twist

    def get_energy(self, skip=0):
        """Get mixed estimate for the energy.

        Returns
        -------
        (energy, error) : tuple
            Mixed estimate for the energy and standard error.
        """
        filename = self.estimators.h5f_name
        from pauxy.analysis import blocking
        try:
            eloc = blocking.reblock_local_energy(filename, skip)
        except IndexError:
            eloc = None
        except ValueError:
            eloc = None
        return eloc

    def setup_timers(self):
        self.tortho = 0
        self.tprop = 0
        self.testim = 0
        self.tpopc = 0
        self.tstep = 0


    def get_one_rdm(self, skip=0):
        """Get back-propagated estimate for the one RDM.

        Returns
        -------
        rdm : :class:`numpy.ndarray`
            Back propagated estimate for 1RMD.
        error : :class:`numpy.ndarray`
            Standard error in the RDM.
        """
        from pauxy.analysis import blocking
        filename = self.estimators.h5f_name
        try:
            bp_rdm, bp_rdm_err = blocking.reblock_rdm(filename)
        except IndexError:
            bp_rdm, bp_rdm_err = None, None
        return (bp_rdm, bp_rdm_err)
