import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import scipy
import scipy.signal
import os
import pathlib
try:
    from . import colorplot
except ImportError:
    import colorplot

CHANDLER_WOBBLE_DAYS = 433
DAYS_PER_YEAR = 365.25
FREQ_ANALYSIS_LIMIT = 800
FREQ_ANALYSIS_SEG_SIZE = 2**11

SCRIPTDIR = pathlib.Path(__file__).parent

# data = pd.read_csv('data\EOP_C01_IAU2000_1846-now.txt', skiprows=13, sep='\s+', header=None) 
data = pd.read_csv(SCRIPTDIR / ".." / "data" / 'EOP_C01_IAU2000_1900-now.txt', skiprows=1, sep=r'\s+', header=None) 

t = data[0]
pos_x = data[1]
pos_y = data[3]

N = int(10 / 0.05)
AVG_CUTOFF = 200
avg_x = np.convolve(pos_x, np.ones(N)/N, mode='valid')
avg_y = np.convolve(pos_y, np.ones(N)/N, mode='valid')
avg_t = t[AVG_CUTOFF-1:]

t_steps = []
for i in range(len(t) - 1):
    step = t[i+1] - t[i]
    t_steps.append(step)

step = t_steps[0]

freq = 1 / (step * DAYS_PER_YEAR)

f_x, p_x = scipy.signal.welch(np.array(pos_x), freq, nperseg=FREQ_ANALYSIS_SEG_SIZE)
f_y, p_y = scipy.signal.welch(np.array(pos_y), freq, nperseg=FREQ_ANALYSIS_SEG_SIZE)

f_x_inv = []
p_x_eff = []
for f, p in zip(f_x, p_x):
    if 1 > FREQ_ANALYSIS_LIMIT * f:
        continue
    f_x_inv.append(1/f)
    p_x_eff.append(p)

f_y_inv = []
p_y_eff = []
for f, p in zip(f_y, p_y):
    if 1 > FREQ_ANALYSIS_LIMIT * f:
        continue
    f_y_inv.append(1/f)
    p_y_eff.append(p)

fig1, ax = plt.subplots()
ax.plot(pos_x, pos_y, alpha=0)
colorplot.colorline(fig1, ax, pos_x, pos_y, alpha=0.5)
print([list(pos_x)[-1]], [list(pos_y)[-1]])
ax.scatter([list(pos_x)[-1]], [list(pos_y)[-1]], zorder=1000, alpha=1, c="r")

ARROW_OFFSET_X = 0.10
ARROW_OFFSET_Y = 0.16
ax.annotate('letzter Wert', xy=(list(pos_x)[-1], list(pos_y)[-1]), xytext=(list(pos_x)[-1] + ARROW_OFFSET_X, list(pos_y)[-1] + ARROW_OFFSET_Y),
            arrowprops=dict(facecolor='black', shrink=0.05),
            )
ax.set_xlabel("X-Koordinate [Bogensekunden]")
ax.set_ylabel("Y-Koordinate [Bogensekunden]")
ax.set_aspect('equal')

# ax.set_title("Pos Overlaid")

fig2, ax = plt.subplots(subplot_kw=dict(projection='3d'))
ax.plot(pos_x, pos_y, t)
ax.plot(avg_x, avg_y, avg_t, c="r")
# ax.set_title("Pos over time")
ax.set_xlabel("X-Abweichung [Bogensekunden]")
ax.set_ylabel("Y-Abweichung [Bogensekunden]")

fig3, ax = plt.subplots()
ax.plot(pos_x, t)
ax.plot(avg_x, avg_t, c="r")
# ax.set_title("Pos over time")
ax.set_xlabel("X-Abweichung [Bogensekunden]")
ax.set_ylabel("Jahr")
ax.set_xlim([-0.65, 0.65])

fig4, ax = plt.subplots()
ax.plot(pos_y, t)
ax.plot(avg_y, avg_t, c="r")
# ax.set_title("Pos over time")
ax.set_xlabel("Y-Abweichung [Bogensekunden]")
ax.set_ylabel("Jahr")
ax.set_xlim([-0.65, 0.65])

fig5, ax = plt.subplots()
ax.plot(f_x_inv, p_x_eff, c="k")
# ax.set_title("Spectrum X Pos")
ax.set_xlabel("Periode [Tage]")
ax.set_ylabel("Intensität")

ax.axvline(DAYS_PER_YEAR, c="b", label="Jährliche Schwingung", zorder=0, alpha=0.5)
ax.axvline(CHANDLER_WOBBLE_DAYS, c="r", label="Chandler Wobble", zorder=0, alpha=0.5)

ax.legend()

fig6, ax = plt.subplots()
ax.plot(f_y_inv, p_y_eff, c="k")
# ax.set_title("Spectrum Y Pos")
ax.set_xlabel("Periode [Tage]")
ax.set_ylabel("Intensität")

ax.axvline(DAYS_PER_YEAR, c="b", label="Jährliche Schwingung", zorder=0, alpha=0.5)
ax.axvline(CHANDLER_WOBBLE_DAYS, c="r", label="Chandler Wobble", zorder=0, alpha=0.5)

ax.legend()

os.makedirs(SCRIPTDIR / ".." / "output", exist_ok=True)
fig1.savefig(SCRIPTDIR / ".." / "output" / "fig2_1.svg")
fig2.savefig(SCRIPTDIR / ".." / "output" / "fig2_2.svg")
fig3.savefig(SCRIPTDIR / ".." / "output" / "fig2_3.svg")
fig4.savefig(SCRIPTDIR / ".." / "output" / "fig2_4.svg")
fig5.savefig(SCRIPTDIR / ".." / "output" / "fig2_5.svg")
fig6.savefig(SCRIPTDIR / ".." / "output" / "fig2_6.svg")


if __name__ == "__main__":
    plt.show()

