# This is a script to convert ms output file(s) (with the '.msout' extension)
# either specified as a single file or as a directory containing multiple files,
# along with the msFlagFile used to create the ms simulations in the ms output files
# and the dummy population coordinate file to create necessary input
# files for SpaceMix, and most of the necessary input files for EEMS (another step
# to calculate the pairwise genetic differences using the R script needs to follow
# using the output of this script (the .4eemsdiffs file) to be fully ready for EEMS.

# Both the SpaceMix genotype counts and sampling counts files and the .4eemsdiffs files
# are formatted with individuals as rows, and SNP loci as columns, which is why they are
# generated at the same time with this script, but the SpaceMix files are aggregated at the 
# population level, and the EEMS file is not (it is still at the individual level). 
# In addition, SpaceMix counts the occurrences of a randomly chosen allele at each
# locus (following the directions from Bradburd et al 2016). The EEMS file does the same
# conversion of allele coding as PCA from the ms output ('0' becomes '1', and '1' becomes '0').

# While this does run conversions for SpaceMix and EEMS in the same script, it runs 
# them using separate portions of code to make it easy to de-couple them if ever necessary in 
# the future. This means the msout files are looped through twice (once for SpaceMix, once for
# EEMS) instead of once.

# This is working correctly for SpaceMix sampling file and allele counts conversion. 021817
# EEMS .4eemsdiffs working correctly 022117 - still need to implement construction of .coords
# and .outer files from the popn coords and the input .outer file. Need to output exact
# copies of the input .coord and .outer file for each iteration, with file stems identical
# to the .4eemsdiffs file stem.

# Fixed counting of alleles for EEMS and SpaceMix output to represent haploid individuals
# 022417.

# Updated to correctly parse haploid/diploid from the ms flag file and then convert the 
# ms output for .4eemsdiffs files EEMS accordingly. The SpaceMix conversion parser
# already was working correctly to accomodate haploid/diploid differences because it 
# aggregates by population using a list the length of <nsam> for ms (ie the number of 
# chromosomes that were simulated, which is parsed from the list after the -I flag in the 
# ms flag file that lists the population number for each individual (# indiv == # chromos) in
# haploid files or for each chromosome for diploid files and aggregates those entries 
# correctly in either case. 031617.

# Test verified both conversions working correctly with diploids, uneven sampling, and 
# extra -e flags in the ms flag file 031717.

# NEEDS TO INCLUDE IPMAP OUTPUT FOR EEMS 040217.

import os
import sys
import argparse
from random import randint

parser = argparse.ArgumentParser()
# Very neat. if calling from the directory with the .msout files to process, can put '-m .' to pass
# all files (that are then screened for .msout extension) in the current dir to the script to process.
parser.add_argument('-m', '--msOutputToConvert', help = "The ms output to convert for PCA. This can either be a single file (with a '.msout' extension) or a directory that contains multiple files with '.msout' extensions to individually convert. TIP: If want to convert all .msout files in the directory that this script is called from, use '.' as the -m flag input value. Note: if this is a directory specified by name, the name cannot contain spaces without escape characters.")
parser.add_argument('-f', '--msFlagFile', help = "The msFlagFile that was used to generate the ms output that is being converted for PCA. This is used to assign each individual to the correct population for the PCA run.")
parser.add_argument('-c', '--popnCoordinates', help = "The population coordinates file generated by the simulation creation script or the simulation sculpting script if populations have been removed from the original rectangular configuration.")
parser.add_argument('-o', '--outputFileStem', help = "The output file stem (no file extension) to use for the files converted for PCA. File extensions ('.eigenstratgeno', '.inp', '.snp', and '.par') will be added.")

args = parser.parse_args()

msFlagFile = args.msFlagFile

inputToConvert = args.msOutputToConvert

popnCoords = args.popnCoordinates

outFileStem = args.outputFileStem


print(inputToConvert)

