#!/usr/bin/python3.12
import argparse
import configparser
import math
import os

def compute_dependent(params):
    model = params["model"]

    # Charge neutrality
    params["charge_B"] = round(params["charge"]/2, 5)
    params["charge_A"] = -2 * params["charge_B"]

    # Relative sigma and epsilon
    params["sigma_A"] = params["sigma"]
    params["epsilon_A"] = params["sigma"]
    sigma_A = params["sigma"]
    epsilon_A = params["epsilon"]
    params["sigma_B"] = round(params["sigma_rel"] * sigma_A, 5)
    params["epsilon_B"] = round(params["epsilon_rel"] * epsilon_A, 5)
    params["d_AB"] = round(params["r_AB"] * sigma_A, 5)
    params["d_AM"] = round(params["r_AM"] * sigma_A, 5)

    # Mass: either specified or computed to give τ ≈ 0.1
    if params["mass"] is None:
        # Let total mass m = 2 * m_A, solve: τ = σ * sqrt(m/ε) ≈ 0.1
        tau_target = 0.05 * 10
        m_total = (tau_target / sigma_A)**2 * epsilon_A
        m_atom = m_total / 2
        params["mass"] = round(m_atom, 5)


    params["d_BB"] = round(2 * params["d_AB"] * math.sin(math.radians(params["angle_ABA"]) / 2), 5)
    params["dummy_ab"] = round(params["d_AM"] / (2 * math.cos(math.radians(params["angle_ABA"]) / 2)), 5)

    return params
    
def clean_params(params):
    if params["model"] != "4point":
        for k in ["r_AM", "d_AM", "d_BB", "dummy_ab"]:
            params.pop(k)
    if params["model"] != "3point" and params["model"] != "4point":
        for k in ["r_AB", "d_AB", "angle_ABA"]:
            params.pop(k)
    return params
    
def format_directory_name(params):
    """Generate a short directory name from active parameters."""
    keys_order = [
        ("sigma", "s"),
        ("epsilon", "e"),
        ("charge", "q"),
        ("sigma_rel", "sr"),
        ("epsilon_rel", "er"),
        ("r_AB", "d"),
        ("angle_ABA", "a"),
        ("r_AM", "m"),
    ]

    parts = [params["model"]]
    parts.append(f"P{params['pressure']}")

    for key, prefix in keys_order:
        if key in params:
            val = params[key]
            parts.append(f"{prefix}{val:.5f}".rstrip("0").rstrip("."))

    parts.append(f"v{params['version']}")
    return "_".join(parts)


def main():
    parser = argparse.ArgumentParser(description="Generate directory and parameter file for A-B molecular systems.")
    parser.add_argument("--version", type=int, default=1, help="Version identifier for parameter sweep")
    parser.add_argument("--model", choices=["atomistic", "3point", "4point"], default='atomistic')
    parser.add_argument("--pressure", type=float, default=1.0)


    # Independent parameters
    parser.add_argument("--sigma", type=float, default=0.3)
    parser.add_argument("--epsilon", type=float, default=0.8)
    parser.add_argument("--q", type=float, default=1.0)

    # Relative parameters
    parser.add_argument("--sigma_rel", type=float, default=0.1)
    parser.add_argument("--epsilon_rel", type=float, default=0.1)

    # Mass (optional, will be calculated if omitted)
    parser.add_argument("--mass", type=float, default=None)

    # Geometry for bonded/virtual_site
    parser.add_argument("--r_AB", type=float, default=0.33334)
    parser.add_argument("--angle_ABA", type=float, default=109.5) # 104.52
    parser.add_argument("--r_AM", type=float, default=0.4)

    # Output
    parser.add_argument("--output", type=str, default="params.ini")

    args = parser.parse_args()

    params = {
        "version": args.version,
        "model": args.model,
        "nmol": 1000,
        "pressure": args.pressure,
        "sigma": args.sigma,
        "sigma_rel": args.sigma_rel,
        "epsilon": args.epsilon,
        "epsilon_rel": args.epsilon_rel,
        "charge": args.q,
        "mass": args.mass,
        "r_AB": args.r_AB,
        "angle_ABA": args.angle_ABA,
        "r_AM": args.r_AM,
    }


    # Compute dependent values
    params = compute_dependent(params)
    params = clean_params(params)
    
    dir_name = format_directory_name(params)
    if os.path.exists(dir_name):
        print(f"Directory {dir_name} already exists. Aborting.")
        return
    os.makedirs(dir_name)

    # Save to file
    config_path = os.path.join(dir_name, args.output)
    with open(config_path, "w") as f:
        config = configparser.ConfigParser()
        config["params"] = params
        config.write(f)

    print(f"Parameter file written to {config_path}")

if __name__ == "__main__":
    main()