# -*- coding: utf-8 -*-
"""
Created on Tue Jan 21 13:06:16 2025

@author: hkleikamp
"""

#%% Modules
from inspect import getsourcefile
import warnings

import sys
from pathlib import Path
import os
import pandas as pd
import numpy as np
from collections import Counter
import pyarrow
import pyarrow.csv as pc
from sklearn.neighbors import KDTree


# %% change directory to script directory (should work on windows and mac)
basedir = str(Path(os.path.abspath(getsourcefile(lambda: 0))).parents[0])
os.chdir(basedir)
print(os.getcwd())
#%% Arguments for execution from Spyder.
#Mass list path
input_file = str(Path(basedir, "test_mass_CASMI2022.txt"))  #: default: CASMI 2022 masses in CartMFP folder default: CartMFP folder/ test_mass_CASMI2022.txt
composition_file= str(Path(basedir,                         #Required: cartesian output file generated by space2cart.py
                           "Cart_Output","H[200]C[75]N[50]O[50]P[10]S[10]_b100000max1000rdbe-5_80_7gr_comp.npy")) 
mass_index_file = ""
mass_defect_file = "" 

mode = "pos"                     # ionization mode. Options: " ", "positive", "negative" # positive substracts electron mass, negative adds electron mass, "" doesn't add anything
adducts =["+H+","+Na+","+K",      # default positive adducts "--","+H+","+Na+","+K"
         "+-","+Cl-","-H+"]            # default negative adducts "+-","-H+","+Cl-" 
charges=[1]                            # default: [1]

#performance arguments
ppm = 5                 # ppm for formula prediction
top_candidates = 20     # only save the best predictions sorted by ppm (default 20)
pre_filter_mass= True   # prefilter masses based top candidates counts
keep_all    = False     # also display mass/ adduct combinations for which no molecular formula was found
add_formula = False     # add formula string
use_params  = True      # use space2cart parameter file if available

#filepaths
mass_table = str(Path(basedir, "mass_table.tsv"))           # table containing element masses, default: CartMFP folder/ mass_table.tsv"
MFP_output_folder = str(Path(basedir, "MFP_Output"))        # default: CartMFP folder / MFP_Output
MFP_output_filename=""                                      # default: CartMFP_ + input_filename + .tsv 

#debug
force_masscalc=False

composition_file="D:/231225_CartMFP/CartMFP_test/C[72]H[134]N[15]P[3]O[30]S[5]_b100000max1000rdbe-5_39_7grCommon_comp.npy"

#%% Arguments for execution from command line.

def parse_path(s):
    raw_s = r'{}'.format(s).replace("\\","/")
    return raw_s

if not hasattr(sys,'ps1'): #checks if code is executed from command line
    
    import argparse
    parser = argparse.ArgumentParser(
                        prog='CartMFP-cart2form',
                        description='molecular formula prediction, see: https://github.com/hbckleikamp/CartMFP')

    #key arguments
    parser.add_argument("-i", "--input_file",       required = False, type=parse_path, default=str(Path(basedir, "test_mass_CASMI2022.txt")),
                        help="Required: input mass file, can be txt or tabular format, can also be single digit or array or dataframe of masses")
    parser.add_argument("-cart", "--composition_file",       required = False, type=parse_path,
                        help="Required: input cartesian file, as constructed by space2cart.py")
    parser.add_argument("--mass_index_file", required = False,type=parse_path, default= "", 
                        help="Optional: cartesian index file, derived from composition input file unless specified")
    parser.add_argument("--mass_defect_file", required = False,type=parse_path, default= "", 
                        help="Optional: mass defect file, derived from composition input file unless specified")
    
    #output and utility filepaths
    parser.add_argument("-mass_table",                 default=str(Path(basedir, "mass_table.tsv")), required = False,type=parse_path, help="list of element masses")  
    parser.add_argument("-mfp_out",  "--MFP_output_folder", default=str(Path(basedir, "MFP_Output")), required = False,type=parse_path, help="Output folder for molecular formula prediction")   
    parser.add_argument("-out_file", "--MFP_output_filename",  default="", required = False,type=parse_path, help="filename of molecular formula prediction output")   
    
    #composition constraints
    parser.add_argument("-mode",  default="pos",   required = False, help="ionization mode: positive, negative or "" (ignore). This will subtract mass based on ion adducts. if "" is used, the exact masses are used")    
    parser.add_argument("-a","--adducts",  default=["+H+","+Na+","+K+","-+","+Cl-"],   
                        required = False, help="The ionization mode will determine used adducts. Syntax: 'sign element charge' eg. gain of H+,Na+,K+ for positive, and  Cl- or loss of H+ for negative ionization mode ")    
    parser.add_argument("-charges",  default=[1],   
                        required = False, help="Charge states considered ")      
    
    #performance arguments
    parser.add_argument("-ppm",  default=5, required = False, help="ppm mass error tolerance of predicted compositions")  
    parser.add_argument("-t","--top_candidates",  default=20, required = False, help="number of best candidates returned (sorted by mass error)",type=int)  
    parser.add_argument("-pre_filter_mass",  default=True, required = False, help="pre trim masses prior to mfp based on top candidates")  
    parser.add_argument("-keep_all",     default=False, required = False, help="keep masses with no predicted formula in output")  
    parser.add_argument("-add_formula",  default=False, required = False, help="add element string outputs")  
    parser.add_argument("-use_params",  default=True, required = False, help="use parameter file for database composition if available")  


    
    args = parser.parse_args()
    args = {k:v for k,v in vars(parser.parse_args()).items() if v is not None}
    
    print("")
    print(args) 
    print("")
    locals().update(args)
    

