import pathlib
import re

import numpy as np
from matplotlib import pyplot as plt

from python.helpers import read_parameter_file

angles = []
tau_cc = []
tau_ss = []
tau2 = []

tevo = [5e-6, 10e-6]


for fit_files in pathlib.Path('.').glob(f'IsotropicAngle/angle=*/Delta/tau=*/ste_fit_*.dat'):
    folder = fit_files.parent

    file_base = fit_files.stem.replace('ste_fit_', '')
    parameter = read_parameter_file(folder / ('ste_' + file_base + '_parameter.txt'))
    angles.append(parameter['angle'])

    with fit_files.open('r') as f:
        # tau of F2 is hidden in the second header line
        for _ in range(2):
            line = f.readline()
        tau2.append(float(re.search('tau=(.+?)\s', line).group(1)))

    fit_values = np.loadtxt(fit_files)
    x = fit_values[:, 0]

    # get indexes for given evolution times
    nearest_idx = [np.searchsorted(x, tt) for tt in tevo]
    tau_cc.append(fit_values[nearest_idx, 1])
    tau_ss.append(fit_values[nearest_idx, 4])


plt.semilogy(angles, tau_cc, '-.')
plt.semilogy(angles, tau_ss, '--')
plt.semilogy(angles, tau2)

np.savetxt(
    'tau_angles.dat',
    np.c_[angles, tau_cc, tau_ss, tau2],
    header=f"Angle dependence of correlation times\nfor evolution times {tevo}\nangle->tau_cc->tauss->tau2"
)

plt.show()