import os
import numpy as np
import csv
# AMS's Qt renders Wayland windows at an integer scale that KWin then shrinks
if os.environ.get("XDG_SESSION_TYPE") == "wayland" and os.environ.get("DISPLAY"):
    os.environ.setdefault("QT_QPA_PLATFORM", "xcb")
import matplotlib.pyplot as plt
from scipy.signal import argrelextrema
from scipy.interpolate import CubicSpline
from scipy.stats import linregress

################################## Functions ######################################

def list_to_matrix(coord_values, n_points):
    if n_points <= 0:
        raise ValueError("The number of path points must be positive.")
    if len(coord_values) % n_points:
        raise ValueError("The number of path coordinates is not divisible by the number of points.")
    n_col = len(coord_values) // n_points  # Calcul du nombre de colonnes
    if n_col == 0:
        raise ValueError("No path coordinates were found.")
    matrix = [coord_values[i:i+n_col] for i in range(0, len(coord_values), n_col)]  # Création de la matrice
    
    return matrix

def input_includes_dispersion(input_text):
    for line in input_text.splitlines():
        content = line.split("#", 1)[0].strip().lower()
        if content and content.split(maxsplit=1)[0] == "dispersion":
            return True
    return False

# extract_diag extracts diagonal elements from upper triangular matrices (IQA2Save)
# extract_off_diag extracts off-diagonal elements from upper triangular matrices (IQA2Save)
# get_indices retrieves original indices (i,j) from off-diagonal elements vector
def extract_diag(Vec, n):
    diag = []
    index = 0
    for i in range(n):
        diag.append(Vec[index])
        index += (n-i)
    return diag

def extract_off_diag(vec, n):
    non_diagonal_elements = []
    k=0
    m = n
    for i in range(int(n*(n+1)/2)):
        if k !=0:
            non_diagonal_elements.append(vec[i])
        k += 1
        if k == m:
            k = 0
            m -= 1
    return non_diagonal_elements

def get_indices(n, k):
    i=1
    j=2
    for m in range(k):
        if j<n:
            j += 1
        else:
            i += 1
            j = i+1
    return i-1, j-1

# Function find_extrema: used to segment the reaction path into small intervals
def find_extrema(values: np.ndarray, seg: int, path_coord: np.ndarray) -> np.ndarray:

    # Find local maxima and minima
    maxima_idx = argrelextrema(values, np.greater)
    minima_idx = argrelextrema(values, np.less)

    # Adjust indices for maxima and minima
    maxima_idx = maxima_idx[0]
    minima_idx = minima_idx[0]

    # Combine the indices of maxima and minima
    extrema_idx = np.concatenate((maxima_idx, minima_idx))

    # Determine points of inflection if choice_seg is not equal to 1
    if seg != 1:
        # IRC points need not be equally spaced along the path.
        first_derivative = np.gradient(values, path_coord)
        second_derivative = np.gradient(first_derivative, path_coord)
        # Find points of sign change in the second derivative
        inflection_points = np.where(np.diff(np.sign(second_derivative)))[0] + 1
        extrema_idx = np.concatenate((extrema_idx, inflection_points))
        # Keep at least three points between each endpoint and the first or last cut
        extrema_idx = extrema_idx[(extrema_idx >= 2) & (extrema_idx <= len(values) - 3)]

    # Add the endpoints of the scan as extrema
    extrema_idx = np.insert(extrema_idx, 0, 0)
    extrema_idx = np.append(extrema_idx, len(values) - 1)

    # Remove duplicates and sort the indices
    extrema_idx = np.unique(extrema_idx)

    return extrema_idx

def regression_bounds(segment_index, extrema_idx):
    start_idx = extrema_idx[segment_index]
    end_idx = extrema_idx[segment_index + 1]
    segment_lengths = np.diff(extrema_idx) + 1
    n_segments = len(segment_lengths)
    excluded_boundaries = np.ones(n_segments - 1, dtype=bool)

    for index, length in enumerate(segment_lengths):
        excluded_left = index > 0 and excluded_boundaries[index - 1]
        excluded_right = index < n_segments - 1 and excluded_boundaries[index]
        if length - excluded_left - excluded_right < 4:
            if index > 0:
                excluded_boundaries[index - 1] = False
            if index < n_segments - 1:
                excluded_boundaries[index] = False

    start = start_idx
    stop = end_idx + 1

    if segment_index > 0 and excluded_boundaries[segment_index - 1]:
        start += 1
    if segment_index < n_segments - 1 and excluded_boundaries[segment_index]:
        stop -= 1

    return start, stop

