from datetime import datetime
import os
import shutil

from ismodelselfconsistent import ismodelselfconsistent
from loadresultsfromcluster import loadresultsfromcluster
from marshall import marshall
from pairoptions import pairoptions
from preqmu import *
from waitonlock import waitonlock


def solve(md, solutionstring, *args):
    """SOLVE - apply solution sequence for this model

    Usage:
        md = solve(md, solutionstring, varargin)
    
    where varargin is a list of paired arguments of string OR enums

    Solution types available comprise:
    - 'Stressbalance'      or 'sb'
    - 'Masstransport'      or 'mt'
    - 'Thermal'            or 'th'
    - 'Steadystate'        or 'ss'
    - 'Transient'          or 'tr'
    - 'Balancethickness'   or 'mc'
    - 'Balancevelocity'    or 'bv'
    - 'BedSlope'           or 'bsl'
    - 'SurfaceSlope'       or 'ssl'
    - 'Hydrology'          or 'hy'
    - 'DamageEvolution'    or 'da'
    - 'Gia'                or 'gia'
    - 'Esa'                or 'esa'
    - 'Sealevelrise'       or 'slr'
    - 'Love'               or 'lv'

    Extra options:
    - loadonly         : does not solve. only load results
    - runtimename      : true or false (default is true), makes name unique
    - checkconsistency : 'yes' or 'no' (default is 'yes'), ensures checks on 
                         consistency of model
    - restart          : directory name (relative to the execution directory) 
                         where the restart file is located

    Examples:
        md = solve(md, 'Stressbalance')
        md = solve(md, 'sb')
    """

    # Recover and process solve options
    if solutionstring.lower() == 'sb' or solutionstring.lower() == 'stressbalance':
        solutionstring = 'StressbalanceSolution'
    elif solutionstring.lower() == 'mt' or solutionstring.lower() == 'masstransport':
        solutionstring = 'MasstransportSolution'
    elif solutionstring.lower() == 'th' or solutionstring.lower() == 'thermal':
        solutionstring = 'ThermalSolution'
    elif solutionstring.lower() == 'st' or solutionstring.lower() == 'steadystate':
        solutionstring = 'SteadystateSolution'
    elif solutionstring.lower() == 'tr' or solutionstring.lower() == 'transient':
        solutionstring = 'TransientSolution'
    elif solutionstring.lower() == 'mc' or solutionstring.lower() == 'balancethickness':
        solutionstring = 'BalancethicknessSolution'
    elif solutionstring.lower() == 'bv' or solutionstring.lower() == 'balancevelocity':
        solutionstring = 'BalancevelocitySolution'
    elif solutionstring.lower() == 'bsl' or solutionstring.lower() == 'bedslope':
        solutionstring = 'BedSlopeSolution'
    elif solutionstring.lower() == 'ssl' or solutionstring.lower() == 'surfaceslope':
        solutionstring = 'SurfaceSlopeSolution'
    elif solutionstring.lower() == 'hy' or solutionstring.lower() == 'hydrology':
        solutionstring = 'HydrologySolution'
    elif solutionstring.lower() == 'da' or solutionstring.lower() == 'damageevolution':
        solutionstring = 'DamageEvolutionSolution'
    elif solutionstring.lower() == 'gia' or solutionstring.lower() == 'gia':
        solutionstring = 'GiaSolution'
    elif solutionstring.lower() == 'lv' or solutionstring.lower() == 'love':
        solutionstring = 'LoveSolution'
    elif solutionstring.lower() == 'esa':
        solutionstring = 'EsaSolution'
    elif solutionstring.lower() == 'slr' or solutionstring.lower() == 'sealevelrise':
        solutionstring = 'SealevelriseSolution'
    else:
        raise ValueError("solutionstring '%s' not supported!" % solutionstring)
    options = pairoptions('solutionstring', solutionstring, *args)

    # Recover some fields
    md.private.solution = solutionstring
    cluster = md.cluster
    if options.getfieldvalue('batch', 'no') == 'yes':
        batch = 1
    else:
        batch = 0

    # Check model consistency
    if options.getfieldvalue('checkconsistency', 'yes') == 'yes':
        if md.verbose.solution:
            print("checking model consistency")
        ismodelselfconsistent(md)

    # First, build a runtime name that is unique
    restart = options.getfieldvalue('restart', '')
    if restart == 1:
        pass  # do nothing
    else:
        if restart:
            md.private.runtimename = restart
        else:
            if options.getfieldvalue('runtimename', True):
                c = datetime.now()
                md.private.runtimename = "%s-%02i-%02i-%04i-%02i-%02i-%02i-%i" % (md.miscellaneous.name, c.month, c.day, c.year, c.hour, c.minute, c.second, os.getpid())
            else:
                md.private.runtimename = md.miscellaneous.name

    # If running QMU analysis, some preprocessing of Dakota files using models 
    # fields needs to be carried out
    if md.qmu.isdakota:
        md = preqmu(md, options)

    # Do we load results only?
    if options.getfieldvalue('loadonly', False):
        md = loadresultsfromcluster(md)
        return md

    # Write all input files
    marshall(md) # bin file
    md.toolkits.ToolkitsFile(md.miscellaneous.name + '.toolkits') # toolkits file
    cluster.BuildQueueScript(md.private.runtimename, md.miscellaneous.name, md.private.solution, md.settings.io_gather, md.debug.valgrind, md.debug.gprof, md.qmu.isdakota, md.transient.isoceancoupling) # queue file

    # Stop here if batch mode
    if options.getfieldvalue('batch', 'no') == 'yes':
        print('batch mode requested: not launching job interactively')
        print('launch solution sequence on remote cluster by hand')
        return md

    # Upload all required files:
    modelname = md.miscellaneous.name
    filelist = [modelname + '.bin ', modelname + '.toolkits ', modelname + '.queue ']
    if md.qmu.isdakota:
        filelist.append(modelname + '.qmu.in')
        
    if not restart:
        cluster.UploadQueueJob(md.miscellaneous.name, md.private.runtimename, filelist)

    # Launch job
    cluster.LaunchQueueJob(md.miscellaneous.name, md.private.runtimename, filelist, restart, batch)

    # Wait on lock
    if md.settings.waitonlock > 0:
        islock = waitonlock(md)
        if islock == 0:  # no results to be loaded
            print('The results must be loaded manually with md = loadresultsfromcluster(md).')
        else:  # load results
            if md.verbose.solution:
                print('loading results from cluster')
            md = loadresultsfromcluster(md)

    return md