# Determine whether the input is a single file or a directory with many files, and 
# parse the input into a list of file names to process.
def checkInput():

    # This is a list of files to convert. If the user passed a single file as input, then
    # that will be the only entry. If they passed a directory, then each file in that dir
    # with the '.msout' extension will be added to this list. This makes it possible to 
    # process the files in the same for loop regardless of whether the input to the script
    # was a single file or a directory.
    fileListToConvert = []

    # Check whether the input to convert is an individual file (has the '.msout' extension) and 
    # is an accessible file
    # or is an accessible directory (which will then loop through the files in that dir that have the .msout extension)
    if inputToConvert[-6:] == ".msout" and os.path.isfile(inputToConvert):
        print("This is an individual file.")
        fileListToConvert.append(inputToConvert)
    
    
    elif os.path.isdir(inputToConvert):
        print("This is a directory.")
        fileList = os.listdir(inputToConvert)
    
        fileCounter = 0
        for file in fileList:
            if file[-6:] == ".msout":
                fileCounter += 1
                fileListToConvert.append(file)
    
        # Print the number of files if there are any, otherwise exit.
        if fileCounter >= 1:
            print("There are {} files with .msout extensions to convert in this directory.".format(fileCounter))
        else:
            print("There are no files with the '.msout' extension to convert in this directory.")
            print("Exiting.")
            sys.exit()
        
    else:
        print("The input was not recognized as an accessible file with an '.msout' extension or as a directory containing accessible '.msout' file(s).")
        print("Exiting.")
        sys.exit()

    #print(fileListToConvert)
    return(fileListToConvert)

# This is a stripped down version of the msFlagFile reader/parser. It does not deal with 
# any of the pairwise migration connections; it parses/counts population sampling information
# from the sampling specifications at the start of the file.    
def readFlagFile():
    
    # Gives something like:
    # ["TESTER_022817_diploidSamples_", "120_screened_2_times_msFlagFile.txt"]
    inputFileSplit1 = msFlagFile.split("ms_nsam_")

    if len(inputFileSplit1[0].split("diploid")) >= 2:
        isDiploidFile = 1
    elif len(inputFileSplit1[0].split("haploid")) >= 2:
        isDiploidFile = 0
    else:
        print("WARNING - Could not correctly parse whether the input ms flag file simulated haploids or diploids. Will parse it as haploid.")
        isDiploidFile = 0
    
    with open(msFlagFile, 'r') as inFile:

        for lineNum, line in enumerate(inFile):
            splitLine = line.split()
            # This is the whole ms flag file call but split by whitespace (any white space,
            # including any number of spaces).
            #print(splitLine)
        
            for index, element in enumerate(splitLine):
                
                if element == '-s':
                    sFlagIndex = index
                    sFlagHolder = splitLine[index + 1]
                if element == '-I':
                    IFlagIndex = index
                    IFlagHolder = splitLine[index + 1]
        
    #print("The s flag location is: %d" % sFlagIndex)
    #print("The I flag location is: %d" % IFlagIndex)

    numPopnsIFlag = int(IFlagHolder)

    #print("There are %s populations to simulate based on the -I flag" % numPopnsIFlag)

    # ----------
    # This finds the num popns following the -I flag by looking for the first '-' flag,
    # not considering the '0.0' before the first flag, which initially sets 
    # all migration rates to 0. In most cases, the first flag will be '-m', but
    # this makes it flexible to cases whether the user has manually added other flags
    # e.g. '-ej' before the first '-m' flag. If this is the case, the first flag MUST
    # still come after the '0.0'. 

    for index, entry in enumerate(splitLine):
        if index != sFlagIndex and index != IFlagIndex and entry.startswith('-'):
            flagStartIndex = index
            break
    
    popnSampleList = splitLine[IFlagIndex + 2 : flagStartIndex - 1]

    # This will be a simple list with a length the same as the total number of sampled
    # individuals, and each entry will be the population number ('collapsed' to consecutively
    # count focal popns only) for that individual.  
    popnSamplingList = []
    
    # NOTE: At this point the population numbers are 'collapsed' to be consecutive among the 
    # focal populations that were simulated. Any intervening ghost populations in the population
    # numbering are disregarded, and the first focal population becomes population '1'.
    # Focal populations are entries > 0 in the populationSample list.
    numFocalPopn = 0
    
    for populationSample in popnSampleList:
        if int(populationSample) > 0:
            numFocalPopn += 1
            for indiv in range(int(populationSample)):
                popnSamplingList.append(numFocalPopn)
    
    #print(popnSamplingList)
    
    return(popnSamplingList, isDiploidFile)


