import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

class Body(object):
    def __init__(self, name):
        self.name = name

names  = ['Sun', 'Earth', 'Mars', 'EMM_Hope', 'Tianwen_1', 'Mars2020']

bodies = []
for name in names:
    fname = name + ' horizons_results.txt'
    with open(fname, 'r') as infile:
        lines = infile.read().splitlines()

    iSOE = [i for i, line in enumerate(lines) if "$$SOE" in line][0]
    iEOE = [i for i, line in enumerate(lines) if "$$EOE" in line][0]

    print(iSOE, iEOE, lines[iSOE], lines[iEOE])
    lines = [line.split(',') for line in lines[iSOE+1:iEOE]]
    JD  = np.array([float(line[0]) for line in lines])
    pos = np.array([[float(item) for item in line[2:5]] for line in lines])
    vel = np.array([[float(item) for item in line[5:8]] for line in lines])
    body = Body(name)
    body.lines = lines
    body.JD = JD
    body.pos = pos.T.copy()
    body.vel = vel.T.copy()
    bodies.append(body)

Sun, Earth, Mars, EMM_Hope, Tianwen_1, Mars2020 = bodies

JDmin, JDmax = min([b.JD[0] for b in bodies]), max([b.JD[-1] for b in bodies])

for b in bodies:
    # https://numpy.org/devdocs/reference/generated/numpy.pad.html
    pad = [ [0, 0], [int(round(b.JD[0] - JDmin)), int(round(JDmax - b.JD[-1]))]]
    b.pos = np.pad(b.pos, pad, constant_values=np.nan)
    b.vel = np.pad(b.vel, pad, constant_values=np.nan)

JD = Earth.JD
JD_2021 = [float(line[0]) for line in Earth.lines if '2021-Jan' in line[1]][0]
days_rel_2021 = Earth.JD - JD_2021

if True:
    million = 1E+06
    days_min =days_rel_2021[0]
    days_max = min(days_rel_2021[-1], 60)
    days_limits = (days_min, days_max)

    fig = plt.figure()

    ax1 = fig.add_subplot(3, 1, 1)

    r_Mars_Earth = np.sqrt(((Mars.pos - Earth.pos)**2).sum(axis=0))
    l, = ax1.plot(days_rel_2021, r_Mars_Earth/million, label='Mars/Earth')
    ax1.set_ylabel('million km')
    ax1.set_xlim(*days_limits)
    ax1.set_ylim(0, 200)
    ax1.legend()

    ax2 = fig.add_subplot(3, 1, 2)

    r_EMM_Hope_Tianwen_1 = np.sqrt(((EMM_Hope.pos - Tianwen_1.pos)**2).sum(axis=0))
    l, = ax2.plot(days_rel_2021, r_EMM_Hope_Tianwen_1/million, label='EMM_Hope/Tianwen-1')

    r_Tianwen_1_Mars2020 = np.sqrt(((Tianwen_1.pos - Mars2020.pos)**2).sum(axis=0))
    l, = ax2.plot(days_rel_2021, r_Tianwen_1_Mars2020/million, label='Tianwen-1/Mars2020')

    r_Mars2020_EMM_Hope = np.sqrt(((Mars2020.pos - EMM_Hope.pos)**2).sum(axis=0))
    l, = ax2.plot(days_rel_2021, r_Mars2020_EMM_Hope/million, label='Mars2020/EMM_Hope')

    ax2.set_ylabel('million km')
    ax2.set_xlim(*days_limits)
    ax2.legend()

    ax3 = fig.add_subplot(3, 1, 3)

    r_EMM_Earth = np.sqrt(((EMM_Hope.pos - Earth.pos)**2).sum(axis=0))
    l, = ax3.plot(days_rel_2021, r_EMM_Earth/million, label='EMM_Hope')
    r_EMM_Mars = np.sqrt(((EMM_Hope.pos - Mars.pos)**2).sum(axis=0))
    ax3.plot(days_rel_2021, r_EMM_Mars/million, color=l.get_color())

    r_Tianwen_1_Earth = np.sqrt(((Tianwen_1.pos - Earth.pos)**2).sum(axis=0))
    l, = ax3.plot(days_rel_2021, r_Tianwen_1_Earth/million, label='Tianwen_1')
    r_Tianwen_1_Mars = np.sqrt(((Tianwen_1.pos - Mars.pos)**2).sum(axis=0))
    ax3.plot(days_rel_2021, r_Tianwen_1_Mars/million, color=l.get_color())

    r_Mars2020_Earth = np.sqrt(((Mars2020.pos - Earth.pos)**2).sum(axis=0))
    l, = ax3.plot(days_rel_2021, r_Mars2020_Earth/million, label='Mars_2020')
    r_Mars2020_Mars = np.sqrt(((Mars2020.pos - Mars.pos)**2).sum(axis=0))
    ax3.plot(days_rel_2021, r_Mars2020_Mars/million, color=l.get_color())
    ax3.legend()

    ax3.set_xlim(*days_limits)
    ax3.set_ylim(0, 10)
    ax3.set_ylabel('million km')
    ax3.set_xlabel('days relative to 00:00 01-Jan-2021 UTC')

    plt.show()