#!/usr/bin/env python
# Metarunlog, experiment management tool.
# Author: Tom Sercu
# Date: 2015-01-23

from metarunlog import cfg # NOTE cfg is modified by MetaRunLog._loadBasedirConfig() with custom configuration.
from metarunlog.exceptions import *
from metarunlog.util import nowstring, sshify, _decode_dict, _decode_list, get_commit
from metarunlog.confParser import ConfParser
import os
import sys
import math
from os import listdir
from os.path import isdir, isfile, join, relpath, expanduser
import argparse
import subprocess
try:
    import simplejson as json # way better error messaging
except:
    import json
import pdb
import datetime
from collections import OrderedDict
from shutil import copy as shcopy
import shutil
import jinja2
from jinja2 import Template, Environment, meta
import itertools
import getpass

DEBUG = True

def initBasedir(basedir, args):
    """ init basedir"""
    # check if it already is a valid basedir, then raise an exception
    try:
        mrlState = MetaRunLog(basedir)
    except NoBasedirConfig:
        pass # should be raised
    else:
        return "{} already has valid .mrl.cfg file, remove it first to re-init.".format(basedir)
    # initialize .mrl.cfg file as a json-dict copy of cfg.py template
    with open(join(basedir, '.mrl.cfg'),"w") as fh:
        # copy the cfg attributes into ordered dict
        bconf = {k:getattr(cfg,k) for k in dir(cfg) if '__' not in k}
        bconf['outdir'] = args.outdir
        json.dump(bconf, fh, indent=2)
        fh.write("\n")
    try:
        mrlState = MetaRunLog(basedir)
    except Exception as e:
        print("initBasedir() - Failed to init basedirconfig")
        raise
    else:
        return join(basedir, ".mrl.cfg")

class MetaRunLog:
    """
    Metarunlog state from which all actions are coordinated.
    Needs to be run from a valid basedir (with .mrl.cfg file).
    Initialization overwrites cfg module from .mrl.cfg: metarunlog.cfg.py should be
    seen as a template for .mrl.cfg.
    """
    def __init__(self, basedir):
        self.basedir = basedir
        self._loadBasedirConfig()
        self.outdir  = join(self.basedir, cfg.outdir)
        if not isdir(self.outdir): raise InvalidOutDirException("Need output directory " + self.outdir + "  Fix your .mrl.cfg file")
        self.expDirList = sorted([x for x in listdir(self.outdir) if self._checkValidExp(x)])
        self.expList = [int(x) for x in self.expDirList]
        self.lastExpId = None if not self.expList else self.expList[-1]

    def _loadBasedirConfig(self):
        try:
            with open(join(self.basedir, '.mrl.cfg')) as fh:
                bconf = json.load(fh, object_hook=_decode_dict)
                for k,v in bconf
.iteritems():
                    setattr(cfg,k,v)
            re