def parse_msOutputForSpaceMix(fileListForConversion,popnSamplingList, isDiploidFile):

    for fileNum, file in enumerate(fileListForConversion):
        fileSplit = file.split('_')
        #print(fileSplit)
        
        # If this is the first file, then set up to create the sample sizes output file. Only needs
        # to be created for this file, because it is common (ie unchanged) input for the rest of the 
        # allele counts files that come from the same ms flag file.
        if fileNum == 0:
            sampleSizeOutputFile = outFileStem + '_commonForAllIterations_sampleSizesForSpaceMix.csv'
        
        # The file split will return the full filename string (and therefore have 
        # length 1) if it does not contain '_'
        if len(fileSplit) > 1:
            iterNumSplit = fileSplit[len(fileSplit) - 1].split('.')[0]
            
            alleleCountsOutputFile = outFileStem + '_Iter_' + str(iterNumSplit) + '_alleleCountsForSpaceMix.csv'
            
        # if the file name does not contain a parsable iteration number, then just make the 
        # input file name into the output name for this file to keep as much identifying information
        # as possible.    
        else:
            print("WARNING - could not parse the iteration number from file: {}.".format(file))
            print("The iteration number in the input file names should be formatted as: <restOfFileName>_Iter_<num>.msout.")
            print("Using the input file name to make the output file name instead of the user-specified stem for the output file name.")
            
            fileStem = file.split('.')[0]
            
            alleleCountsOutputFile = fileStem + '_alleleCountsForSpaceMix.csv'
            
        #print(alleleCountsOutputFile)
        #print(sampleSizeOutputFile)
        
        numIndividuals = 0

        with open(file,'r') as inputFile:
            
            # keep track of the number of individuals simulated in each of the ms reps.
            individualCounter = 0
            
            # This keeps the sum of the transformed allele values for all individuals
            # in the population for each locus 
            # (can range from 0 to 2*num individuals per population). This is tabulated 
            # after randomly selecting, for each locus, whether the '0' or the '1' allele
            # will be counted (as a diploid) for that locus.
            populationAlleleSum = 0
            
            # This keeps the sum of the number of chromosomes (as a diploid) sampled 
            # per population, and will always be 2*num individuals per population, but 
            # is still tabulated and is not hard-coded.
            populationSamplingSum = 0
            
            populationCounter = 0
            
            prevPopulationCounter = 0
            
            SNPDict = {}
            
            # If this is the first file, then set up to create the sample sizes output file. Only needs
            # to be created for this file, because it is common (ie unchanged) input for the rest of the 
            # allele counts files that come from the same ms flag file.
            if fileNum == 0:
                samplingDict = {}
            
            SNPcounter = 0
            
            # Keep track of whether the '0' or the '1' allele was randomly selected to 
            # count for each locus
            refAlleleList = []
            
            for line in inputFile:
                #print(line.strip())
            
                if line.strip() == '//':
                    
                    # Following the SpaceMix manual, need to randomly select the reference allele for each locus.
                    # Do that here.
                    refAllele = randint(0,1)
                    refAlleleList.append(refAllele)
                    #print("The reference allele is {}".format(refAllele))
                    
                    # If the SNP counter is not 0, then this is at least the second locus (SNP)
                    # entry in the ms output, and need to first add the entry for the last 
                    # population for the last locus
                    # to the SNPDict before starting to parse the entry for the current
                    # allele.
                    if SNPcounter != 0:
                        
                        if populationCounter not in SNPDict.keys():
                            #print("Adding.")
                            SNPDict[prevPopulationCounter] = str(populationAlleleSum) + ','
                            #print(SNPDict)
                            # If this is the first file, then process for the sampling counts file.
                            if fileNum == 0:
                                samplingDict[prevPopulationCounter] = str(populationSamplingSum) + ','
                            
                        else:
                            SNPDict[prevPopulationCounter] += str(populationAlleleSum) + ','
                            
                            # If this is the first file, then process for the sampling counts file.
                            if fileNum == 0:
                                samplingDict[prevPopulationCounter] += str(populationSamplingSum) + ','
                
                    
                    # This is reset here for each new allele ('//') and is copied to the correct
                    # dict entry and then erased
                    populationAlleleSum = 0
                    
                    # If this is the first file, then process for the sampling counts file.
                    if fileNum == 0:
                        populationSamplingSum = 0
                    
                    #print('Yes for //')
                    
                    individualCounter = 0
                    
                    populationCounter = 0
                    
                    prevPopulationCounter = 0
                    
                    # This counts the number of the current ms locus in the simulation (denoted by '//' rows in the file)
                    SNPcounter += 1
                    #print('The SNP counter is: {}'.format(str(SNPcounter)))
                    
                    # Make a list of all individual
                    # counter entries that resets for each file. This will be passed
                    # to the output creation nested function where it is sorted
                    # in increasing order and used to call the dict entries (with
                    # keys defined by the populationCounter entries) in the
                    # order they were encountered in the input file (this makes the
                    # output in identical order to the input).
                    
                    populationCounterList = []
                    
                    continue
    
                if line.startswith('segsites:'):
                    continue
                
                if line.startswith('positions:'):
                    continue
                
                # If this line is within the ms output entry (and therefore is a simulated seq)
                if line.strip() == '1' or line.strip() == '0':
                    
                    individualCounter += 1
                
                    #print("individ counter is: {}".format(individualCounter))
                
                    # Because of how the popnSamplingList is set up as a list the length of the
                    # number of individuals with each entry being the population for the individual,
                    # can use individualCounter as an index into the popnSamplingList to get the 
                    # population number for the current individual.
                
                    # Need the -1 indexing for the 0-based Python list.
                    # Note - this is converted to an int, which is critical for 
                    # correct sorting of the populationCounterList when making the output file.
                    populationCounter = int(popnSamplingList[individualCounter - 1])
                
                    #print("The population number is: {}".format(populationCounter))
                
                    # Add the population number to the list if it isn't there already.
                    if populationCounter not in populationCounterList:
                        populationCounterList.append(populationCounter)
                
                    # If this is at least the 2nd population (the prev counter > 0), and 
                    # is different from the previous popn, then we're done aggregating 
                    # samples into population-level sums for the previous population, so 
                    # add the entry for the prev popn to the SNPDict here before clearing the
                    # holders and starting to aggregate for the current popn.    
                    if prevPopulationCounter > 0 and prevPopulationCounter != populationCounter:
            
                    
                        if prevPopulationCounter not in SNPDict.keys():
                        
                            SNPDict[prevPopulationCounter] = str(populationAlleleSum) + ','
                        
                            # If this is the first file, then process for the sampling counts file.
                            if fileNum == 0:
                                samplingDict[prevPopulationCounter] = str(populationSamplingSum) + ','
                        
                        else:
                            SNPDict[prevPopulationCounter] += str(populationAlleleSum) + ','
                        
                            # If this is the first file, then process for the sampling counts file.
                            if fileNum == 0:
                                samplingDict[prevPopulationCounter] += str(populationSamplingSum) + ','
                        
                        #print("Before locus end.")
                        #print(SNPDict)
                        
                        # clear the holder after adding to dict.    
                        populationAlleleSum = 0
                    
                        # If this is the first file, then process for the sampling counts file.
                        if fileNum == 0:
                            populationSamplingSum = 0
                    
                    # If the allele for the current individual matches the randomly
                    # selected allele to count for this locus (defined above), then
                    # count it as 1 occurrence of the allele to count (assuming a haploid)
                    if int(line.strip()) == refAllele:
                        populationAlleleSum += 1
                    
                    # If this is the first file, then process for the sampling counts file.
                    # Any sampled allele at this locus represents 1 sampled chromosome 
                    # (assuming a haploid)
                    if fileNum == 0:
                        populationSamplingSum += 1
                    
                    # update the previous population counter with the curr population now.
                    prevPopulationCounter = populationCounter
            
            # Cleaning up at the end of the for loop (ie the last population value at the 
            # last locus)
            if SNPcounter != 0:
                        
                if populationCounter not in SNPDict.keys():
                    
                    SNPDict[prevPopulationCounter] = str(populationAlleleSum) + ','
                    
                    # If this is the first file, then process for the sampling counts file.
                    if fileNum == 0:
                        samplingDict[prevPopulationCounter] = str(populationSamplingSum) + ','
                    
                else:
                    SNPDict[prevPopulationCounter] += str(populationAlleleSum) + ','
                    
                    # If this is the first file, then process for the sampling counts file.
                    if fileNum == 0:
                        samplingDict[prevPopulationCounter] += str(populationSamplingSum) + ','
            
            # Now call the nested function to generate the SpaceMix output file from within the
            # for loop going through the file list (so it generates 1 set of output files for
            # each input file).
            createSpaceMixOutputFile(SNPDict, samplingDict, populationCounterList, alleleCountsOutputFile, sampleSizeOutputFile)
    

