"""
This source file is part of the Avogadro project.
This source code is released under the 3-Clause BSD License, (see "LICENSE").
"""

import argparse
import json
import sys


def getOptions():
    userOptions = {
        "labels": {
            "type": "stringList",
            "label": "Atom Label",
            "default": 0,
            "values": [
                "Electrophilicity",
                "Nucleophilicity",
                "f-",
                "f+",
                "f0",
                "s+",
                "s-",
                "s0",
            ],
        },
        "file": {
            "type": "filePath",
            "label": "File Path",
            "default": ""
        }
    }

    opts = {"userOptions": userOptions}
    return opts


def parse_cdft(filePath):
    """
    Parse a CDFT.txt file and return a dictionary mapping property names
    to lists of values (one per atom).
    """
    # Map of property name -> list of values per atom
    properties = {}

    with open(filePath, "r") as f:
        lines = f.readlines()

    current_section = None
    header_columns = []

    for line in lines:
        line = line.strip()
        if not line:
            current_section = None
            header_columns = []
            continue

        # Detect section headers
        if line.startswith("Hirshfeld charges"):
            current_section = "hirshfeld"
            continue
        elif line.startswith("Condensed local electrophilicity"):
            current_section = "electrophilicity"
            continue
        elif line.startswith("Condensed local softness"):
            current_section = "softness"
            continue

        # Parse header line for each section
        if current_section and line.startswith("Atom"):
            # Extract column names from header
            parts = line.split()
            header_columns = parts[1:]  # Skip "Atom"
            # Initialize property lists
            for col in header_columns:
                if col not in properties:
                    properties[col] = []
            continue

        # Parse data lines (start with atom number)
        if current_section and header_columns:
            # Check if line starts with an atom index like "1(O )" or "17(H )"
            parts = line.split()
            if parts and '(' in parts[0]:
                # Extract values (skip the atom identifier)
                values = parts[1:]
                for i, col in enumerate(header_columns):
                    if i < len(values):
                        try:
                            properties[col].append(float(values[i]))
                        except ValueError:
                            properties[col].append(0.0)

    return properties


def generate(opts):
    cjson = opts["cjson"]
    requested_label = opts["labels"]

    # try to open the file
    filePath = opts["file"]

    # Parse the CDFT file
    properties = parse_cdft(filePath)

    if requested_label not in properties:
        # Return unchanged if property not found
        return cjson

    values = properties[requested_label]
    numAtoms = len(cjson["atoms"]["elements"]["number"])

    # Get or create atom labels list
    if "labels" not in cjson["atoms"]:
        cjson["atoms"]["labels"] = [""] * numAtoms
    atomLabels = cjson["atoms"]["labels"]

    # Expand labels if needed
    while len(atomLabels) < numAtoms:
        atomLabels.append("")

    # Assign labels for each atom
    for i, value in enumerate(values):
        if i < numAtoms:
            atomLabels[i] = f"{value:.4f}"

    cjson["atoms"]["labels"] = atomLabels
    return cjson


def runCommand():
    # Read options from stdin
    stdinStr = sys.stdin.read()

    # Parse the JSON strings
    opts = json.loads(stdinStr)

    # Replace this molecule with the new labels
    result = {}
    result["moleculeFormat"] = "cjson"
    result["cjson"] = generate(opts)
    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Read Fukui Labels")
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--print-options", action="store_true")
    parser.add_argument("--run-command", action="store_true")
    parser.add_argument("--display-name", action="store_true")
    parser.add_argument("--menu-path", action="store_true")
    parser.add_argument("--lang", nargs="?", default="en")
    args = vars(parser.parse_args())

    debug = args["debug"]

    if args["display_name"]:
        print("Assign Fukui Labels")
    if args["menu_path"]:
        print("&Extensions")
    if args["print_options"]:
        print(json.dumps(getOptions()))
    elif args["run_command"]:
        print(json.dumps(runCommand()))