import numpy as np
from sgp4.api import Satrec, SatrecArray
import matplotlib.pyplot as plt
import random
from sgp4.io import twoline2rv
from sgp4.earth_gravity import wgs72
from sgp4.conveniences import check_satrec
import pathlib
import math
try:
    from . import convlib
    from . import skyplot
except ImportError:
    import convlib
    import skyplot

LAT, LON = math.radians(48.775556), math.radians(9.182778)

SCRIPTDIR = pathlib.Path(__file__).parent

CURRENT_JD = 2460874.5 + 2

N = 10000

with open(SCRIPTDIR / ".." / "data" / "tles.txt") as f:
    lines = f.readlines()

satellites = []

tle = []
for line in lines:
    tle.append(line)
    if len(tle) > 2:
        s = tle[1]
        t = tle[2]
        #print(s, t)
        twoline2rv(s, t, wgs72)
        sat = Satrec.twoline2rv(s, t)
        check_satrec(sat)
        assert sat.jdsatepoch > (CURRENT_JD - 14) 
        if not sat.jdsatepoch < (CURRENT_JD):
            print(sat.jdsatepoch)
            continue
        satellites.append(sat)
        tle = []
    if len(satellites) >= 1:
        break

print(len(satellites))

a = SatrecArray(satellites)

jd = np.array([CURRENT_JD for _ in range(N)])
fr = np.linspace(0, 1, N)

e, r, v = a.sgp4(jd, fr)


fig1, ax1 = plt.subplots(subplot_kw=dict(projection='3d'))
fig2, ax2 = plt.subplots()

for sat_e, sat_r, sat_v in zip(e, r, v):
    if not all([error == 0 for error in sat_e]):
        continue
    x = sat_r[:, 0]
    y = sat_r[:, 1]
    z = sat_r[:, 2]
    x, y, z, e, a = convlib.to_local(x, y, z, LAT, LON, jd, fr)
    assert len(x) == N
    col = random.choice(["r", "g", "b", "c", "m", "y", "k"])
    #print(col)
    ax1.plot(x, y, z, c=col)
    ax2.plot(e)


    skyplot.skyplot(plt, e, a)

    total_duration = 0
    entered_fov = None
    for elev, fr in zip(e, fr):
        if math.degrees(elev) > 80 and entered_fov is None:
            entered_fov = fr
        elif math.degrees(elev) < 80 and entered_fov is not None:
            total_duration += fr - entered_fov
            entered_fov = None
    if entered_fov:
        total_duration += fr - entered_fov
    print("Total: ", total_duration)


ax1.set_xlabel("X Pos")
ax1.set_ylabel("Y Pos")
ax1.set_zlabel("Z Pos")


fig1.savefig(SCRIPTDIR / ".." / "output" / "fig6.svg")

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