# This is called from within parse_msOutputForSpaceMix.        
def createSpaceMixOutputFile(SNPDict, samplingDict, populationCounterList, alleleCountsFileName, sampleSizeFileName):
    
    # sort this list and use this to access the dict entries (ie the SNP calls for each
    # individual) in the same order they appear in the input file.
    sortedPopulationCounterList = sorted(populationCounterList)
    #print(sortedPopulationCounterList)
    
    with open(alleleCountsFileName,'w') as alleleOutFile, \
        open(sampleSizeFileName, 'w') as sampleOutFile:
        for popnNum in sortedPopulationCounterList:
            #print('pop Num is: {}'.format(str(popnNum)))
            #print(SNPDict)
            #print(samplingDict)
            # Omit the trailing comma from the last value in the dict entry
            # when writing to file
            alleleOutFile.write(SNPDict[popnNum][:-1] + '\n')
            sampleOutFile.write(samplingDict[popnNum][:-1] + '\n')

# This is essentially the same output as the SpaceMix conversion above, but is still at
# the individual sample level, and not aggregated at the population level like the 
# SpaceMix output.
def parse_msOutputForEEMS(fileListForConversion,popnSamplingList, isDiploidFile):

    
    for fileNum, file in enumerate(fileListForConversion):
        fileSplit = file.split('_')
        #print(fileSplit)
        
        # The file split will return the full filename string (and therefore have 
        # length 1) if it does not contain '_'
        if len(fileSplit) > 1:
            iterNumSplit = fileSplit[len(fileSplit) - 1].split('.')[0]
            
            eems4DiffsOutputFile = outFileStem + '_Iter_' + str(iterNumSplit) + '.4eemsdiffs'
            
        # if the file name does not contain a parsable iteration number, then just make the 
        # input file name into the output name for this file to keep as much identifying information
        # as possible.    
        else:
            print("WARNING - could not parse the iteration number from file: {}.".format(file))
            print("The iteration number in the input file names should be formatted as: <restOfFileName>_Iter_<num>.msout.")
            print("Using the input file name to make the output file name instead of the user-specified stem for the output file name.")
            
            fileStem = file.split('.')[0]
            
            eems4DiffsOutputFile = fileStem + '.4eemsdiffs'
            
        #print(alleleCountsOutputFile)
        #print(sampleSizeOutputFile)

        # This is used to aggregate by pairs of chromosomes for each individual
        if isDiploidFile:
            diploidHolder = 0
            prevLine = 1

        with open(file,'r') as inputFile:

            # This is a boolean for whether looping through the ms output has reached the actual 
            # seq output block (which starts with the first '//' in the file)
            entriesStarted = 0
    
            # keep track of the number of individuals (populations) simulated in each of the ms reps (should
            # be the same in all the reps, but this is not enforced)
            individualCounter = 0
        
            SNPDict = {}
        
            # for degugging only
            SNPcounter = 0
        
            for line in inputFile:
                #print(line.strip())
        
                if line.strip() == '//':
                
                    #print('Yes for //')
                
                    individualCounter = 0
                
                    SNPcounter += 1
                    #print('The SNP counter is: {}'.format(str(SNPcounter)))
                
                    # This is wasteful computation, but make a list of all individual
                    # counter entries that resets for each file. This will be passed
                    # to the output creation nested function where it is sorted
                    # in increasing order and used to call the dict entries (with
                    # keys defined by the individualCounter entries) in the
                    # order they were encountered in the input file (this makes the
                    # output in identical order to the input).
                
                    individualCounterList = []
                
                    continue

                if line.startswith('segsites:'):
                    #print('in segsites')
                    numSNPs = int(line.strip().split(' ')[1])
                
                    continue
            
                if line.startswith('positions:'):
                    continue
            
                # If this line is within the ms output entry (and therefore is a simulated seq)
                if line.strip() == '1' or line.strip() == '0':
                
                    # For haploid parsing
                    if isDiploidFile == 0:
                        individualCounter += 1
                        individualCounterList.append(individualCounter)
                        alleleHolder = ""
                
                        # Change any '1' entries from ms (indicating derived allele at this locus)
                        # to '0' entries indicating 0 copies of the reference allele .
                        if line.strip() == '0':
                            alleleHolder = '1'
                    
                        # Change any '1' entries from ms (indicating derived allele at this locus)
                        # to '0' entries indicating 0 copies of the reference allele .
                    
                        if line.strip() == '1':
                            alleleHolder = '0'
                
                    # For diploid parsing
                    if isDiploidFile == 1:
                        if prevLine == 1:
                        
                            individualCounter += 1
                            individualCounterList.append(individualCounter)
                            alleleHolder = ""
                            
                            if line.strip() == '0':
                                diploidHolder += 1
                            if line.strip() == '1':
                                diploidHolder += 0
                            
                            prevLine = 2
                            #print(prevLine)
                            # This continue is critical.
                            continue
                        
                        if prevLine == 2:
                    
                            if line.strip() == '0':
                                diploidHolder += 1
                            if line.strip() == '1':
                                diploidHolder += 0
                            
                            alleleHolder = str(diploidHolder)
                            
                            diploidHolder = 0
                            
                            prevLine = 1
                            
                            #print(prevLine)
                        #print(alleleHolder)
                    
                    if individualCounter not in SNPDict.keys():
                        SNPDict[individualCounter] = alleleHolder + ','
                
                    else:
                        SNPDict[individualCounter] += alleleHolder + ','
                    
