"""Plots the development of the smectic order parameter over time

This script calculates the smectic order parameter based on the position
of the central atom of each molecule. The order function is caluclated
using the fourier components of the position distribution.

Plots the evolution of this order parameter over time, alongside the time
evolution of the volume fraction. 

Also plots the molecule's centre-of-mass distribution over time (the
formation of a smaectic phase is shown by the emergence of periodic
components in this plot). This transition can also be characterised by 
changes to the variation and mean of the centre of mass distribution.
"""

import numpy as np
import scipy
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.ndimage import uniform_filter1d  # for rolling average
from phase_plot import vol_frac

file_root = "output_T_0.5_time_"
sampling_freq = 5  # only samples one in X files (must be integer) #30
plotting_freq = 5  # only plots on in X of the sampled distributions

plt.rcParams.update({"font.size": 13})  # for figures to go into latex at halfwidth


def order_param(data):
    """Input data in array of length Molecule Number 

    Input data will be rod_positions array which stores input data   
    First index gives molecule number
    Third index gives the component of the position of CoM (x,y,z)

    Method for calculation given by Polson (1997) 10.1103/PhysRevE.56.R6260
    """
    k = 2 * np.pi / 10  # divide by length of molecule
    exp_values = np.exp(1j * k * data)
    return np.abs(np.sum(exp_values)) / N_molecules


# READ PARAMETER VALUES FROM LOG FILE

file_name = "log.lammps"
log_file = open(file_name, "r")
mix_steps_values = []

for i, line in enumerate(log_file):
    """For loop iteratres over every line in file to find the required variables.

    However, because the dump_interval occurs last (in the main body not the preamble)
    we break at this point to avoid reading the whole file unnecessarily."""

    if "variable N" in line:  # to extract independant variable value of N
        for t in line.split():  # separate by whitespace
            try:
                N_molecules = int(t)
            except ValueError:
                pass  # any non-floats in this line are ignored

    if "dump" and "all custom" in line:  # to extract time interval for dump
        for t in line.split():  # separate by whitespace
            try:
                dump_interval = int(t)
                # interval is the last integer in that line
            except ValueError:
                pass
        break  # got all data, break from for loop

    if "variable mix_steps" in line:  # to extract mix_steps (for shrinking)
        run_num = 0  # counts number of runs
        for t in line.split():
            try:
                mix_steps_values.append(int(t))
                run_num += 1
            except ValueError:
                pass

    if "variable run_steps" in line:  # to extract run_steps (for equilibration)
        for t in line.split():
            try:
                equilibrium_time = int(t)  # time per run
            except ValueError:
                pass
        # this comes up first in file so no searching variable here

log_file.close()

tot_mix_time = sum(mix_steps_values)
run_time = tot_mix_time + run_num * equilibrium_time
time_range = range(0, int(run_time), int(dump_interval * sampling_freq))
print(
    "N_molecules, run_time, dump_interval = "
    + str((N_molecules, run_time, dump_interval))
)


# time_range = range(0, 3300000, 100000)  # FOR SIMPLICITY IN TESTING


def CoM_dist(data):
    """Input data in array of size Molecule Number x 3

    Input data will be rod_positions array which stores input data   
    First index gives molecule number
    Second index gives particle number within molecule (1/5/6/10)
    Third index gives the component of the position (x,y,z)

    """
    return data[:, 1]  # return y coordinate


# READ MOLECULE POSITIONS

