# -*-coding:utf-8-*-
from __future__ import print_function as p2
useFuture = True
### FILE FOR DIRECT LAUNCH IN A STANDALONE FOLDER
import subprocess
import os, shutil
import scripts.metaSimu as MS
import scripts.mergeStats as mSt
import scripts.statCrunch as SC
import scripts.fileManip as FM

import time
import sys

# HANDLE LOCAL COMPUTATION; HAVE OUTPUTS FOR IT ##
# Generate 10 replications of simulation based on the current input lua

USE_DEPRECATED_ABSOLUTE_PATHFINDING = False
DISPLAY_PROGRESS = False

# generate folders environment
verbose = False
overwrite = "blackhole"  # TODO set as parameter
removeInterm = False  # delete grep files and temporar informations at the end of the execution

def getNRep(csv):
    t1 = 'outputs/R'
    p1 = csv.find(t1) + len(t1)
    p2 = csv.find("/", p1)
    try:
        return int(csv[p1:p2])
    except:
        if verbose:
            print("can't return the R value for csv", csv, "(p1:", p1, ",", "p2:", p2, ")")
        return 0

def getNsteps(outputDir):
    fileRun = outputDir + "/process.log"
    #TODO REAL PROCESSING OF LINES TO GET LAST VALUE
    f = open(fileRun, "r")
    lf = f.readlines()
    f.close()
    lastStep = 0
    lastLine = ""
    for l in lf:
        if 'Clock' in l:
            try:
                lastStep = int(l.split()[2])
            except:
                if "nit" not in l:
                    print("can't parse",l,"to get step")
    lastLine = lf[-1][:-1]
    return (lastStep, lastLine)


def delRawoutputsContent(rawout):
    for dir in ["genotypes", "quantitative", "pRocessed"]:
        folder = rawout + "/" + dir
        for the_file in os.listdir(folder):
            file_path = os.path.join(folder, the_file)
            try:
                if os.path.isfile(file_path):
                    os.unlink(file_path)
                    print(file_path, "unlinked")
            except Exception, e:
                print(e)