#charges and adducts need command line string parsing
if type(charges)==str: charges=[int(i.strip()) for i in charges.split(",")]
if type(adducts)==str: adducts=[i.strip("[").strip("]").strip() for i in adducts.strip("[").strip("]").split(",")]
        
# %% Get elemental metadata (or replace with utils table)
if os.path.exists(mass_table): mdf = pd.read_csv(mass_table, index_col=[0], sep="\t")
else:                         raise ValueError("Mass table "+mass_table+" not found!")
mdf=mdf["mass"]

emass = 0.000548579909  # electron mass
mdf.loc["+"]=-emass
mdf.loc["-"]=+emass


#%% Functions

# read input table (dynamic delimiter detection)
def read_input(tabfile, *,
               Keyword=[], #rewrite multiple keywords,
               xls=False,
               dlim=""
               ):
    
    
    if type(Keyword)==str: Keyword=[i.strip() for i in Keyword.split(",")]
    
    if type(tabfile)!=str: #no string -> no filepath
    
        #if is DataFrame
        if isinstance(tabfile, pd.DataFrame): 
            print("DataFrame detected")
            if len(Keyword):
                if any(tabfile.columns.isin(Keyword)):
                    if len(Keyword)>1:
                        print("more than one Keyword found, using first Keyword")
                    return True,pd.DataFrame(tabfile[Keyword[0]].values,columns=["mz"])
            else:
                print("no Keyword found, using first column")
                return True,pd.DataFrame(tabfile.iloc[:,0].values,columns=["mz"])
                
        
        if isinstance(tabfile, pd.Series):
            print("Series detected")
            return True,pd.DataFrame(tabfile.values,columns=["mz"])
        

        if isinstance(tabfile,np.ndarray) or type(tabfile)==list() or  type(tabfile)==tuple():
            print("iterable detected")
            return True,pd.DataFrame(np.array(tabfile).astype(float),columns=["mz"])

        if type(tabfile)==int or type(tabfile)==float:
            print("single digit detected")
            return True,pd.DataFrame(np.array([tabfile]).astype(float),columns=["mz"])
    
    
    if os.path.exists(tabfile):
    
        #numpy format
        if tabfile.endswith(".npy"): 
            tab=np.load(tabfile)#,allow_pickle=True)
            if len(Keyword): tab=pd.DataFrame(tab,columns=Keyword)
            return True,tab
        
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore")
    
            if len(dlim):
                try:
                    tab=pd.read_csv(tabfile,sep=dlim)
                    return True,tab
                except:
                    pass
    
            #try opening with xlsx
            if tabfile.endswith(".xls") or tabfile.endswith(".xlsx") or xls:
                try:
                    tab = pd.read_excel(tabfile, engine='openpyxl')
                    return True,tab
                except:
                    pass
            
            # dynamic delimiter detection: if file delimiter is different, split using different delimiters until the desired column name is found
            if len(Keyword):
                
                with open(tabfile, "r") as f:
                    tab = pd.DataFrame(f.read().splitlines())
                
                
                if not tab.columns.isin(Keyword).any():
                    delims = [i[0] for i in Counter(
                        [i for i in str(tab.iloc[0]) if not i.isalnum()]).most_common()]
                    for delim in delims:
                        if delim == " ":
                            delim = "\s"
                        try:
                            tab = pd.read_csv(tabfile, sep=delim)
                            if tab.columns.isin(Keyword).any():
                                return True,tab
                        except:
                            pass
    
    else:
        print("Filepath not found!")

    return False,tab