volume_values = np.full(len(time_range), np.nan)  # new array of NaN
order_param_values = np.full(len(time_range), np.nan)
CoM_mean_values = []
CoM_mean_times = []
for i, time in enumerate(time_range):  # interate over dump files
    data_file = open(file_root + str(time) + ".dump", "r")
    extract_atom_data = False  # start of file doesn't contain particle values
    extract_box_data = False  # start of file doesn't contain box dimension

    box_volume = 1
    rod_positions = np.zeros((N_molecules, 3))
    """Indices are Molecule Number;  Positional coord index"""

    for line in data_file:
        if "ITEM: BOX" in line:  # to start reading volume data
            extract_box_data = True
            extract_atom_data = False
            continue  # don't attempt to read this line

        if "ITEM: ATOMS" in line:  # to start reading particle data
            extract_box_data = False
            extract_atom_data = True
            continue

        if extract_box_data and not extract_atom_data:
            # evaluate before particle values
            # each line gives box max and min in a single axis
            box_dimension = []
            for d in line.split():  # separate by whitespace
                try:
                    box_dimension.append(float(d))
                except ValueError:
                    pass  # any non-floats in this line are ignored
            box_volume *= box_dimension[1] - box_dimension[0]
            # multiply box volume by length of this dimension of box

        if extract_atom_data and not extract_box_data:
            # evaluate after box dimension collection
            # each line is in the form "id mol type x y z vx vy vz"
            particle_values = []
            for t in line.split():  # separate by whitespace
                try:
                    particle_values.append(float(t))
                except ValueError:
                    pass  # any non-floats in this line are ignored

            # Save positional coordatinates of centre particle
            if int(particle_values[2]) == 5:
                rod_positions[int(particle_values[1]) - 1, :] = particle_values[3:6]

    data_file.close()  # close data_file for time step t
    volume_values[i] = box_volume
    order_param_values[i] = order_param(rod_positions[:, 1])

    CoM_data = rod_positions[:, 1]  # select y coordinate
    # kde_data = scipy.stats.gaussian_kde(CoM_data)

    plot_num = 0
    tot_plot_num = len(time_range) // plotting_freq
    colors = plt.cm.cividis(np.linspace(0, 1, tot_plot_num))
    if i % plotting_freq == 0 and time != 0:
        if i == plotting_freq or time >= run_time - (
            dump_interval * sampling_freq * plotting_freq
        ):

            my_kde = sns.kdeplot(
                CoM_data,
                label="T = " + str(int(time)),
                color=colors[i // plotting_freq - 1],
                bw_adjust=0.5,  # adjusts smoothing (default is 1)
                cut=0,
                # clip=(0, 50), # cuts off x limit
                # gridsize=50,  #adjusts points in average (default is 200)
            )
        else:
            my_kde = sns.kdeplot(
                CoM_data,
                color=colors[i // plotting_freq - 1],
                alpha=1,
                bw_adjust=0.5,
                cut=0,
                # clip=(0, 50), # cuts off x limit
                # gridsize=50
            )
            # alpha may be used to adjust transparency
        line = my_kde.lines[plot_num - 1]
        x, y = line.get_data()
        CoM_mean_values.append(max(y) - min(y))
        CoM_mean_times.append(time)  # evaluate order param at time t
        plot_num += 1
        print(max(y), min(y))

    print("T = " + str(time) + "/" + str(run_time))

plt.title("Evolution of CoM distribution over time")
plt.xlabel("Mean CoM")
plt.ylabel("Normalised Frequency")
plt.legend()
plt.savefig("CoM_dist.png")
plt.show()

plt.plot(CoM_mean_times, CoM_mean_values)
plt.plot(
    CoM_mean_times, uniform_filter1d(CoM_mean_values, size=int(10)), linestyle="--",
)
plt.xlabel("Time (arbitrary units)")
plt.ylabel("Mean CoM")
plt.title("Evolution of Mean CoM")
plt.savefig("CoM_mean.png")
plt.show()

fig, ax1 = plt.subplots()

color = "tab:red"
ax1.set_xlabel("Time (arbitrary units)")
ax1.set_ylabel("CoM density variation", color=color)
ax1.plot(CoM_mean_times, uniform_filter1d(CoM_mean_values, size=int(10)), color=color)
ax1.tick_params(axis="y", labelcolor=color)

ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis

color = "tab:blue"
ax2.set_ylabel(
    "Volume Fraction", color=color
)  # we already handled the x-label with ax1
ax2.plot(time_range, vol_frac(volume_values), color=color)
ax2.tick_params(axis="y", labelcolor=color)

plt.title("Evolution of CoM Density")
fig.tight_layout()  # otherwise the right y-label is slightly clipped
plt.savefig("CoMdensity_and_volfrac.png")
plt.show()

plt.plot(time_range, order_param_values)
plt.plot(
    time_range, uniform_filter1d(order_param_values, size=int(10)), linestyle="--",
)
plt.xlabel("Time (arbitrary units)")
plt.ylabel("Order Parameter")
plt.title("Evolution of Smectic Order Parameter")
plt.savefig("smectic_order_evo.png")
plt.show()


fig, ax1 = plt.subplots()

color = "tab:red"
ax1.set_xlabel("Time (arbitrary units)")
ax1.set_ylabel("Smectic Order Parameter", color=color)
ax1.plot(time_range, order_param_values, color=color)
ax1.tick_params(axis="y", labelcolor=color)

ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis

color = "tab:blue"
ax2.set_ylabel(
    "Volume Fraction", color=color
)  # we already handled the x-label with ax1
ax2.plot(time_range, vol_frac(volume_values), color=color)
ax2.tick_params(axis="y", labelcolor=color)

plt.title("Evolution of Order Parameter")
fig.tight_layout()  # otherwise the right y-label is slightly clipped
plt.savefig("s_order_and_volfrac2.png")
plt.show()
