#!/usr/bin/env python

import argparse
import pandas as pd
import numpy as np

def getOptions():
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="Subset and reformat FlyBase ortholog file for comparison to a single species.")

    # Input data
    parser.add_argument(
        "-i",
        "--input",
        dest="inFile",
        required=True,
        help="Input FlyBase ortholog .tsv or .tsv.gz file."
    )
    parser.add_argument(
        "-s",
        "--species",
        dest="inS",
        required=True,
        help="Species from input FlyBase file to select Dmel orthologs for."
    )
    parser.add_argument(
        "-n",
        "--species-name",
        dest="inSname",
        required=False,
        help="Species name to give prefix for output file columns (e.g. sim for Dsim). Default is species given -s."
    )

    # Output data
    parser.add_argument(
        "-o",
        "--output",
        dest="outFile",
        required=True,
        help="Output CSV of ortholog pairs of Dmel and selected species."
    )

    args = parser.parse_args()
    return args

def main():
    # Get flybase ortholog file
    orthoDF = pd.read_csv(
        args.inFile,
        sep="\t", 
        comment="#",
        names=["FBgn_ID",
               "GeneSymbol",
               "Arm/Scaffold",
               "Location",
               "Strand",
               "Ortholog_FBgn_ID",
               "Ortholog_GeneSymbol",
               "Ortholog_Arm/Scaffold",
               "Ortholog_Location",
               "Ortholog_Strand",
               "OrthoDB_Group_ID"],
        keep_default_na=False
    )

    species = args.inS
    if args.inSname is not None:
        name = args.inSname
    else:
        name = species
    # Select for given species
    if orthoDF["Ortholog_GeneSymbol"].str.contains(species).any():
        orthoDF = orthoDF[orthoDF["Ortholog_GeneSymbol"].str.contains(species)]
    else:
        print("!!!ERROR: Species name {} not found in Ortholog_GeneSymbol column.".format(
                species
        ))
        exit()

    # Remove species from gene symbol of ortholog
    orthoDF[name+"_geneSymbol"] = orthoDF["Ortholog_GeneSymbol"].str[len(species)+1:]

    # Split coordinates into separate columns
    orthoDF["mel_start"] = orthoDF["Location"].str.split(".").str[0].str.strip()
    orthoDF["mel_end"] = orthoDF["Location"].str.split(".").str[-1].str.strip()
    orthoDF[name+"_start"] = orthoDF["Ortholog_Location"].str.split(".").str[0].str.strip()
    orthoDF[name+"_end"] = orthoDF["Ortholog_Location"].str.split(".").str[-1].str.strip()

    # Set strand to + if Strand value is 1 and - if strand value is -1
    strandChoices = ["+", "-"]
    strandConditions = [orthoDF["Strand"]==1, orthoDF["Strand"]==-1]
    strandConditionsOrtho = [orthoDF["Ortholog_Strand"]==1, orthoDF["Ortholog_Strand"]==-1]
    orthoDF["mel_strand"] = np.select(
            strandConditions,
            strandChoices,
            ""
    )
    orthoDF[name+"_strand"] = np.select(
            strandConditionsOrtho,
            strandChoices,
            ""
    )
    if (orthoDF["mel_strand"]=="").any() or (orthoDF[name+"_strand"]=="").any():
        print("!!!WARNING: Strand missing for at least one entry.")

    # Rename columns
    orthoDF = orthoDF.rename(columns={
            "FBgn_ID": "mel_geneID",
            "GeneSymbol": "mel_geneSymbol",
            "Arm/Scaffold": "mel_chrom",
            "Ortholog_FBgn_ID": name+"_geneID",
            "Ortholog_Arm/Scaffold": name+"_chrom",
            "OrthoDB_Group_ID": "orthoDB_groupID"
    })

    # flag_one2one_ortholog
    orthoDF["num_mel_geneID_2_"+name+"_geneID"] = orthoDF.groupby(name+"_geneID")["mel_geneID"].transform("nunique")
    orthoDF["num_"+name+"_geneID_2_mel_geneID"] = orthoDF.groupby("mel_geneID")[name+"_geneID"].transform("nunique")
    orthoDF["flag_one2one_ortholog"] = np.where(
            (orthoDF["num_mel_geneID_2_"+name+"_geneID"]==1)
            & (orthoDF["num_"+name+"_geneID_2_mel_geneID"]==1),
            1,
            0
    )

    # Output final file
    orthoDF[['mel_geneID',
             'mel_geneSymbol',
             'mel_chrom',
             'mel_start',
             'mel_end',
             'mel_strand',
             name+'_geneID',
             name+'_geneSymbol',
             name+'_chrom',
             name+'_start',
             name+'_end',
             name+'_strand',
             'orthoDB_groupID',
             "num_mel_geneID_2_"+name+"_geneID",
             "num_"+name+"_geneID_2_mel_geneID",
             "flag_one2one_ortholog"
         ]].to_csv(args.outFile, index=False)

if __name__ == '__main__':
    # Parse command line arguments
    global args
    args = getOptions()
    main()