# term_label returns the atom (intra) or atom pair (inter) label of column j, or None to skip it
def term_label(j, intra):
    if intra:
        return f"{mol.symbols[j]}{j+1}"
    k, l = get_indices(numAtoms, j)
    if k == l:
        return None
    return f"{mol.symbols[k]}{k+1}-{mol.symbols[l]}{l+1}"

# Function process_LR to perform linear regression for each IQA descriptor (so far: EiqaTotal, ECTot and EeeX)
# intra is a boolean :  true (diagonal elements in array), false (off diagonal elements in array)
def process_LR(array, label, all_reg, all_r2, all_Pearson, results, labels_list, intra):
    for j in range(array.shape[1]):
        label2 = term_label(j, intra)
        if label2 is None:
            continue
        slope, intercept, r_value, _, _ = linregress(x, array[:,j])
        r2 = r_value**2
        coeff = slope
        pearson = r_value
        result_row = [label,label2, f"{r2:>10.3f}", f"{pearson:>10.3f}",f"{coeff:>10.3f}"]
        results.append(result_row)
        labels_list.append(label+label2)
        all_reg.append(coeff)
        all_r2.append(r2)
        all_Pearson.append(pearson)

#################################### End of functions ##########################################


# Make sure this script was called correctly
resultsdir = globals().get("resultsdir")
AMSJob = globals().get("AMSJob")
if not resultsdir:
    print("\nUSAGE: $AMSBIN/plams ADFREG.py -v resultsdir=[path to .results folder]")
    print("For instance: $AMSBIN/plams ADFREG.py -v resultsdir=./dimerH2O.results\n") 
    raise SystemExit(2)
if AMSJob is None:
    print("Run this script with $AMSBIN/plams so AMSJob is available.")
    raise SystemExit(2)

# Names of the two output files
resultsdir = os.path.normpath(os.fspath(resultsdir))
if not resultsdir.endswith(".results"):
    print("Error: resultsdir must point to a .results folder.")
    raise SystemExit(2)
results_base = resultsdir[:-len(".results")]
results_filename = results_base + "_ADFREG.txt"
csv_filename = results_base + "_ADFREG.csv"

# Load the job associated with that .results folder
amsrkf = AMSJob.load_external(resultsdir)

# Detect either a PES Scan or an IRC profile?
jobScan = False
jobIRC = False
scan_error = None
try:
    direction = amsrkf.results.get_history_property("IRCDirection")
    jobIRC = direction is not None
except KeyError:
    direction = None

if jobIRC:
    print("\nIRC detected!")
else:
    try:
        # Get the number of scan coordinates
        nPoints = amsrkf.results.readrkf('PESScan','nPoints')
        # Get the scan coordinates and their values
        coordValues = amsrkf.results.readrkf('PESScan','PESCoords')
        nbCoord=len(coordValues)
        scanCoord =  amsrkf.results.readrkf('PESScan','ScanCoord(1)')
        # Get the converged bond energies
        beValues = amsrkf.results.readrkf('PESScan','PES')
        if nPoints < 2:
            raise ValueError(f"PES scan has {nPoints} points; at least two are required.")
        if nbCoord % nPoints:
            raise ValueError("PES coordinate count is not divisible by the number of scan points.")
        if len(beValues) != nPoints:
            raise ValueError(f"PES scan declares {nPoints} points but contains {len(beValues)} energies.")
        # store coordValues in coord_matrix (A column or more)
        ncol = nbCoord // nPoints
        coord_matrix = list_to_matrix(coordValues, nPoints)
        coord_matrix = np.array(coord_matrix)
        jobScan= True
        print("\nPES Scan detected!")
    except KeyError as exc:
        scan_error = exc

if not jobIRC and not jobScan:
    print("\nNo PES Scan or IRC detected: have you saved the data along the reaction path?\n")
    if scan_error is not None:
        print(f"PES scan data could not be read: {scan_error}")
    print("Python script stopped.")
    raise SystemExit(2)

# Get the number of atoms
mol = amsrkf.results.get_input_molecule()
numAtoms = len(mol)

