#!/usr/bin/env python3
import argparse

import pandas as pd
from rdkit import Chem


def sdf_to_tab(vars):
    mols = Chem.SDMolSupplier(vars.inp, sanitize=False)
    df = pd.DataFrame()  # for output

    for n in range(len(mols)):
        if mols[n]:
            d = mols[n].GetPropsAsDict()
            # filter dict for desired props
            if vars.props.strip() == "":  # none specified, return all
                d = {
                    prop: val
                    for (prop, val) in d.items()
                    if not any(x in str(val) for x in ["\n", "\t"])
                }  # remove items containing newlines or tabs
            else:
                d = {
                    prop: val
                    for (prop, val) in d.items()
                    if prop in vars.props.replace(" ", "").split(",")
                }  # remove items not requested via CLI
            if vars.name:
                d["SDFMoleculeName"] = mols[n].GetProp("_Name")
            if vars.smiles:
                d["SMILES"] = Chem.MolToSmiles(mols[n], isomericSmiles=False)
            d["Index"] = int(n)

            df = df.append(d, ignore_index=True)
        else:
            print("Molecule could not be read - skipped.")

    df = df.astype({"Index": int}).set_index("Index")
    sorted_cols = sorted(df.columns.values.tolist())
    df.to_csv(vars.out, sep="\t", header=vars.header, columns=sorted_cols)


def main():
    parser = argparse.ArgumentParser(description="Convert SDF to tabular")
    parser.add_argument("--inp", "-i", help="The input file", required=True)
    parser.add_argument("--out", "-o", help="The output file", required=True)
    parser.add_argument(
        "--props",
        "-p",
        help="Properties to filter (leave blank for all)",
        required=True,
    )
    parser.add_argument(
        "--header",
        "-t",
        action="store_true",
        help="Write property name as the first row.",
    )
    parser.add_argument(
        "--smiles", "-s", action="store_true", help="Include SMILES in output."
    )
    parser.add_argument(
        "--name", "-n", action="store_true", help="Include molecule name in output."
    )
    sdf_to_tab(parser.parse_args())


if __name__ == "__main__":
    main()