#                     individualCounter += 1
#                 
#                     individualCounterList.append(individualCounter)
#                     
#                     alleleHolder = ""
                    
                    # Do the same replacement as for PCA:
                    # Change any '0' entries from ms (indicating ancestral allele at this locus)
                    # to '1' entries indicating 1 copy of the reference allele (assuming haploid).
                    
#                     if line.strip() == '0':
#                         alleleHolder = '1'
#                     
#                     # Change any '1' entries from ms (indicating derived allele at this locus)
#                     # to '0' entries indicating 0 copies of the reference allele (assuming haploid).
#                     
#                     if line.strip() == '1':
#                         alleleHolder = '0'
                
#                     if individualCounter not in SNPDict.keys():
#                         SNPDict[individualCounter] = alleleHolder + ','
#                 
#                     else:
#                         SNPDict[individualCounter] += alleleHolder + ','
                        
            # Now call the nested function to generate the .4eemsdiffs output file from within the
            # for loop going through the file list (so it generates 1 set of output files for
            # each input file).
            create4EEMSDiffsOutputFile(SNPDict, individualCounterList, eems4DiffsOutputFile)
    
    
def create4EEMSDiffsOutputFile(SNPDict, individualCounterList, eems4DiffsOutputFileName):
    
    # sort this list and use this to access the dict entries (ie the SNP calls for each
    # individual) in the same order they appear in the input file.
    sortedIndividualCounterList = sorted(individualCounterList)
    
    outputHolder = ''
    
    with open(eems4DiffsOutputFileName,'w') as outFile:
        for indivNum in sortedIndividualCounterList:
            outFile.write(SNPDict[indivNum][:-1] + '\n')



# ------------
fileListToProcess = checkInput()
popnSamplingList, isDiploidFile = readFlagFile()
parse_msOutputForSpaceMix(fileListToProcess,popnSamplingList, isDiploidFile)
parse_msOutputForEEMS(fileListToProcess,popnSamplingList, isDiploidFile)

# print("The SNPDict is ")
# print(SNPDict)
# 
# print("The sampling dict is")
# print(samplingDict)
# 
# print("The popnCounterList is")
# print(populationCounterList)