def removeAll(folder, recreate = []):
    if os.path.isdir(folder):
        print("ERASE", folder)
        for the_file in os.listdir(folder):
            file_path = os.path.join(folder, the_file)
            try:
                if os.path.isfile(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                print(e)
        for new_fol in recreate:
            os.makedirs(os.path.join(folder, new_fol))
    else:
        print("NO FOLDER")

def createStructure(outputFolder):
    try:
        os.makedirs(outputFolder)
        if verbose:
            print(outputFolder, "created (a).")
    except:
        print("folder", outputFolder, "already exists.")
    dfPath = {}
    for f in ["settings", "genotypes", "quantitative", "pRocessed"]:
        dfPath[f] = os.path.join(outputFolder, f)
        try:
            os.makedirs(dfPath[f])
        except:
            print("folder", dfPath[f], "already exists.")
            # DEL OUTPUTS CONTENT ?
    rez = "results"
    try:
        os.makedirs(rez)
    except:
        if verbose:
            print("folder", rez, "already exists.")

if USE_DEPRECATED_ABSOLUTE_PATHFINDING:
    localFol = os.path.dirname(os.path.abspath(__file__))
else:
    localFol = "."

mtRef = localFol
outRef = localFol + "/outputs"
inputDir = localFol + "/user_input"

treatDir = localFol + "/data_processing/"

def printState(txt, isOver = False): #DOES NOT REALLY WORK: DIRECTLY WRITE THE END
    if not isOver:
        print('\r' + txt, end='')
    else:
        print('\r' + txt)

def launchSimu(mtpDir, outputDir="outputs", nReplicates=10, verboseSimu=False):
    recompile = True
    if recompile:
        MS.checkBin(MTPdir=".", doRecomp=True, standalone=True)
        #MS.compileBin(MTPdir=".", isRedirect=False) # IF false and no binary, the first computation time will include the compiling time

    if overwrite == "blackhole":
        removeAll(outputDir, recreate = ["pRocessed"])
    for rep in range(1, nReplicates+1):
        if nReplicates == 1:
            outputDirRep = outputDir
            #subprocess.call(["sh", localFol + "/scripts/runMP.sh"]) #TODO test this part
        else:
            outputDirRep = outputDir + "/R"+str(rep)

        if overwrite == "replace":
            removeAll(outputDirRep)
        createStructure(outputDirRep)

        if USE_DEPRECATED_ABSOLUTE_PATHFINDING:
            nTotStep = MS.switchConfFile(os.path.join(mtpDir, outputDirRep), "user_input/simuConf.lua", isReg=False)
        else:
            nTotStep = MS.switchConfFile("../" + outputDirRep, "user_input/simuConf.lua", isReg=False)
        shutil.copytree("user_input", outputDirRep + "/settings/user_input")

        print("Simulation launched (may take a while)")

        if verboseSimu or nReplicates == 1:
            t0 = time.time()
            subprocess.call(["sh", localFol + "/scripts/runMP.sh"])
            t1 = time.time()
            print("EXEC TIME (may include compilation time:" + str(t1 - t0))
        else:
            txtRep = "Replication " +  str(rep) + "/" + str(nReplicates)

            MS.checkBin(MTPdir=".", doRecomp=True, standalone=True)
            t0 = time.time()
            printState(" ".join([txtRep, "running..."]))

            subprocess.call(["sh", localFol + "/scripts/launchDuplicate.sh", outputDirRep])
            nst = getNsteps(outputDirRep)
            messageEnd = ""
            if nst[1] != "End of simulation.":
                messageEnd = "("+nst[1]+")"
            printState(" ".join([txtRep, "over at step", str(nst[0]), "/", str(nTotStep), ".", messageEnd]), True)
            t1 = time.time()
            print("\tSimulation time: "+str(round(t1-t0, 2))+"s.")


def computeData(outRef, inputDir, cStamp, mode="All"):
    prefixPath = outRef + "/"
    processedDir = prefixPath + "pRocessed/"


    calcPheno = True
    calcGenet = True
    ltypes = {}#{"T1":["a"]}
    if mode[0] == "P":
        calcGenet = False
    elif mode == "G":
        calcPheno = False

    species = FM.getSpecies(inputDir+"/user_input/")
    #print('SPECIES:'+str(species))
    traitsByS = {}
    for S in species:
        try:
            locS = S+".lua"  # "S1.lua"
            traitsByS[S] = FM.checkTfromS(inputDir + "/user_input/species/" + locS)
            print("Fetch species "+S)  # traits", traits
        except:
            print("Tracking in rep did not find "+S)
            exit(0)  # blocking issue if S file not found. Can't continue
    for S in traitsByS:
        for t in traitsByS[S]:
            ltypes[t] = FM.checkLTfromT(inputDir + "/user_input/traits/" + t + ".lua")


    if calcGenet:
        print("* GENETIC COMPUTATION *")
        oCsvList = prefixPath + "GENESCSV_" + str(cStamp) + ".txt"
        subprocess.call(["sh", treatDir + "Sh_scripts/getGenesCsv.sh", prefixPath, oCsvList])
        #print "SH", treatDir + "Sh_scripts/getGenesCsv.sh", genePath, oCsvList
        fcsv = open(oCsvList, "r")
        lcsv = fcsv.readlines()
        fcsv.close()
        i, tot = 0, len(lcsv)
        for csv in lcsv:
            #print("csv"+csv)

            genePath = "/".join(csv.split("/")[:-1])
            lLoci = []

            sp, ti = csv.split("_")[-2], csv.split("_")[-1].split(".")[0]

            for t in traitsByS[sp]:
                lLoci += [SC.getLociList(csv, absol=False, trait=t, mode="standalone", splitLt=True)]
            ##print "LL", lLoci
            txtLociWithTrait = "-".join(traitsByS[sp]) + " " + "-".join(lLoci)

            #print "=== TREAT GENES csv # >", csv, "< ==="

            if DISPLAY_PROGRESS:
                txtComp = " ".join([csv.split("/")[-1][:-1] + " (", csv.split("/")[-3] + ")"])
                printState(" ".join([txtComp, "computing..."]))

            fileN = outRef.split("/")[-1].replace("_", "") + str(cStamp)
            nrep = getNRep(csv)
            dtcsv = [fileN, "0", str(nrep), sp, str(ti)]
            #print "dtcsv", dtcsv
            fileTag = "-".join(dtcsv)
            #print "TAG:", fileTag
            # launcher = prefixPath + "/" + "launchG-"+fileTag+".sh"
            launcher = prefixPath + "launchG-" + fileTag + ".sh"

            nameRout = "GRez_" + fileTag + "_" + str(cStamp) + ".rout"
            if not os.path.isfile(processedDir + nameRout):
                amorce2 = "cd " + treatDir + "R_scripts/\n"
                if USE_DEPRECATED_ABSOLUTE_PATHFINDING:
                    SC.generateMonoRLauncher(launcher, genePath, processedDir + nameRout, txtLociWithTrait, ti, sp, newamorce=amorce2)
                else:
                    SC.generateMonoRLauncher(launcher, "../../"+genePath, "../../"+processedDir + nameRout, txtLociWithTrait, ti, sp, newamorce=amorce2)
                subprocess.call(["sh", launcher])

                if DISPLAY_PROGRESS:
                    printState(" ".join([txtComp, "processed. (", str(round((i + 1) * 100. / tot, 1)), '% |',str( i + 1), '/', str(tot), ')']), True)
            else:

                if DISPLAY_PROGRESS:
                    printState(
                    " ".join([txtComp, "skipped. (", str(round((i + 1) * 100. / tot, 1)), '% |',str( i + 1), '/', str(tot), ')']),
                    True)
            i += 1


    if calcPheno:
        print("* PHENO COMPUTATION *")
        oCsvList = prefixPath + "PHENOCSV_" + str(cStamp) + ".txt"
        subprocess.call(["sh", treatDir + "Sh_scripts/getPhenoCsv.sh", prefixPath, oCsvList])
        fcsv = open(oCsvList, "r")
        lcsv = fcsv.readlines()
        i, tot = 0, len(lcsv)
        for csv in lcsv:
            phenPath = "/".join(csv.split("/")[:-1])

            sp, ti = csv.split("/")[-1].split("_")[0], csv.split("/")[-1].split("_")[1]

            if DISPLAY_PROGRESS:
                txtComp = " ".join([csv.split("/")[-1][:-1] + " (", csv.split("/")[-3] + ")"])
                printState(" ".join([txtComp, "computing..."]))
            for tr in traitsByS[sp]:
                lLoci = SC.getLociList(csv, absol=False, trait=tr, mode="standalone")
                #print "LL", lLoci

                #print "=== TREAT PHENO csv # >", csv,"< ==="

                fileN = outRef.split("/")[-1].replace("_", "") + str(cStamp)

                nrep = getNRep(csv)
                dtcsv = [fileN, "0", str(nrep), sp, tr, str(ti)]
                fileTag = "-".join(dtcsv)
                launcher = prefixPath + "launchP-"+fileTag+".sh"
                nameRout = "PRez_" + fileTag +"_" + str(cStamp)+ ".rout"
                if not os.path.isfile(processedDir + nameRout):
                    amorce2 = "cd " + treatDir + "R_scripts/\n"

                    if USE_DEPRECATED_ABSOLUTE_PATHFINDING:
                        SC.generateMonoRLauncher(launcher, phenPath, processedDir+nameRout, lLoci, ti, sp, mode="P",
                                         trait=tr, newamorce=amorce2)
                    else:
                        SC.generateMonoRLauncher(launcher, "../../"+phenPath, "../../"+processedDir + nameRout, lLoci, ti, sp, mode="P",
                                                 trait=tr, newamorce=amorce2)
                    subprocess.call(["sh", launcher])

                # PHENO BY LTYPE
                if len(ltypes[tr]) > 1:
                    #for lt in ltypes[tr]:
                    dtcsv = [fileN, "0", str(nrep), sp, tr, str(ti), "lt"]
                    fileTag = "-".join(dtcsv)
                    launcher = prefixPath + "launchPlt-" + fileTag + ".sh"
                    nameRout = "PRez_" + fileTag + "_" + str(cStamp) + ".rout"
                    if not os.path.isfile(processedDir + nameRout):
                        amorce2 = "cd " + treatDir + "R_scripts/\n"
                        if USE_DEPRECATED_ABSOLUTE_PATHFINDING:
                            SC.generateMonoRLauncher(launcher, phenPath, processedDir + nameRout, lLoci, ti, sp, mode="Plt",
                                                 trait=tr, newamorce=amorce2)
                        else:
                            #print("Generates MONORLauncher with new path :D")
                            SC.generateMonoRLauncher(launcher, "../../"+phenPath, "../../"+processedDir + nameRout, lLoci, ti, sp,
                                                     mode="Plt",
                                                     trait=tr, newamorce=amorce2)
                        subprocess.call(["sh", launcher])
                        #print "LT:", launcher

            if DISPLAY_PROGRESS:
                printState(" ".join([txtComp, "processed. (", str(round((i + 1) * 100. / tot, 1)), '% |',str( i + 1), '/', str(tot), ')']),
                                   True)
            i += 1

    if calcPheno:
        i = 0
        for S in traitsByS:
            if len(traitsByS[S]) > 1:
                print("* MULTITRAIT PHENO COMPUTATION *")

                oCsvList = prefixPath + "PHENOCSV_" + str(cStamp) + ".txt"
                subprocess.call(["sh", treatDir + "Sh_scripts/getPhenoCsv.sh", prefixPath, oCsvList])
                fcsv = open(oCsvList, "r")
                lcsv = fcsv.readlines()
                doneCSV = []
                for csv in lcsv:
                    sp, ti = csv.split("/")[-1].split("_")[0], csv.split("/")[-1].split("_")[1]

                    nrep = getNRep(csv)
                    if (sp, ti, nrep) not in doneCSV:

                        doneCSV += [(sp, ti, nrep)]

                        if DISPLAY_PROGRESS:
                            txtComp = csv.split("/")[-1][:-1] + " (" + csv.split("/")[-3] + ")"
                            printState(" ".join([txtComp, "ongoing..."]))

                        phenPath = "/".join(csv.split("/")[:-1])
                        fileN = outRef.split("/")[-1].replace("_", "") + str(cStamp)

                        dtcsv = [fileN, "0", str(nrep), sp, "".join(traitsByS[S]), str(ti)]
                        fileTag = "-".join(dtcsv)
                        launcher = prefixPath + "launchPmc-" + fileTag + ".sh"
                        nameRout = "PRezMC_" + fileTag + "_" + str(cStamp) + ".rout"


                        if not os.path.isfile(processedDir + nameRout):
                            amorce2 = "cd " + treatDir + "R_scripts/\n"
                            if USE_DEPRECATED_ABSOLUTE_PATHFINDING:
                                SC.generateMonoRLauncher(launcher, phenPath, processedDir + nameRout, step=ti, species=sp, mode="Pmc",
                                                 trait=traitsByS[S], newamorce=amorce2)
                            else:
                                SC.generateMonoRLauncher(launcher, "../../"+phenPath, "../../"+processedDir + nameRout, step=ti, species=sp,
                                                         mode="Pmc",
                                                         trait=traitsByS[S], newamorce=amorce2)
                            subprocess.call(["sh", launcher])

                        if DISPLAY_PROGRESS:
                            printState(" ".join([txtComp, "processed."]),
                            True)
                        i += 1


cStamp = int(time.time()) % 1000000

isrun = True
isstat = True
ismerge = True

modeComp = "All"


if len(sys.argv) > 1:
    if sys.argv[1] == "STAT":
        if "y" in input("Are you sure you want to destroy the potential old results? (y/n)"):
            isrun = False
        else:
            exit(0)
    elif sys.argv[1] in ["G", "Gst", "GST"]:
        isrun = False
        modeComp = "G"
    elif sys.argv[1] == "MERGE":
        isrun = False
        isstat = False
        ismerge = True
    else:
        print("arg not recognized", sys.argv[1])
        exit(0)

if isrun:
    print("RUN")
    nRep = FM.getReplicationsNumber(mtRef+"/user_input/simuConf.lua")
    launchSimu(mtRef, nReplicates=nRep)
if isstat:
    print("STATISTICS", modeComp)
    if overwrite == "blackhole":
        removeAll(mtRef + "/outputs/pRocessed/")
        removeAll(mtRef + "/results/")
    currPath = "."
    computeData(outRef, currPath, cStamp, mode=modeComp)#outRef = foo/outputs
    # merge rez
    if (verbose):
        print("outRef", outRef)
        print("localFol", localFol)

if ismerge:
    currPath = "."
    species = FM.getSpecies(currPath+"/user_input/")
    ds, de = mSt.getStats(outRef, prefixSh=localFol + "/scripts/", typ="standalone", species=species)
    pRez = mSt.prez(ds, pathInfos=outRef, widthMax=65)
    #print("* SUMMARY *")
    #print(pRez)
    pRezFull = mSt.prez(ds, pathInfos=outRef)
    #mSt.outCalc(ds, "outputs/Merged_values.txt", prefix=pRezFull)
    #outfol = "outputs"
    outfol = "results"
    mSt.outSplit(ds, outfol + "/Summary", species=species)#M"+str(cStamp))

    cStamp = int(time.time()) % 1000000
    mSt.outCalc(ds, "MERGEOUT_" + str(cStamp) + ".txt", prefix=mSt.prez(ds, pathInfos=outRef), displayVar=True)
    print ("* SUMMARY SPLITTED *")

    if removeInterm:
        subprocess.call(["sh", "scripts/cleanInterm.sh"])
    print("End of program. Results can be found in folder "+outfol)
