# ####################################################################
# #                       Required Libraries                         #
# ####################################################################
# This is the import section that is causing a ModuleNotFoundError.
# Trying the most basic import structure for NanoLanguage.
from NanoLanguage import *
from Analyser import NematicOrderParameter, OrientationFunction
from IO import nlread
import numpy
import pylab
# ####################################################################
# #                        !!! SETTINGS !!!                          #
# # PLEASE CAREFULLY EDIT THIS SECTION TO MATCH YOUR SIMULATION      #
# ####################################################################
# 1. Full name of your results file
results_filename = "PureAndC60_8CB_Resorvuar_Pro_NoE_results.py" # Enter your filename here
# 2. LOCAL indices of atoms defining the molecular axis
#    - In the Viewer, select the FIRST 8CB molecule (e.g., atoms with indices 0 to 46).
#    - Note the Index of the atom at the start of the axis (e.g., 5).
#    - Note the Index of the atom at the end of the axis (e.g., 25).
#    - Subtract 1 from these numbers and write them in the list below (e.g., [4, 24]).
axis_atom_indices_local = [4, 24] # THIS IS AN EXAMPLE, YOU MUST ENTER YOUR OWN VALUES!
# 3. Molecule and Atom Counts (Updated based on new information)
first_8cb_atom_index = 0              # 8CB molecules are now at the beginning
atoms_per_8cb_molecule = 47           # Number of atoms in one 8CB molecule (C21H25N)
total_atom_count = 3820               # New total atom count: (80 * 47) + 60
# 4. Temperature range to be analyzed (should be the same as in your simulation)
temperatures = range(293, 315)
# 5. Filenames for the output plots
s_parameter_plot_filename = "S_vs_Temperature_NEW.png"
angular_dist_plot_filename = "Angular_Distribution_NEW.png"
# ####################################################################
# #                  ANALYSIS SCRIPT (DO NOT MODIFY)                 #
# ####################################################################
# Create empty lists to store the results
calculated_s_parameters = []
first_temp_data = None
last_temp_data = None
# Create the list of all 8CB atom indices (excluding C60)
last_8cb_atom_index = first_8cb_atom_index + (80 * atoms_per_8cb_molecule)
all_8cb_indices = list(range(first_8cb_atom_index, last_8cb_atom_index))
print("+" + "-"*50 + "+")
print("ANALYSIS SCRIPT STARTED")
print(f"File: {results_filename}")
print(f"Total Atoms: {total_atom_count}, 8CB Atoms: {len(all_8cb_indices)}")
print("+" + "-"*50 + "+")
print(f"{'Temperature (K)':<15} | {'Order Parameter (S)':<25}")
print("-" * 52)
# Loop over temperatures
for T in temperatures:
    try:
        # Read the final configuration for each temperature from the HDF5 file
        object_id = f'last_image_3_md_reservoir_temperature_{float(T)}.0_Kelvin'
        configuration = nlread(results_filename, object_id=object_id)[0]
        # Calculate Nematic Order Parameter
        analysis_s = NematicOrderParameter(
            configuration=configuration,
            axis_atom_indices=axis_atom_indices_local,
            molecule_indices=[all_8cb_indices]
        )
        s_value = analysis_s.evaluate()
        calculated_s_parameters.append(s_value)
        print(f"{T:<15.1f} | {s_value:<25.4f}")
        
        # Calculate Orientation Function (only for the first and last temperatures for plotting)
        if T == temperatures[0] or T == temperatures[-1]:
            molecular_axes = []
            for i in range(first_8cb_atom_index, last_8cb_atom_index, atoms_per_8cb_molecule):
                atom1_index = i + axis_atom_indices_local[0]
                atom2_index = i + axis_atom_indices_local[1]
                
                pos1 = configuration.cartesianCoordinates()[atom1_index]
                pos2 = configuration.cartesianCoordinates()[atom2_index]
                vector = pos2 - pos1
                unit_vector = vector / numpy.linalg.norm(vector)
                molecular_axes.append(unit_vector)
            
            analysis_angular = OrientationFunction(vectors=molecular_axes, axis=[0,0,1])
            theta, probability = analysis_angular.evaluate()
            
            if T == temperatures[0]:
                first_temp_data = (theta, probability, T)
            else:
                last_temp_data = (theta, probability, T)
    except Exception as e:
        print(f"ERROR: Could not read or analyze data for T={T}K. Check the object_id. Error: {e}")
        calculated_s_parameters.append(None) # Append None in case of an error
print("\nGenerating plots...")
# Plot 1: S Parameter vs. Temperature
pylab.figure()
pylab.plot(temperatures, calculated_s_parameters, 'o-', label='S Parameter')
pylab.xlabel("Temperature (K)")
pylab.ylabel("Nematic Order Parameter (S)")
pylab.title("Order Parameter vs. Temperature")
pylab.ylim(0, 1)
pylab.grid(True)
pylab.legend()
pylab.savefig(s_parameter_plot_filename)
print(f"-> Saved plot: '{s_parameter_plot_filename}'")
# Plot 2: Angular Distribution
pylab.figure()
if first_temp_data:
    theta_first, prob_first, T_first = first_temp_data
    pylab.plot(theta_first * 180 / numpy.pi, prob_first, label=f'{T_first} K (Low Temperature)')
if last_temp_data:
    theta_last, prob_last, T_last = last_temp_data
    pylab.plot(theta_last * 180 / numpy.pi, prob_last, label=f'{T_last} K (High Temperature)')
pylab.xlabel("Angle with Z-axis (Degrees)")
pylab.ylabel("Probability Density")
pylab.title("Molecular Orientational Distribution")
pylab.grid(True)
pylab.legend()
pylab.savefig(angular_dist_plot_filename)
print(f"-> Saved plot: '{angular_dist_plot_filename}'")
print("\nANALYSIS COMPLETE!")