if jobIRC:
    beValues = amsrkf.results.get_history_property("Energy")
    coordValues = amsrkf.results.get_history_property("PathLength")
    converged = amsrkf.results.get_history_property("Converged")
    histories = {
        "Energy": beValues,
        "PathLength": coordValues,
        "IRCDirection": direction,
        "Converged": converged,
    }
    missing_histories = [name for name, values in histories.items() if values is None]
    if missing_histories:
        raise ValueError(f"IRC history is missing: {', '.join(missing_histories)}.")
    history_lengths = {name: len(values) for name, values in histories.items()}
    if len(set(history_lengths.values())) != 1:
        raise ValueError(f"IRC history lengths do not match: {history_lengths}.")
    # Create lists for bond energy values, directions(1:fwd,2:bwd) and coordValues
    new_be = []
    new_direction = []
    new_coordValues = []
    # Scan lists and add new elements if the point is converged
    for e, c, d, is_converged in zip(beValues, coordValues, direction, converged):
        if is_converged:
            new_be.append(e)
            new_direction.append(d)
            if d == 1:
                new_coordValues.append(c)
            else:
                new_coordValues.append(-c)
    # Replace old lists with sorted values (converged points only)
    beValues = new_be
    direction = new_direction
    coordValues = new_coordValues
    nPoints=len(beValues)    
    if nPoints < 2:
        raise ValueError(f"IRC has only {nPoints} converged points; at least two are required.")
    # Swap nforward first elements with nbackward last éléments
    nforward = direction.count(1)
    nbackward = len(beValues) - nforward 
    if nforward == 0 or nbackward == 0:
        raise ValueError("IRC data must contain converged points in both directions.")
    beValues[:nforward], beValues[-nbackward:] = beValues[-nbackward:], beValues[:nforward]
    direction[:nforward], direction[-nbackward:] = direction[-nbackward:], direction[:nforward]
    coordValues[:nforward], coordValues[-nbackward:] = coordValues[-nbackward:], coordValues[:nforward]
    # Invert first elements in temporary lists
    first_be = beValues[:nbackward]
    first_coordValues = coordValues[:nbackward]
    first_be_inverse = first_be[::-1]
    first_coordValues_inverse = first_coordValues[::-1]
    # Now lists are in the right order:
    beValues[:nbackward] = first_be_inverse
    coordValues[:nbackward] = first_coordValues_inverse
    # Store coordValues in coord_matrix for final display (and compatibility with a PES Scan)
    ncol = 1 # 1 column for this matrix (a vector in this case); used for final printing
    coord_matrix = list_to_matrix(coordValues, nPoints)
    coord_matrix = np.array(coord_matrix)

# Create empty lists to store IQA descriptors
iqaKin_matrix = []
iqaeN_matrix = []
iqaeeC_matrix = []
iqaeeXintra_matrix = []
iqaeeXinter_matrix = []
iqaDisp_matrix = []
iqaeeTotal_matrix = []
iqaCTotalintra_matrix = []
iqaCTotalinter_matrix = []
iqaTotal_matrix = []
totalEnergyValues = []
point_files = []

# Read IQA data (and coordinate values + converged energies) 
# A loop over all points in the scan - Some data are unused so far but can be needed later...

if jobIRC: 
    nTS = nbackward # determine index for TS
    part = 1 # before TS for IRC