#parse form
def parse_form(form): #chemical formular parser
    e,c,fc="","",[]
    for i in form:
        if i.isupper(): #new entry   
            if e: 
                if not c: c="1"
                fc.append([e,c])
            e,c=i,""         
        elif i.islower(): e+=i
        elif i.isdigit(): c+=i
    if e: 
        if not c: c="1"
        fc.append([e,c])
    
    cdf=pd.DataFrame(fc,columns=["elements","counts"]).set_index("elements").T.astype(int)
    cdf["+"]=form.count("+")
    cdf["-"]=form.count("-")
    return cdf

def getMz(form): #this could be vectorized for speed up in the future
    cdf=parse_form(form)
    return (cdf*mdf.loc[cdf.columns].T.values).sum().sum()

#https://stackoverflow.com/questions/47125697/concatenate-range-arrays-given-start-stop-numbers-in-a-vectorized-way-numpy
def create_ranges(a):
    l = a[:,1] - a[:,0]
    clens = l.cumsum()
    ids = np.ones(clens[-1],dtype=int)
    ids[0] = a[0,0]
    ids[clens[:-1]] = a[1:,0] - a[:-1,1]+1
    return ids.cumsum()

#vectorized find nearest mass
#https://stackoverflow.com/questions/8914491/finding-the-nearest-value-and-return-the-index-of-array-in-python
def find_closest(A, target): #returns index of closest array of A within target
    #A must be sorted
    idx = A.searchsorted(target)
    idx = np.clip(idx, 1, len(A)-1)
    left = A[idx-1]
    right = A[idx]
    idx -= target - left < right - target
    return idx
