#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jul 26 12:11:11 2023.

@author: placais
"""
import numpy as np
import matplotlib.pyplot as plt

from palettable.colorbrewer.qualitative import Dark2_8
from cycler import cycler
font = {'family': 'serif', 'size': 25}
plt.rc('font', **font)
plt.rcParams['axes.prop_cycle'] = cycler(color=Dark2_8.mpl_colors)
plt.rcParams["figure.figsize"] = (19.2, 11.24)
plt.rcParams["figure.dpi"] = 100

# Constants
c = 299_792_458

# =============================================================================
# tracewin.out file
# =============================================================================
# Load results
filepath = "tracewin.out"
out = {'z': [],
       'gama-1': [],
       }
n_header = 10

with open(filepath, 'r', encoding='utf-8') as file:
    for i, line in enumerate(file):
        splitted_line = line.split()

        if i == 1:
            freq_mhz = float(splitted_line[1])
            continue

        if i < n_header:
            continue

        out['z'].append(float(splitted_line[1]))
        out['gama-1'].append(float(splitted_line[2]))

# Compute phase
n_points = len(out['z'])
omega = 2. * np.pi * freq_mhz * 1e6

out['z'] = np.array(out['z'])
out['gamma'] = np.array(out['gama-1']) + 1.
out['beta'] = np.sqrt(1. - (out['gamma'])**-2)
out['delta_z'] = np.diff(out['z'])

averaged_beta_over_a_step = (out['beta'][1:] + out['beta'][:-1]) * 0.5
out['delta_phi'] = omega * out['delta_z'] / (averaged_beta_over_a_step * c)

out['phi_abs'] = np.zeros(n_points)
for i in range(1, n_points):
    out['phi_abs'][i] = out['phi_abs'][i - 1] + out['delta_phi'][i - 1]
out['phi_abs'] = np.rad2deg(out['phi_abs'])

# =============================================================================
# data.txt file (Export function from Data tab)
# =============================================================================
filepath = "data.txt"
data = {'z': [0.],
        'phi_abs': [0.]
        }

with open(filepath, 'r', encoding='utf-8') as file:
    for i, line in enumerate(file):
        if line[0] in ['#', '-']:
            continue

        splitted_line = line.split()
        data['z'].append(float(splitted_line[11]))
        data['phi_abs'].append(float(splitted_line[12]))

data['phi_abs'] = np.array(data['phi_abs'])
data['z'] = np.array(data['z']) * 1e-3

# =============================================================================
# Interpolate .out over data.txt
# =============================================================================
out['phi_abs_interp'] = np.interp(x=data['z'],
                                  xp=out['z'],
                                  fp=out['phi_abs'])
delta_phi = out['phi_abs_interp'] - data['phi_abs']

# =============================================================================
# Plot
# =============================================================================
plt.close('all')
fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True)

ax1.grid(True)
ax1.set_ylabel(r'Abs. $\phi$ [deg]')
ax1.plot(out['z'], out['phi_abs'], label='Calculated from .out file')
ax1.plot(data['z'], data['phi_abs'], label='Extracted from Data tab', ls='--')
ax1.legend()

ax2.grid(True)
ax2.set_ylabel(r'$\phi_{out} - \phi_{data}$ [deg]')
ax2.set_xlabel('z [m]')
ax2.plot(data['z'], delta_phi)