for i in range(nPoints):
    if jobScan: # Get the name of the file for each PES Scan step
        filename='PESPoint({:d})'.format(i+1) 
    else: # This is a little bit more tricky for an IRC...
        if i == nTS: 
            filename='TS'
            part = 2 # after TS for IRC
        elif part == 1:
            k = nbackward - i
            filename='IRC_Backward_{:d}_converged'.format(k)
            if not os.path.isfile(os.path.join(resultsdir, filename + '.rkf')):
                filename='IRC_Backward_{:d}_converged_sp'.format(k) 
        else:
            k = i - nbackward
            filename='IRC_Forward_{:d}_converged'.format(k)
            if not os.path.isfile(os.path.join(resultsdir, filename + '.rkf')):
                filename='IRC_Forward_{:d}_converged_sp'.format(k)

    point_files.append(filename)
    iqaKin_point = amsrkf.results.readrkf('Properties','IQA kinetic energy',file=filename)            
    iqaKin_matrix.append(iqaKin_point)
    iqaeN_point = amsrkf.results.readrkf('Properties','IQA electron-nucleus',file=filename)
    iqaeN_matrix.append(iqaeN_point)
    iqaeeC_point = amsrkf.results.readrkf('Properties','IQA e-e Coulomb',file=filename)
    iqaeeC_matrix.append(iqaeeC_point)
    iqaeeX_point = amsrkf.results.readrkf('Properties','IQA e-e exchange',file=filename)
    iqaeeXintra_matrix.append(extract_diag(iqaeeX_point,numAtoms))
    iqaeeXinter_matrix.append(extract_off_diag(iqaeeX_point,numAtoms))
    iqaeeTotal_point = amsrkf.results.readrkf('Properties','IQA e-e total',file=filename)
    iqaeeTotal_matrix.append(iqaeeTotal_point)
    iqaCTotal_point = amsrkf.results.readrkf('Properties','IQA Coulomb total',file=filename)
    iqaCTotalintra_matrix.append(extract_diag(iqaCTotal_point,numAtoms))
    iqaCTotalinter_matrix.append(extract_off_diag(iqaCTotal_point,numAtoms))
    iqaTotal_point = amsrkf.results.readrkf('Properties','IQA atom-atom total',file=filename) # Diagonal : intra terms
    iqaTotal_matrix.append(extract_diag(iqaTotal_point,numAtoms))
    totalEnergyValue_point = amsrkf.results.readrkf('Total Energy','Total energy',file=filename)
    totalEnergyValues.append(totalEnergyValue_point)

# End Reading IQA DATA

# Convert lists to np.arrays
iqaKin_matrix = np.array(iqaKin_matrix)
iqaeN_matrix = np.array(iqaeN_matrix)
iqaeeC_matrix = np.array(iqaeeC_matrix)
iqaeeXintra_matrix = np.array(iqaeeXintra_matrix)
iqaeeXinter_matrix = np.array(iqaeeXinter_matrix)
iqaeeTotal_matrix = np.array(iqaeeTotal_matrix)
iqaCTotalintra_matrix = np.array(iqaCTotalintra_matrix)
iqaCTotalinter_matrix = np.array(iqaCTotalinter_matrix)
iqaTotal_matrix = np.array(iqaTotal_matrix)
beValues = np.array(beValues)
coordValues = np.array(coordValues)
totalEnergyValues = np.array(totalEnergyValues)

try:
    user_input = amsrkf.results.readrkf('General', 'user input')
except KeyError:
    user_input = ""
dispersion_enabled = input_includes_dispersion(user_input)
if dispersion_enabled:
    for filename in point_files:
        iqaDisp_point = amsrkf.results.readrkf('Properties', 'IQA pairs disp', file=filename)
        iqaDisp_matrix.append(extract_off_diag(iqaDisp_point, numAtoms))
    iqaDisp_matrix = np.array(iqaDisp_matrix)

# Title on screen + ask the user to choose a segmentation option
print("\n***********************************************************************\n")
print("ADFREG 2.4: A REG ANALYSIS ALONG A PATH COORDINATE (1D PES Scan or IRC)\n")
print("            (October 2026, L. Joubert & V. Tognetti, Univ. Rouen, France )\n")
print("***********************************************************************\n")
print("Use defaults values (recommended for beginners)?")
print("* Original EDA: Intra: no decomposition; Inter: Vcl and VX")
print("* Segmentation: extrema only")
print("* R2 threshold: 0.9")
print("* m(REG) limits: below -0.2 or above +0.2")
if dispersion_enabled:
    print("* Dispersion correction detected in calculation")
print()
# Read user's choice
default = input("Type yes (default) or no: ")

if (default == "yes" or default == "y" or not default):
    scheme=1
    seg=1
    R2_thresh = 0.9
    MREG_thresh = 0.2
    disp = 0