#%%
def predict_formula(
    input_file           = input_file,
    composition_file     = composition_file,
    MFP_output_folder    = MFP_output_folder,
    MFP_output_filername = MFP_output_filename,
    
    mass_index_file  = mass_index_file,
    mass_defect_file = mass_defect_file,
    mass_table       = mass_table,

    mode     =   mode,                    
    adducts  =   adducts,        
    charges  =   charges,
    ppm      =   ppm,  
    top_candidates  = top_candidates,  
    keep_all        = keep_all, 
    add_formula      = add_formula,
    pre_filter_mass = pre_filter_mass,


    #is there a more elegant way to pass all these arguments?        
        
        ):


    #%% Check filepaths
    if not os.path.exists(composition_file):  composition_file=os.path.normpath(composition_file)
    if not os.path.exists(composition_file):  composition_file=composition_file.replace("[0,","[") #min element counts 
    if not os.path.exists(composition_file):  composition_file=composition_file.replace(".0,","")  #float and int
    if not len(mass_index_file):              mass_index_file=composition_file.replace("_comp.npy","_m2g.npy")
    if not os.path.exists(mass_index_file):   mass_index_file=mass_index_file[:-1]+"z"             #sparse index
    if not len(mass_defect_file):             mass_defect_file=composition_file.replace("_comp.npy","_mass.npy")
    if use_params:                            params_file=composition_file.replace("_comp.npy",".params")
    
    if not os.path.exists(composition_file):  raise ValueError("Composition file "+composition_file+" not found!, can be found on ")
    if not os.path.exists(mass_index_file):   raise ValueError("Mass index file " +mass_index_file +" not found!, run space2cart.py")
    if not os.path.exists(mass_defect_file):  warnings.warn("Mass defect file " +mass_defect_file +" not found!, run space2cart with write_mass=True for better performance")
    

    #% Parse information from composition output
    print("")
    if os.path.exists(params_file) and use_params:
        import json
        print("Parsing database information from params file "+params_file)
        jdf=json.load(open(params_file)) #locals().update(jdf)
        elements,fs,mass_blowup,max_mass=[jdf[i] for i in ["elements","composition","mass_blowup","max_mass"]]
        elements=np.array(elements)
        
    
        
    else:
        if not os.path.exists(params_file): print("Params file "+params_file+" not found!")
        else: print("No params file supplied")
        print("")
        print("Parsing database information from database title.")
        fs=Path(composition_file).stem
        mass_blowup=int(fs.split("_b")[-1].split("max")[0])
        elements=np.array([i.split("]")[-1] for i in fs.split("[")][:-1])
        max_mass=float(composition_file.split("max")[1].split("rdbe")[0])
        
    # Parse RDBE information
    Xrdbe = np.argwhere(elements == "C").flatten()
    Yrdbe = np.argwhere(np.in1d(elements,["H", "F", "Cl", "Br", "I"])).flatten()
    Zrdbe = np.argwhere(np.in1d(elements,["N", "P"])).flatten()
    
    ###% read input masses 
    print("")
    if   type(input_file)==int or type(input_file)==float:                        masses=[input_file]    #single numeric mass
    elif isinstance(input_file,pd.Series) or isinstance(input_file,pd.DataFrame): masses=input_file      #when used as function
    elif isinstance(input_file,np.ndarray):                                       masses=input_file.flatten()
    elif type(input_file)==str: #str -> filepath  
        print("Reading table: "+str(input_file))
        print("")
        if not os.path.exists(input_file):        raise ValueError("Input file  file "+input_file+" not found!, test file is available from https://github.com/hbckleikamp/CartMFP")
        #read table input
        check,masses = read_input(input_file,Keyword="mz") 
        if check: masses=masses["mz"].astype(float).values
        else:
            check,masses = read_input(input_file,Keyword="mass")
            if check: masses=masses["mass"].astype(float)
            else: masses=pd.read_csv(input_file).iloc[:,-1]
            
    else: masses=list(input_file) #assumes numeric format

    if isinstance(masses,pd.Series) or isinstance(masses,pd.DataFrame):  mass_ix=masses.index #keep original index
    else:                                                                mass_ix=np.arange(len(masses)) #make new index
    
    masses,umass_ix=np.unique(masses,return_inverse=True)
    map_umass=pd.DataFrame(np.vstack([mass_ix,umass_ix]).T,columns=["original_index","index"])
 

    if (masses > max_mass).sum():
        print("masses above maximum mass detected!, filtering masses")
        masses = masses[masses <= max_mass]
    print("")
    lm=len(masses)
    
    ### Add adducts ###
    if   "n" in mode or "-" in mode: adducts=[a for a in adducts  if a.count("-")==1]
    elif "p" in mode or "+" in mode: adducts=[a for a in adducts  if a.count("-")!=1]
    else: adducts=[]
    if (not len(adducts)) & ("n" in mode or "-" in mode): adducts=["+-"]
    if (not len(adducts)) & ("p" in mode or "+" in mode): adducts=["--"]

    acomps=[]
    if len(adducts):
        adduct_sign    =np.array([-1  if a[0]=="-" else  1  for a in adducts])
        
        #handle mass based adducts
        adduct_mass=[]
        for a in adducts:
            if a[1] in mdf.index: am=getMz(a[1:]) #check if elemet or digit
            else:
                if   a[-1]=="+": am=float(a[1:-1])-emass
                elif a[-1]=="-": am=float(a[1:-1])+emass
                else           : am=float(a[1:-1])
            adduct_mass.append(am)
        adduct_mass=np.array(adduct_mass)
        
        adf=pd.DataFrame([adducts,adduct_mass]).T
        adf.columns=["adduct","adduct_mass"]
        adf["adduct_mass"]=adf["adduct_mass"]*adduct_sign
        acomps=pd.concat([parse_form(a[1:]) for a in adducts]).fillna(0)*np.array(adduct_sign).reshape(-1,1)
        [acomps.pop(i) for i in ["+","-"] if i in acomps.columns]
        acomps.index=adducts
        
        if "n" in mode or "-" in mode:  print("mode is negative,")
        if "p" in mode or "+" in mode:  print("mode is positive,")
        
        print("adducts used: "+", ".join([i+" ("+str(adduct_mass[ix].round(4)*adduct_sign[ix]) +") " for ix,i in enumerate(adducts)]))
        la=len(adducts)
        mass_df=pd.DataFrame(np.tile(np.vstack([np.arange(lm),masses]).T,(la, 1)),columns=["index","input_mass"])
        mass_df["index"]=mass_df["index"].astype(int)
    
        mass_df["adduct_mass"]=np.repeat(adduct_mass.flatten()*adduct_sign.flatten(),len(masses))
        mass_df["adduct_ix"]=np.repeat(np.arange(len(adducts)),len(masses))
        mass_df["adduct"]=np.repeat(adducts,len(masses))
   
        if len(acomps):
            alements=acomps.columns
            aXrdbe = alements[alements.isin(["C"])].tolist()
            aYrdbe = alements[alements.isin(["H", "F", "Cl", "Br", "I"])].tolist()
            aZrdbe = alements[alements.isin(["N","P"])].tolist()
            ardbe=(acomps[aXrdbe].sum(axis=1)-acomps[aYrdbe].sum(axis=1)/2+acomps[aZrdbe].sum(axis=1)/2).values

    else:
        print("no adducts used!")
        adduct_mass=np.array([0])
        mass_df=pd.DataFrame([np.arange(lm),masses],index=["index","input_mass"]).T
        mass_df["index"]=mass_df["index"].astype(int)
        mass_df["adduct_mass"]=0
        mass_df["adduct_ix"]=0
        mass_df["adduct"]=""

    
    ### Add charges ###
    mass_df["charge"]=[charges]*len(mass_df)
    mass_df=mass_df.explode("charge")
    
    #this implies that multiple charged species should have the same adduct each time, so no hybrid adducts
    # mass_df["mass"]=mass_df["input_mass"]*mass_df["charge"]-mass_df["adduct_mass"]*mass_df["charge"]

    #this allows for one "special" adduct, and the rest of the adducts are  +/- emass
    mass_df["mass"]=mass_df["input_mass"]*mass_df["charge"]-mass_df["adduct_mass"]
    
    
    if mode=="pos": mass_df["mass"]+=emass*(mass_df["charge"]-1) #correct for regular adduct
    if mode=="neg": mass_df["mass"]-=emass*(mass_df["charge"]-1) #correct for regular adduct
    mass_df=mass_df[mass_df["mass"]<max_mass].reset_index(drop=True) #filter on max mass
    m=mass_df["mass"].values.astype(float)
    aix=mass_df.pop("adduct_ix").values


    
    #% MFP
    print("Loading index files")
    print("")
    
    #load mass file
    if mass_index_file[-1]=="z":
        from scipy import sparse 
        emp = sparse.load_npz(mass_index_file) #sparse
    else:
        emp   = np.load(mass_index_file, mmap_mode="r")
        
    comps = np.load(composition_file, mmap_mode="r")
    

    peak_mass = m*mass_blowup
    pmi=np.round(peak_mass,0).astype(np.int64)
    d=np.ceil(pmi/1e6*ppm).astype(np.int64)
    dd=d*2
    


    #compute all
    um=(np.repeat(pmi-d,dd)+(np.arange(dd.sum()) - np.repeat(np.cumsum(dd)-dd, dd))).astype(np.uint64)
    a_ix=np.repeat(np.arange(len(m)),dd) 
    
    #total solutions
    sols=emp[um,1]
    if mass_index_file[-1]=="z": sols=np.array(sols.todense()) #sparse
    tot_sols=pd.DataFrame(np.vstack([a_ix,sols.flatten()]).T,columns=["index","solutions"]).groupby("index")["solutions"].sum()



    if pre_filter_mass:
    
        c=np.array([1]*len(pmi)).astype(int)
        
        
        if mass_index_file[-1]=="z":
            t=np.array([top_candidates]*len(pmi))-np.array(emp[pmi,1].todense()).flatten()
        else:
            t=np.array([top_candidates]*len(pmi))-emp[pmi,1]
        tog=np.array([True]*len(pmi))
        res=[np.vstack([np.arange(len(pmi)),c-1]).T]
        
        while any(tog): 
            
            if mass_index_file[-1]=="z":
                r=np.array(emp[pmi+c,1].todense()).flatten()
                l=np.array(emp[pmi-c,1].todense()).flatten()
            else:
                r=emp[pmi+c,1]
                l=emp[pmi-c,1]
                
            ql=np.argwhere((l>0) & (tog))[:,0]
            qr=np.argwhere((r>0) & (tog))[:,0]
            
            ### Future: add any filters here ### (in case they differ from db)
            
            res.append(np.vstack([ql,-c[ql]]).T)
            res.append(np.vstack([qr,c[qr]]).T)
            c+=1
            t[ql]-=l[ql]
            t[qr]-=r[qr]
            tog[t<0]=False
            tog[c>d]=False
            
        res=np.vstack(res)
        a_ix=res[:,0]
        um=pmi[a_ix]+res[:,1]
   
    else:
        um=(np.repeat(pmi-d,dd)+(np.arange(dd.sum()) - np.repeat(np.cumsum(dd)-dd, dd))).astype(np.uint64)
        a_ix=np.repeat(np.arange(len(m)),dd)
        
    ## Formula prediction 
    x = emp[um].astype(np.int64)#.copy()
    if mass_index_file[-1]=="z": x=np.array(x.todense()) #convert sparse 
    
    mr=np.arange(x[:,1].sum()) - np.repeat(np.cumsum(x[:,1])-x[:,1], x[:,1])
    cq=np.repeat(x[:,0],x[:,1])+mr
    cs= comps[cq]#.copy()
    us=np.repeat(a_ix,x[:,1].astype(int))
    
    if not os.path.exists(mass_defect_file) or force_masscalc:
        ms=np.sum(cs*mdf.loc[elements].values.T,axis=1)                                  #direct float mass calculation
    else:  ms = np.load(mass_defect_file, mmap_mode="r")[cq] #precomputed float mass 
 

    #subtract adduct mass, and composition (is this needed?)
    ms+=mass_df.iloc[us]["adduct_mass"].values.astype(ms.dtype) #compatible with multiple charges

    if len(adducts)>0:
        #this goes wrong
        
        if acomps.values.any():
            missing=list(set(acomps.columns)-set(elements))
            cs=np.hstack([cs,np.zeros([len(cs),len(missing)],cs.dtype)]) #update cs
            elements=np.hstack([elements,missing])                       #update elements
            
            ac=np.hstack([np.argwhere(elements==e)[0] for e in acomps.columns])
            av=acomps.values[aix[us]].astype(int)
         
            #account for negative adducts
            q=~np.any((cs.astype(np.int16)[:,ac]+av)<0,axis=1)
            cs, us, ms, cq = cs[q].astype(comps.dtype), us[q], ms[q], cq[q]

            #add adduct to composition
            cs=cs.astype(np.int16)
            cs[:,ac]+=av[q]

 
    if len(ms):
        print("")

        uu,uc=np.unique(mass_df["index"].iloc[us],return_counts=True)
        q=uc>top_candidates
        lt,mt=uu[~q],uu[q] #less than, more than
        ltq=np.isin(mass_df["index"].iloc[us],lt)
        f_cs,f_us,f_ms=cs[ltq],us[ltq],ms[ltq] #final

              
        if len(mt): 
            print("Picking best "+str(top_candidates)+" candidates.")
            
            #unique formulas
            cqu,cqx,cqi=np.unique(cq,return_index=True,return_inverse=True)
            ms,cs=ms[cqx],cs[cqx]
            
            tree = KDTree(ms.reshape(1,-1).T, leaf_size=200)  #should be unique masses?
            r = tree.query(masses[mt].reshape(1,-1).T, return_distance=False,k=top_candidates).flatten()
            f_cs=np.vstack([f_cs,cs[r]])
            f_ms=np.hstack([f_ms,ms[r]])
            f_us=np.hstack([f_us,np.repeat(mt,top_candidates)])

        res=pd.DataFrame(mass_df.iloc[f_us,:])
        res["pred_mass"]=f_ms
        res["total_solutions"]=tot_sols.loc[res["index"].values].values
        
        #correct for charge
        if "charge" in res.columns:        
            if mode=="neg": res["pred_mass"]+=(res.charge-1)*emass
            if mode=="pos": res["pred_mass"]-=(res.charge-1)*emass
            res["pred_mass"]/=res.charge
    
        res["ppm"]=(res["input_mass"]-res["pred_mass"])/res["pred_mass"]*1e6 #slow
        res["appm"]=res["ppm"].abs()
        res["rdbe"]=(  f_cs[:, Xrdbe].sum(axis=1)*2
                      -f_cs[:, Yrdbe].sum(axis=1)
                      +f_cs[:, Zrdbe].sum(axis=1))/2+1
        
        
        
        if len(acomps): res["rdbe"]+=ardbe[aix[f_us]] #update rdbe with adducts
        res[elements]=f_cs

        #add charge
        if mode=="neg": 
            if "charge" in res.columns: res["-"]=res["charge"]
            else: res["-"]=1    
            elements=np.hstack([elements,["-"]])
        if mode=="pos": 
            if "charge" in res.columns: res["+"]=res["charge"]
            else: res["+"]=1    
            elements=np.hstack([elements,["+"]])
     
        res["index"]=res["index"].astype(int)
        

   

        #add formula
        if add_formula:     # add formula string
            
            q=res.columns.isin(mdf.index)
            hill=res.columns[q].sort_values().tolist()
            hill=[i for i in hill if i not in ["+","-"]]
            col_order=res.columns[~q].tolist()+hill
            if mode=="neg": col_order+=["-"]
            if mode=="pos": col_order+=["+"]
            res=res[col_order]
            
            #with adduct
            ecounts=res[hill]
            e_arr=np.tile(hill,len(res)).reshape(len(res),-1) 
            e_arr=np.where(ecounts==0,"",e_arr)
            eles=ecounts.applymap(str).replace("0","").replace("1","")
            res["formula+adduct"]=["".join(i) for i in np.hstack([e_arr,eles])[:,np.repeat(np.arange(len(hill)),2)+np.tile(np.array([0,len(hill)]),len(hill))]]

            #without adduct
            if len(adducts): 
                ecounts[alements]-=acomps.values[aix[f_us]].astype(int)
                e_arr=np.tile(hill,len(res)).reshape(len(res),-1) 
                e_arr=np.where(ecounts==0,"",e_arr)
                eles=ecounts.applymap(str).replace("0","").replace("1","")
                res["formula"]=["".join(i) for i in np.hstack([e_arr,eles])[:,np.repeat(np.arange(len(hill)),2)+np.tile(np.array([0,len(hill)]),len(hill))]] 

   
        if not pre_filter_mass or acomps.values.any(): res=res[res["appm"]<=ppm]

        res=map_umass.merge(res.drop_duplicates(),on="index",how="inner") 

        if keep_all: 
            missing_index=np.argwhere(~np.in1d(np.arange(lm),np.unique(us))).tolist()
            
            if len(missing_index):
                missing_rows=mass_df[["index","input_mass"]].drop_duplicates().set_index("index").loc[missing_index,:].reset_index()
                missing_rows[res.columns[2:]]=0
                missing_rows["adduct"]=""
                res=pd.concat([res,missing_rows])
        
        res=res.sort_values(by=["original_index","appm"]).reset_index(drop=True)
       
    
#%%
    else:
        rc=['index', 'input_mass', 'adduct', 'adduct_mass', 'charge', 'mass','pred_mass', 'ppm', 'appm']
        res=pd.DataFrame(columns=rc)
    
    return res
    
        
#%% Execute from command line/interpreter

if __name__=="__main__":

    res=predict_formula()
    
    if not os.path.exists(MFP_output_folder): os.makedirs(MFP_output_folder)
    if not len(MFP_output_filename): 
        if os.path.exists(input_file): MFP_output_filename=Path(input_file).stem+"_mfp.tsv"
        else:
            from datetime import datetime
            MFP_output_filename=datetime.now().strftime("%Y-%m-%d_%H-%M-%S")+"_mfp.tsv"
    MFP_outpath=str(Path(MFP_output_folder,MFP_output_filename))
    # res.to_csv(MFP_outpath,sep="\t", index=False) #without pyarrow

    
    #faster writing than pandas
    new_pa_dataframe = pyarrow.Table.from_pandas(res)
    write_options = pc.WriteOptions(delimiter="\t",batch_size=10000)
    pc.write_csv(new_pa_dataframe, MFP_outpath,write_options)

    


