"""
This file is part of PUQ
Copyright (c) 2013 PUQ Authors
See LICENSE file for terms.
"""

import copy
import time, os, re, h5py, sys 
import numpy as np
from puq.testprogram import TestProgram
from numpy import ndarray
from puq.hdf import get_output_names
from logging import debug
from puq.util import vprint
from puq.options import options
from puq.jpickle import pickle, unpickle
from socket import gethostname
from puq.parameter import get_psamples
from puq.calibrate import calibrate

_vcache = {}
_dcache = {}

class Sweep(object):
    """
    Creates an object that contains all the information about
    a parameter sweep.

    Args:
      psweep: Parameter Sweep object. See :class:`PSweep`.
      host: Host object.  See :class:`Host`.
      prog: TestProgram object: See :class:`TestProgram`.
      caldata(array: Experimental data for calibration. Optional.
      calerr(float): Measurement error in the experimental data.
      description(string): Optional description of this run.
      sweepid: an optional string that uniquely identifies this sweep.
      This string is passed to the TestProgram by the Host object
      using the jobinfo parameter.
    """

    def __init__(self, psweep, host, prog, caldata=None, calerr=None, description='',sweepid=''):
        self.host = host
        self.host.sweepid=sweepid
        
        self._reinit = False

        if isinstance(prog, TestProgram):
            self.prog = prog
        else:
            self.prog = TestProgram(prog)

        if description == '':
            description = self.prog.desc
        self.description = description

        # optional calibration data
        self.caldata = caldata
        self.err = calerr

        # trying to get 10Hz resolution, 1 year clock
        secperyear = 365*24*60*60
        self.fname = 'sweep_%s' % int((time.time() % secperyear) * 10)
        self.psweep = psweep
        self.host.prog = copy.copy(self.prog) #else jsonpickle can't unpickle without error
        self.input_script = os.path.abspath(sys.argv[0])

    def _save_hdf5(self):
        debug('')
        h5 = h5py.File(self.fname + '.hdf5')

        # write HDF5 header information, once only
        if not 'version' in h5.attrs:
            h5.attrs['MEMOSA_UQ'] = 'MEMOSA'
            h5.attrs['version'] = 201
            #h5.attrs['id'] = self.id
            h5.attrs['date'] = time.strftime("%b %d %H:%M %Z %Y", time.localtime())
            h5.attrs['hostname'] = gethostname()
            h5.attrs['username'] = "puq user" #FR
            h5.attrs['UQtype'] = self.psweep.__class__.__name__.lower()
            h5.attrs['description'] = self.description

        # overwrite previous
        if 'input' in h5:
            del h5['input']
        if 'private' in h5:
            del h5['private']

        hp = h5.require_group('private')
        hp['sweep'] = pickle(self,max_depth=6)

        # in /input write the input params in json and regular arrays
        h = h5.require_group('input')

        # basic parameter table for non-python reading of the hdf5 file
        h.create_dataset('param_array',data=np.column_stack([p.values for p in self.psweep.params]),
                         compression='gzip',compression_opts=9)
        h['param_array'].attrs['name'] = [str(p.name) for p in self.psweep.params]
        h['param_array'].attrs['description'] = [str(p.description) for p in self.psweep.params]

        # json-pickled parameters
        h = h.require_group('params')
        for p in self.psweep.params:
            h[p.name] = pickle(p)
            if p.attrs!=None and len(p.attrs)>0:
                for attrtuple in p.attrs:
                    h[p.name].attrs[attrtuple[0]]=attrtuple[1]

        # input script
        if hasattr(self, 'input_script'):
            h5['input/scriptname'] = str(self.input_script)
            try:
                h5['input/script'] = open(self.input_script).read()
            except:
                h5['input/script'] = "Source was unavailable."
            h5.close()

    def _save_and_run(self,dryrun=False):
        self._save_hdf5()
        res = self.host.run(dryrun)
        if res:
            self._save_hdf5()
        return res

    def run(self, fn=None, overwrite=False, dryrun=False):
        """
        Calls PSweep.run() to run all the jobs in the Sweep.  Collect the data
        from the outputs and call the PSweep analyze method.  If the PSweep method
        has an iterative callback defined, call it, otherwise return.

        Args:
          fn(string): HDF5 filename for output. '.hdf5' will be
            appended to the filename if necessary. If fn is None,
            a filename will be generated starting with "sweep\_"
            followed by a timestamp.
          overwrite(boolean): If True and fn is not None, will
            silently overwrite any previous files of the same name.
          dryrun(boolean): Simulates what would happen if UQ is run, without
            actually running the TestProgram.
        Returns:
          True on success.
        """

        if fn is not None:
            self.fname = os.path.splitext(fn)[0]
            fn = self.fname + '.hdf5'
            if os.path.exists(fn):
                if not overwrite:
                    done = False
                    while 1:
                        ans = raw_input('%s already exists.  Replace (Y/N):' % fn)
                        try:
                            if ans.upper() == 'N':
                                done = True
                                break
                            elif ans.upper() == 'Y':
                                break
                        except:
                            pass
                        print "Please answer with 'Y' or 'N'\n"
                    if done:
                        sys.exit(-1)
                os.remove(fn)
        vprint(1, 'Saving run to %s.hdf5' % self.fname)
        return self.psweep.run(self,dryrun)

    def extend(self, num=None):
        return self.psweep.extend(num)

    def collect_data(self, hf=None):
        """ Collects data from captured stdout files and puts it in arrays
        in 'output/data'. Returns True on success.
        """

        need_to_close = False
        if hf is None:
            hf = h5py.File(self.fname + '.hdf5')
            need_to_close = True

        finished_jobs = self.host.collect(hf)
        self._extract_hdf5(hf, finished_jobs)

        has_data = 'data' in hf['output']
        if has_data:
            outd = hf['output/data']
            data = dict([(x, outd[x].value) for x in outd])
            params = dict([(p.name, p.values) for p in self.psweep.params])

        if need_to_close:
            hf.close()

        if not has_data and not self._reinit:
            print "WARNING: There is no data in the output section!"
            print "--Check that your runs completed successfully."
            return False
        return params, data

    def analyze_errors(self, hf):
        p = re.compile('Command exited with non-zero status \d+')
        for job in hf['output/jobs']:
            if job == 'time':
                continue
            err = hf['output/jobs/%s/stderr' % job].value
            res = p.findall(err)
            if res:
                print "Job %s: %s" % (job, res[0])
                for line in err.split('\n'):
                    if line != res[0] and not line.startswith('HDF5:{'):
                        print line
            elif len(err) == 0:
                print "Job %s never completed. Walltime exceeded?" % job

            results = False
            out = hf['output/jobs/%s/stdout' % job].value
            for line in out.split('\n'):
                if line.startswith('HDF5:{'):
                    results = True
                    break
            if not results:
                print "ERROR: Job %s has no output data in stdout." % job

    def analyze(self, verbose=False):
        """
        If not done already, collects the output from all the jobs into an HDF5 file.
        Parses any tagged data in the output and puts it in
        the output/data group in the HDF5 file.
        
        This function is called from the puq command line interface script automatically.
        """
        
        debug('')
        hf = h5py.File(self.fname + '.hdf5')
        if not self.host.status(quiet=1)[1]:
            print "Cannot collect data or perform analysis until all jobs are completed."
            print "You should do 'puq resume' to resume jobs."
            sys.exit(-1)

        # collect the data if it has not already been collected.
        has_data = 'output' in hf and 'data' in hf['output']
        if not has_data:
            print('No data found. Attempting to collect data')
            try:
                if not self.collect_data(hf):
                    raise Exception()
                self.psweep.analyze(hf)
            except:
                print 'Warning: analysis failed.'
                errors = 1

        # quick error check
        if 'data' in hf['output']:
            errors = 0
            try:
                options[self.psweep.__class__.__name__]['verbose'] = verbose
            except KeyError:
                options[self.psweep.__class__.__name__] = {'verbose': verbose}

            for var in hf['output/data']:
                if not isinstance(hf['output/data/%s' % var], h5py.Group):
                    tlen = len(hf['output/data/%s' % var].value)
                    num_jobs = len(hf['output/jobs'])-1
                    if tlen != num_jobs:
                        errors += 1
                        print "Expected %s data points for variable %s, but got %s." % (num_jobs, var, tlen)
                        self.analyze_errors(hf)
                        return errors

        if not 'psamples' in hf:
            s = get_psamples(self.psweep.params)
            if s is not None:
                hf['psamples'] = s

        # FIXME check for correlation if multiple outputs

        # calibrate
        if hasattr(self, 'caldata') and self.caldata is not None:
            self._calibrate(hf)

        hf.close()
        self._save_hdf5()
        return errors

    # Bayesian Calibration
    def _calibrate(self, hf):
        ovar = get_output_names(hf)[0]
        rs = unpickle(hf["/smolyak/%s/response" % ovar].value)
        self.psweep.params = calibrate(self.psweep.params, self.caldata, self.err, rs.eval)

    def _dump_hdf5_cache(self, hf, d):
        global _vcache, _dcache
        if len(_vcache):
            if d:
                dgrp = hf.require_group('output/data')
            else:
                dgrp = hf.require_group('output/jobs')
            for n in _vcache:
                if n in dgrp:
                    del dgrp[n]
                adata = _vcache[n]
                if d and len(adata.shape) > 1:
                    # Data is a multidimensional array and we want to do analysis
                    # on each array element individually.  So we write them
                    # individually to /output/data
                    numvals = np.prod(adata.shape[1:])
                    for i, index in enumerate(np.ndindex(adata.shape[1:])):
                        name = '%s%s' % (n, [ind for ind in index])
                        data = adata.flatten()[i::numvals]
                        ds = dgrp.create_dataset(name, data=data)
                        ds.attrs["description"] = _dcache[n]
                else:
                    ds = dgrp.create_dataset(n, data=adata)
                    ds.attrs["description"] = str(_dcache[n])
            _vcache = {}
            _dcache = {}

    def _dump_hdf5(self, grp, line, job, mjob):
        debug("Dump %s : %s" % (job, line))
        #print "Dump %s : %s" % (job, line)
        global _vcache, _dcache

        # old format used single quotes.
        if line.startswith("{'"):
            #sometimes nans can still slip through. if so set it to null 
            #which gets converted to None when loading the json
            if "'value': nan" in line:
                line=line.replace("'value': nan","'value': null")                
                print('warning: output value for job {} was nan'.format(job))
            
            line = line.replace("'", '"')

        
        x = unpickle(line)
        v = x['value']
        n = x['name']

        if not n in _vcache:
            if isinstance(v, ndarray):
                _vcache[n] = np.empty([mjob] + list(v.shape))
            else:
                _vcache[n] = np.empty((mjob))
            _vcache[n].fill(np.nan)
            _dcache[n] = x['desc']

        _vcache[n][job] = v

    # Extract tagged data to hdf5
    def _extract_hdf5(self, hf, jobs):
        debug("Extract")
        mjob = np.max(jobs) + 1
        run_grp = hf.require_group('output/jobs')
        for ext in ['out', 'err']:
            for j in jobs:
                grp = run_grp.require_group(str(j))
                if not 'std%s' % ext in grp:
                    continue
                f = grp['std%s' % ext].value
                cont = False
                for line in f.splitlines():
                    if cont:
                        line = line.strip()
                        cline += line
                        if line.endswith(':5FDH'):
                            cont = False
                            cline = cline[:-5]
                            self._dump_hdf5(grp, cline, j, mjob)
                    elif line.startswith('HDF5:'):
                        line = line[5:].strip()
                        if line.endswith(':5FDH'):
                            line = line[:-5]
                            self._dump_hdf5(grp, line, j, mjob)
                        else:
                            cont = True
                            cline = line
                    elif ext == 'err':
                        print 'STDERR[job %d]: %s' % (j, line)
            self._dump_hdf5_cache(hf, ext == 'out')

    def resume(self):
        if hasattr(self.host, 'jobs'):
            self.host.run()
            self._save_hdf5()
            self.analyze()
        else:
            print "All jobs finished."