else:
    print("\nDecomposition scheme:")
    print("1) Intra: no decomposition; Inter: Vcl and VX - Original EDA (default)")
    print("2) Intra: Vkin, Vcl, VX; Inter: Vcl and VX")
    # Read user's choice
    scheme = input("Enter your choice (1 or 2): ")
    if not scheme: scheme = 1
    else: scheme = int(scheme) 
    if dispersion_enabled:
        print("\nInclude dispersion contributions (Vdisp)?")
        disp = input("Type yes or no (default): ")
        if (disp == 'no' or not disp): disp = 0
        else: disp = 1
    else:
        disp = 0
    print("\nSegmentation of the path:")
    print("1) Extrema only (default)")
    print("2) Extrema and inflection points")
    # Read user's choice
    seg = input("Enter your choice (1 or 2): ")
    if not seg: seg = 1
    else: seg = int(seg)
    # ask the user to choose a threshold value for R2 coefficients
    print("\nR2 threshold (minimum value):")
    # Read user's choice
    R2_thresh = input("Enter a value (default:0.9): ")
    if not R2_thresh: R2_thresh = 0.9
    else: R2_thresh = float(R2_thresh)
    # Ask for the signed m(REG) cutoff magnitude; both signed limits are applied.
    print("\nm(REG) cutoff (keep values below -cutoff or above +cutoff):")
    MREG_thresh = input("Enter a value (default:0.2): ")
    if not MREG_thresh: MREG_thresh = 0.2
    else: MREG_thresh = float(MREG_thresh)

def passes_reg_filters(r2, mreg):
    return r2 > R2_thresh and (mreg < -MREG_thresh or mreg > MREG_thresh)

# Look for extrema on the path (and inflection point if seg == 2)
path_coord = coord_matrix[:, 0]
extrema_idx = find_extrema(beValues, seg, path_coord)
regression_ranges = [regression_bounds(i, extrema_idx) for i in range(len(extrema_idx) - 1)]
# Per segment: {IQA term: (R2, m(REG))}, used by the IQA plot option
segment_stats = []

