#!/usr/bin/env python3
# (C) 2026 Cadence Design Systems, Inc. (Cadence)
# All rights reserved.
# TERMS FOR USE OF SAMPLE CODE The software below ("Sample Code") is
# provided to current licensees or subscribers of Cadence products or
# SaaS offerings (each a "Customer").
# Customer is hereby permitted to use, copy, and modify the Sample Code,
# subject to these terms. Cadence claims no rights to Customer's
# modifications. Modification of Sample Code is at Customer's sole and
# exclusive risk. Sample Code may require Customer to have a then
# current license or subscription to the applicable Cadence offering.
# THE SAMPLE CODE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED.  OPENEYE DISCLAIMS ALL WARRANTIES, INCLUDING, BUT
# NOT LIMITED TO, WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
# PARTICULAR PURPOSE AND NONINFRINGEMENT. In no event shall Cadence be
# liable for any damages or liability in connection with the Sample Code
# or its use.

"""Parse a CSV file and reorder it in the OEChem TK CSV file format."""

import argparse
import csv
import os
import sys
from pathlib import Path

from openeye import oechem
from rich.console import Console
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Parse a CSV file and reorder it in the OEChem TK CSV file format."
__SCRIPT_TOOLKITS__ = ["oechem"]


def parse_options() -> argparse.Namespace:
    """Set up command line options."""
    parser = argparse.ArgumentParser(
        add_help=True,
        formatter_class=RichHelpFormatter,
        description="[yellow]" + __SCRIPT_DESC__ + "[/yellow]",
    )
    parser.add_argument(
        "--in",
        dest="in_file",
        type=str,
        required=True,
        help="Input CSV filename",
    )
    parser.add_argument(
        "--out",
        dest="out_file",
        type=str,
        required=True,
        help="Output CSV filename",
    )
    parser.add_argument(
        "--no-header",
        action="store_true",
        default=False,
        help="Set if the CSV does not have a header",
    )
    parser.add_argument("--help-image", action=HelpPreviewAction)
    return parser.parse_args()


def main() -> int:
    """Parse and reorder a CSV file."""
    args = parse_options()
    console = Console()

    in_file: Path = Path(args.in_file)
    out_file: Path = Path(args.out_file)
    has_header: bool = not args.no_header

    if in_file.suffix != ".csv":
        oechem.OEThrow.Fatal(f"Input {in_file.name} must be CSV format!")

    if out_file.suffix != ".csv":
        oechem.OEThrow.Fatal(f"Output {out_file.name} must be CSV format!")

    header = _get_header(in_file, has_header)

    with in_file.open() as inf, out_file.open("w", newline="") as ofile:
        writer = csv.DictWriter(ofile, fieldnames=header)
        writer.writeheader()
        for row in csv.DictReader(inf):
            writer.writerow(row)

    console.print(f"[green]Reordered CSV written to:[/green] {out_file.name}")
    return os.EX_OK


def _get_header(iname: Path, has_header: bool) -> list[str]:
    """Extract and reorder the CSV header."""
    with iname.open() as f:
        header = list(csv.DictReader(f).fieldnames)

    if has_header:
        if not any("smiles" in h.lower() for h in header):
            oechem.OEThrow.Fatal("Input CSV file must have a SMILES column!")

        # find the SMILES column and move it to the front
        smiles = header.pop(
            next(i for i, h in enumerate(header) if "smiles" in h.lower())
        )
        header.insert(0, smiles)

        # move the TITLE column to be right after the SMILES
        if any("title" in h.lower() for h in header):
            title = header.pop(
                next(i for i, h in enumerate(header) if "title" in h.lower())
            )
            header.insert(1, title)
        else:
            header.insert(1, "TITLE")

    else:
        # find column that can be interpreted as SMILES
        smiles = None
        for i, h in enumerate(header):
            if _is_valid_smiles(h):
                smiles = header.pop(i)
                break
        if smiles is not None:
            header.insert(0, smiles)
        header.insert(1, "")

    return header


def _is_valid_smiles(field: str) -> bool:
    """Check if a field is a valid SMILES string."""
    mol = oechem.OEGraphMol()
    opts = oechem.OEParseSmilesOptions()
    opts.SetQuiet(True)
    return oechem.OEParseSmiles(mol, field, opts)


setattr(main, "__SCRIPT_NAME__", __SCRIPT_NAME__)
setattr(main, "__SCRIPT_DESC__", __SCRIPT_DESC__)
setattr(main, "__SCRIPT_TOOLKITS__", __SCRIPT_TOOLKITS__)

if __name__ == "__main__":
    sys.exit(main())