from kuibit.simdir import SimDir
import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import interp1d
from kuibit import argparse_helper as kah
from kuibit.visualize_matplotlib import (
    add_text_to_corner,
    get_figname,
    save_from_dir_filename_ext,
    set_axis_limits_from_args,
    setup_matplotlib,
)
import os

fname = "reftrajectories.csv"
if not os.path.exists(fname):
    print(f"Getting file {fname} from website")
    import requests
    r = requests.get(f"http://einsteintoolkit.org/gallery/bbh/{fname}")
    assert r.status_code == 200
    with open(fname,"w") as fd:
        print(r.content.decode(),file=fd)

data = np.genfromtxt("reftrajectories.csv",delimiter=",")
plt.plot(data[:,0],data[:,1],label='Zenodo')

if __name__ == "__main__":
    setup_matplotlib()

    desc = f"""\
{kah.get_program_name()} plots the trajectories.
"""

    parser = kah.init_argparse(desc)
    kah.add_figure_to_parser(parser, add_limits=True)
    args = kah.get_args(parser)

    sd = SimDir(args.datadir)

    ah1 = sd.horizons.get_apparent_horizon(1).ah
    ah2 = sd.horizons.get_apparent_horizon(2).ah

    t1 = ah1.cctk_iteration.values
    t2 = ah2.cctk_iteration.values

    cen_x1, cen_y1 = ah1.centroid_x, ah1.centroid_y
    cen_x2, cen_y2 = ah2.centroid_x, ah2.centroid_y

    fx2 = interp1d(t2, cen_x2.y, kind='cubic', fill_value='extrapolate')
    fy2 = interp1d(t2, cen_y2.y, kind='cubic', fill_value='extrapolate')

    vx2 = fx2(t1)
    vy2 = fy2(t1)

    plt.plot(cen_x1.y - vx2, cen_y1.y - vy2, "--", label='Recent Run')
    plt.title("Trajectory Comparison")
    plt.xlabel("X")
    plt.ylabel("Y")
    plt.axis('equal')

    plt.legend()
    plt.savefig("traj-diff.png",dpi=300)