with open(results_filename, 'w', encoding='utf-8') as f:
    f.write(f"***********************************************************************\n")
    f.write(f"ADFREG 2.4: A REG ANALYSIS ALONG A PATH COORDINATE (1D PES Scan or IRC)\n")
    f.write(f"            (October 2026, L. Joubert & V. Tognetti, Univ. Rouen, France) \n")
    f.write(f"***********************************************************************\n")
    f.write(f"Original reference for REG Analysis: Thacker, J.C.R., Popelier, P.L.A.\n")
    f.write(f"The ANANKE relative energy gradient (REG) method to automate IQA analysis\n")
    f.write(f"over configurational change. Theor Chem Acc 136, 86 (2017).\n\n")
    if scheme == 1:
        f.write(f"IQA terms: Eintra (total intra-atomic energy)\n")
    else:
        f.write(f"IQA terms: Vkin (intra-atomic kinetic energy)\n")
    f.write(f"           VCl (total coulombic intra- or inter-atomic energy)\n")
    f.write(f"           VX (intra or inter-atomic exchange energy)\n")
    if disp: f.write(f"           Vdisp (dispersion interaction energy)\n")
    f.write(f"\nR2: coefficient of determination / r: Pearson coefficient\n")
    f.write(f"m(REG): relative energy gradient (slope)\n\n")
    if jobScan:
        f.write(f"ADF Job: PES Scan (1D)\n")
        f.write(f"Scan coordinate 1: {scanCoord}\n")
    else:
        f.write(f"ADF Job: IRC\n")
    if scheme == 1:
        f.write(f"Decomposition scheme: Intra: no decomposition; Inter: Vcl and VX\n")
    else:
        f.write(f"Decomposition scheme: Intra: Vkin, Vcl, VX; Inter: Vcl and VX\n")
    if seg == 1:
        f.write(f"Segmentation: extrema only\n")
    else:
        f.write(f"Segmentation: extrema and inflection points\n")
    f.write(f"R2 threshold: {R2_thresh:.2f}\n") 
    f.write(f"m(REG) limits: below {-MREG_thresh:.2f} or above +{MREG_thresh:.2f}\n")
    if dispersion_enabled:
        f.write("Dispersion correction enabled\n")
        if disp:
            f.write("IQA dispersion terms included\n\n")
        else:
            f.write("IQA dispersion terms not included\n\n")
    else:
        f.write("\n")

    # main loop for each segment
    for i in range(len(extrema_idx) - 1):
        start_idx = extrema_idx[i]
        end_idx = extrema_idx[i+1]
        f.write(f"Segment {i+1}\n")
        f.write(f"Path coordinate range: {coordValues[start_idx]:.4f} -> {coordValues[end_idx]:.4f}\n") 
        nbPoints=end_idx+1-start_idx
        f.write(f"{nbPoints} points selected for this segment.\n")
        f.write(f"First point: {start_idx+1}\n")
        f.write(f"Last point: {end_idx+1}\n\n")

        regression_start, regression_stop = regression_ranges[i]
        f.write(f"{regression_stop - regression_start} points used for REG analysis.\n\n")
        x = totalEnergyValues[regression_start:regression_stop]
        y1, y2, y3 = iqaTotal_matrix[regression_start:regression_stop,:], iqaCTotalinter_matrix[regression_start:regression_stop,:], iqaeeXinter_matrix[regression_start:regression_stop,:]
        
        if scheme == 2:
            y4, y5, y6 = iqaKin_matrix[regression_start:regression_stop,:], iqaCTotalintra_matrix[regression_start:regression_stop,:], iqaeeXintra_matrix[regression_start:regression_stop,:]

        if disp:
            y7 = iqaDisp_matrix[regression_start:regression_stop,:]

        # Initialization
        results, all_reg, all_r2, all_Pearson = [], [], [], []
        labels_list1, labels_list2,labels_list3 = [], [], []
        if scheme == 2:
            labels_list4, labels_list5,labels_list6 = [], [], []
        if disp: labels_list7 = []

        # Process linear regressions for selected descriptors
        label='VCl    '
        process_LR(y2, label, all_reg, all_r2, all_Pearson, results, labels_list2, False)
        label='VX     '
        process_LR(y3, label, all_reg, all_r2, all_Pearson, results, labels_list3, False)
        if scheme == 1:
            label='Eintra '
            process_LR(y1, label, all_reg, all_r2, all_Pearson, results, labels_list1, True)
        else:
            label='Vkin'
            process_LR(y4, label, all_reg, all_r2, all_Pearson, results, labels_list4, True)
            label='VCl    '
            process_LR(y5, label, all_reg, all_r2, all_Pearson, results, labels_list5, True)
            label='VX     '
            process_LR(y6, label, all_reg, all_r2, all_Pearson, results, labels_list6, True)
        if disp:
            label='Vdisp  '
            process_LR(y7, label, all_reg, all_r2, all_Pearson, results, labels_list7, False)

        segment_stats.append({row[0] + row[1]: (r2, m) for row, r2, m in zip(results, all_r2, all_reg)})

        # Sort m(REG) for negative values
        negative_indices = [i for i in range(len(all_reg))
                    if all_reg[i] < 0 and passes_reg_filters(all_r2[i], all_reg[i])]
        negative_sorted_indices = sorted(negative_indices, key=lambda k: all_reg[k])

        # Sort m(REG) for positive values
        positive_indices = [i for i in range(len(all_reg))
                    if all_reg[i] >= 0 and passes_reg_filters(all_r2[i], all_reg[i])]
        positive_sorted_indices = sorted(positive_indices, key=lambda k: all_reg[k], reverse=True)

        # Reorganize lines for results
        sorted_results = [results[i] for i in positive_sorted_indices]
        header = ['IQA term', 'Atom(s)', 'R2', 'r', 'm(REG)']
        formatted_results = [header] + [[r[header.index(col)] for col in header] for r in sorted_results]

        # Write positive results to file
        f.write('\n'.join(['{:<10}  {:^10}  {:>10}  {:>10}  {:>10}'.format(*row) for row in formatted_results]) + '\n\n')

        sorted_results = [results[i] for i in negative_sorted_indices]
        formatted_results = [header] + [[r[header.index(col)] for col in header] for r in sorted_results]

        # Write negative results to file
        f.write('\n'.join(['{:<10}  {:^10}  {:>10}  {:>10}  {:>10}'.format(*row) for row in formatted_results]) + '\n\n')

