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

def round_dynamic(x, rel_tol=1e-3, max_decimals=6, max_negative_decimals=3):
    """
    Dynamically round x so that the relative error is within rel_tol.
    Supports negative decimals for large values.
    """
    if x == 0:
        return 0.0

    # Determine the required number of decimals (can be negative)
    decimals = -int(math.floor(math.log10(rel_tol * abs(x))))

    # Clamp within allowed range
    decimals = max(-max_negative_decimals, min(max_decimals, decimals))

    if decimals >= 0:
        return round(x, decimals)
    else:
        # Round to nearest 10^(-decimals)
        factor = 10 ** (-decimals)
        return round(x / factor) * factor

def generate_temperature_list_dynamic(dT=10, Tmin=10, Tmax=5000, rel_tol=5e-3):
    r = (300 - dT) / 300.0

    temps_down = []
    T = 300
    while T > Tmin:
        T *= r
        if T < Tmin:
            break
        temps_down.append(T)

    inv_r = 1.0 / r
    temps_up = []
    T = 300
    while T < Tmax:
        T *= inv_r
        if T > Tmax:
            break
        temps_up.append(T)

    all_T = list(reversed(temps_down)) + [300.0] + temps_up

    # Apply dynamic rounding
    rounded_T = [round_dynamic(t, rel_tol=rel_tol) for t in all_T]

    return rounded_T
    
def generate_temperature_list(dT=10, Tmin=10, Tmax=5000, rel_tol=5e-3):
    dT = (Tmax - Tmin) / 20
    all_T = [i*dT+Tmin for i in range(21)]

    # Apply dynamic rounding
    rounded_T = [round_dynamic(t, rel_tol=rel_tol) for t in all_T]

    return rounded_T

def generate_times_list(npoints, start_time, end_time):
    if npoints == 1:
        return [0.0]
    duration = end_time - start_time
    interval = duration / (npoints - 1)
    return [round(start_time + i * interval, 3) for i in range(npoints)]

def main():
    parser = argparse.ArgumentParser(description="Generate annealing times and temperatures for GROMACS.")
    parser.add_argument("-d", "--dT", type=float, default=20.0, help="Temperature step size")
    parser.add_argument("-l", "--Tmin", type=float, default=10.0, help="Minimum temperature")
    parser.add_argument("-u", "--Tmax", type=float, default=5000.0, help="Maximum temperature")
    parser.add_argument("-r", "--rel_tol", type=float, default=5e-3, help="Relative tolerance for rounding")
    parser.add_argument("-s", "--start_time", type=float, default=0, help="Start time in fs")
    parser.add_argument("-e", "--end_time", type=float, required=True, help="End time in fs")
    parser.add_argument("-c", "--cooling", action='store_true', help="Do cooling instead of heating")

    args = parser.parse_args()

    #temps = generate_temperature_list_dynamic(args.dT, args.Tmin, args.Tmax, args.rel_tol)
    temps = generate_temperature_list(args.dT, args.Tmin, args.Tmax, args.rel_tol)
    if args.cooling:
        temps = list(reversed(temps))

    if args.start_time > 0:
        temps = [temps[0]] + temps  # hold first temperature
        times = [0.0] + generate_times_list(len(temps) - 1, args.start_time, args.end_time)
    else:
        times = generate_times_list(len(temps), args.start_time, args.end_time)

    print(f"annealing-npoints        = {len(times)}")
    print("annealing-time           = " + "  ".join(f"{t:.3f}" for t in times))
    print("annealing-temp           = " + "  ".join(f"{t:.2f}" for t in temps))

if __name__ == "__main__":
    main()