turn True
        except Exception as e:
            raise NoBasedirConfig(self.basedir, str(e))

    def new(self,args):
        expId = 1 if not self.lastExpId else self.lastExpId+1
        expConfig = OrderedDict()
        expConfig ['expId'] = expId
        expConfig['basedir'] = self.basedir
        gitclean = not bool(args.notclean)
        expConfig['gitFailUntracked']= args.gitFailUntracked
        untracked = 'no' if args.gitFailUntracked == 'no' else 'normal'
        uncommited = subprocess.check_output("git status --porcelain --untracked=" + untracked, shell=True)
        #print "uncommited: ", uncommited
        if gitclean and uncommited:
            raise NoCleanStateException("new: uncommited files -- please commit changes first\n" + uncommited)
        expConfig['gitHash'], expConfig['gitDescription'] = get_commit()
        if not gitclean and uncommited: expConfig['gitHash'] += '-sloppy'
        expConfig['timestamp'] = nowstring()
        expConfig['user'] = getpass.getuser()
        expConfig['description'] = args.description if args.description else ""
        # make dir and symlink
        expDir = self._getExpDir(expId, True)
        if self.lastExpId and os.path.lexists(join(self.basedir,'last')): 
            os.remove(join(self.basedir,'last'))
        os.symlink(relpath(expDir,self.basedir), join(self.basedir, 'last'))
        os.mkdir(expDir)
        # After this point: expDir is made, no more exception throwing! copy config files
        try:
            if args.copyConfigFrom == 'no':
                pass
            else:
                srcDir = self._getExpDir(self._resolveExpId(args.copyConfigFrom))
                self._copyConfigFrom(srcDir, expDir)
        except (InvalidExpIdException, IOError) as e:
            print("Can't copy config files. ")
            print(e)
            print("Still succesfully created new experiment directory.")
        self._saveExpDotmrl(expDir, expConfig)
        self._putEmptyNote(expDir, expConfig['description'])
        self.lastExpId = expId
        self.expList.append(expId)
        return expDir

    def info(self, args):
        """ load info from experiment id and print it """
        expId, expDir, expConfig = self._loadExp(args.expId)
        subExpList = self._getSubExperiments(expDir)
        items = ['',expDir,'']
        maxkeylen = max(len(k) for k,v in expConfig.iteritems())
        items += ["%*s : %.80s" % (maxkeylen,k,str(v)) for k,v in expConfig.iteritems()]
        items += ["%*s : %.80s" % (maxkeylen, 'subExperiments', str(subExpList))]
        return "\n".join(items)

    def ls(self, args):
        for expId in self.expList[::-1]:
            expConfig = self._getExpConf(expId)
            row = self._fmtSingleExp(expId)
            if args.tm:
                row += "\t" + expConfig['timestamp']
            if args.ghash:
                row += "\t" + expConfig['gitHash']
            if args.gdesc:
                row += "\t" + expConfig['gitDescription']
            if args.desc:
                expDir = self._getExpDir(expId)
                row += "\t" + self._expIsDoneIndicator(expDir) + expConfig['description']
            print row
        return ""

    def makebatch(self, args):
        expId, expDir, expConfig = self._loadExp(args.expId)
        # check if already expanded in batch and cancel
        oldSubExpList = self._getSubExperiments(expDir)
        if oldSubExpList:
            if args.replace:
                # just remove all these subfolders. Drastic but it wouldn't make sense to keep results from old config files.
                print "Will remove all old subexperiments: {}".format(oldSubExpList)
                for subExp in oldSubExpList:
                    shutil.rmtree(join(expDir, subExp))
            else:
                raise BatchException("Experiment {} is already expanded into subexperiments: {}".\
                        format(expId, str(oldSubExpList)))
        # make ConfParser object
        confP = ConfParser(join(expDir, cfg.confTemplFile))
        # check if ConfParser output is non-empty
        if not confP.output:
            err = "ConfParser output is empty, are you sure {} is a batch template?"
            err = err.format(join(expDir, cfg.confTemplFile))
            raise BatchException(err)
        # generate output directories, write the output config files
        for i, params, fileContent in confP.output:
            self._newSubExp(expDir, i, {'params': params}, fileContent)
        # update the current .mrl file
        subExpList = self._getSubExperiments(expDir)
        return "Generated subExperiments {} for expId {}.\n".format(subExpList, expId)

    def analyze(self, args):
        ## Load modules only needed for analyzing and rendering the html file
        import pandas as pd
        import renderHtml
        try:
            import mrl_analyze
        except Exception as e:
            print("Could not load local mrl_analyze module: {}".format(str(e)))
        expId, expDir, expConfig = self._loadExp(args.expId)
        print "Analyze expId {} in path {}".format(expId, expDir)
        outdir = args.outdir if args.outdir else join(expDir, cfg.analysis_outdir)
        if not os.path.exists(outdir): os.mkdir(outdir)
        # load the params into dataframe
        subExpIds = self._getSubExperiments(expDir)
        if not subExpIds:
           