with open(csv_filename, "w", newline="", encoding='utf-8') as file_csv:
    writer = csv.writer(file_csv)
    # Write data for each row
    if scheme == 1:
        # Write column headers
        line = [f"Path coord. {value+1:2d}" for value in range(ncol)] + ['Bond energy', 'Total energy'] + labels_list1 + labels_list2 + labels_list3
        if disp: line = line + labels_list7
        writer.writerow(line)
        for i in range(nPoints):
            line = [f"{value:.6f}" for value in coord_matrix[i,:]] + [f"{beValues[i]:.6f}", f"{totalEnergyValues[i]:.6f}"] + \
            [f"{value:.6f}" for value in iqaTotal_matrix[i,:]] + \
            [f"{value:.6f}" for value in iqaCTotalinter_matrix[i,:]] + \
            [f"{value:.6f}" for value in iqaeeXinter_matrix[i,:]] 
            if disp: line = line + [f"{value:.6f}" for value in iqaDisp_matrix[i,:]]     
            writer.writerow(line)
    else:
        # Write column headers
        line = [f"Path coord. {value+1:2d}" for value in range(ncol)] + ['Bond energy', 'Total energy'] + labels_list4 + labels_list5 + labels_list6 + labels_list2 + labels_list3 
        if disp: line = line + labels_list7
        writer.writerow(line)
        for i in range(nPoints):
            line = [f"{value:.6f}" for value in coord_matrix[i,:]] + [f"{beValues[i]:.6f}", f"{totalEnergyValues[i]:.6f}"] + \
            [f"{value:.6f}" for value in iqaKin_matrix[i,:]] + \
            [f"{value:.6f}" for value in iqaCTotalintra_matrix[i,:]] + \
            [f"{value:.6f}" for value in iqaeeXintra_matrix[i,:]] + \
            [f"{value:.6f}" for value in iqaCTotalinter_matrix[i,:]] + \
            [f"{value:.6f}" for value in iqaeeXinter_matrix[i,:]]
            if disp: line = line + [f"{value:.6f}" for value in iqaDisp_matrix[i,:]]
            writer.writerow(line)

print("\nTwo files written: ")
print(results_filename)
print(csv_filename,"\n")

# Optional plots: one color per segment
plt.rcParams.update({
    "font.size": 10,
    "axes.labelsize": 11,
    "axes.titlesize": 11,
    "legend.fontsize": 9,
    "lines.linewidth": 1.0,
    "lines.markersize": 4,
    "axes.grid": True,
    "grid.alpha": 0.3,
})
plot_title = os.path.basename(results_base)
path_coord_label = "IRC path length (Angstrom)" if jobIRC else "Path coordinate (a.u.)"
n_segments = len(extrema_idx) - 1

# Full path values of every IQA term, keyed like segment_stats
descriptors = [(iqaCTotalinter_matrix, 'VCl    ', False), (iqaeeXinter_matrix, 'VX     ', False)]
if scheme == 1:
    descriptors.append((iqaTotal_matrix, 'Eintra ', True))
else:
    descriptors += [(iqaKin_matrix, 'Vkin', True), (iqaCTotalintra_matrix, 'VCl    ', True),
                    (iqaeeXintra_matrix, 'VX     ', True)]
if disp:
    descriptors.append((iqaDisp_matrix, 'Vdisp  ', False))
iqa_series = {}
for matrix, label, intra in descriptors:
    for j in range(matrix.shape[1]):
        label2 = term_label(j, intra)
        if label2 is not None:
            iqa_series[label + label2] = matrix[:, j]

def plot_segments(values, ylabel, segment_ids):
    # CubicSpline needs strictly increasing abscissae
    order = np.argsort(path_coord)
    spline = CubicSpline(path_coord[order], values[order])
    fig, ax = plt.subplots()
    for i in segment_ids:
        start_idx, end_idx = extrema_idx[i], extrema_idx[i+1]
        color = f"C{i % 10}"
        x_seg = path_coord[start_idx:end_idx+1]
        y_seg = values[start_idx:end_idx+1]
        x_fit = np.linspace(x_seg.min(), x_seg.max(), 200)
        ax.plot(x_fit, spline(x_fit), color=color, label=f"Segment {i+1}")
        ax.plot(x_seg, y_seg, "o", color=color)
    ax.set_xlabel(path_coord_label)
    ax.set_ylabel(ylabel)
    ax.set_title(plot_title)
    ax.ticklabel_format(axis="y", style="plain", useOffset=False)
    ax.legend()
    fig.tight_layout()
    plt.show()

def ask_index(prompt, n_choices, first=1):
    answer = input(prompt).strip()
    if answer.isdigit() and first <= int(answer) <= n_choices:
        return int(answer)
    print("Invalid choice.\n")
    return None

