# 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_arg
ument('-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 followin