def plot_vs_energy(term, energies, xlabel, segment_ids):
    values = iqa_series[term]
    # Energy is not monotonic along the path, so both axes are splined against the path coordinate
    order = np.argsort(path_coord)
    spline_x = CubicSpline(path_coord[order], energies[order])
    spline_y = CubicSpline(path_coord[order], values[order])
    fig, ax = plt.subplots()
    for i in segment_ids:
        start_idx, end_idx = extrema_idx[i], extrema_idx[i+1]
        color = f"C{i % 10}"
        x_seg = energies[start_idx:end_idx+1]
        y_seg = values[start_idx:end_idx+1]
        s_fit = np.linspace(path_coord[start_idx], path_coord[end_idx], 200)
        ax.plot(spline_x(s_fit), spline_y(s_fit), color=color)
        ax.plot(x_seg, y_seg, "o", color=color)
        regression_start, regression_stop = regression_ranges[i]
        regression_x = energies[regression_start:regression_stop]
        regression_y = values[regression_start:regression_stop]
        fit = linregress(regression_x, regression_y)
        x_line = np.array([regression_x.min(), regression_x.max()])
        ax.plot(x_line, fit.slope * x_line + fit.intercept, "--", color=color,
                label=f"Segment {i+1}: y = {fit.slope:.3f} x {fit.intercept:+.3f}, R\u00b2 = {fit.rvalue**2:.3f}")
    ax.set_xlabel(xlabel)
    ax.set_ylabel(f"{' '.join(term.split())} (a.u.)")
    ax.set_title(plot_title)
    ax.ticklabel_format(axis="both", style="plain", useOffset=False)
    ax.locator_params(axis="x", nbins=5)
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.15), frameon=False)
    fig.tight_layout()
    plt.show()

def select_term():
    print("\nSegments:")
    print("0) All segments")
    for i in range(n_segments):
        print(f"{i+1}) Segment {i+1}: path coordinate {path_coord[extrema_idx[i]]:.4f} -> {path_coord[extrema_idx[i+1]]:.4f}")
    seg_choice = ask_index(f"Enter your choice (0-{n_segments}): ", n_segments, first=0)
    if seg_choice is None:
        return None
    segment_ids = list(range(n_segments)) if seg_choice == 0 else [seg_choice - 1]

    # A term must pass both thresholds on every selected segment.
    terms = [term for term in segment_stats[segment_ids[0]]
             if all(passes_reg_filters(*segment_stats[i][term]) for i in segment_ids)]
    if not terms:
        print(f"\nNo IQA term with R2 > {R2_thresh} and m(REG) outside [{-MREG_thresh}, {MREG_thresh}] on the selected segment(s).\n")
        return None
    terms.sort(key=lambda term: -min(segment_stats[i][term][0] for i in segment_ids))

    print(f"\nIQA terms with R2 > {R2_thresh} and m(REG) outside [{-MREG_thresh}, {MREG_thresh}]:")
    for n, term in enumerate(terms, start=1):
        stats = "  ".join(f"S{i+1}: R2={segment_stats[i][term][0]:.3f} m(REG)={segment_stats[i][term][1]:.3f}"
                          for i in segment_ids)
        print(f"{n:3d}) {' '.join(term.split()):<16} {stats}")
    term_choice = ask_index(f"Enter the term to plot (1-{len(terms)}): ", len(terms))
    if term_choice is None:
        return None
    return segment_ids, terms[term_choice - 1]

plot_choices = {
    "1": (totalEnergyValues, "Total energy (a.u.)"),
    "2": (beValues, "Bond energy (a.u.)"),
}
regression_choices = {
    "4": (totalEnergyValues, "Total energy (a.u.)"),
    "5": (beValues, "Bond energy (a.u.)"),
}
while True:
    print("Plot options:")
    print("1) Total energy vs path coordinate")
    print("2) Bond energy vs path coordinate")
    print("3) IQA term vs path coordinate")
    print("4) IQA term vs total energy")
    print("5) IQA term vs bond energy")
    print("6) Quit")
    choice = input("Enter your choice (1-6): ").strip()
    if choice == "6" or not choice:
        break
    if choice in plot_choices:
        energies, ylabel = plot_choices[choice]
        plot_segments(energies, ylabel, range(n_segments))
        continue
    if choice != "3" and choice not in regression_choices:
        print("Invalid choice.\n")
        continue

    selection = select_term()
    if selection is None:
        continue
    segment_ids, term = selection
    if choice == "3":
        plot_segments(iqa_series[term], f"{' '.join(term.split())} (a.u.)", segment_ids)
    else:
        energies, xlabel = regression_choices[choice]
        plot_vs_energy(term, energies, xlabel, segment_ids)