import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

def main():
    def F(x1, x2):
        return np.array([x1 + x2, x2 + 0.5])

    def G(x1, x2):
        return np.array([x1, x2])

    def H(x1, x2, t):
        return (1 - t) * G(x1, x2) + t * F(x1, x2)

    # x = np.linspace(-1, 1, 10)
    # y = np.linspace(-1, 1, 10)
    # X1, X2 = np.meshgrid(x, y)
    #
    # fig, axes = plt.subplots(1, 4, figsize=(20, 4))
    #
    # for i, t in enumerate(np.linspace(0, 1, 4)):
    #     H_ = H(X1, X2, t)
    #     axes[i].quiver(X1, X2, *H_, color='r')
    #     axes[i].set_title(f't = {t}')
    #
    # plt.show()

    df_dict = {"x1": [], "x2": [], "t": [], "H1": [], "H2": [], "Hmag": []}

    for x1 in np.linspace(-1, 1, 10):
        for x2 in np.linspace(-1, 1, 10):
            for t in [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]:
                df_dict["x1"].append(x1)
                df_dict["x2"].append(x2)
                df_dict["t"].append(t)
                H_ = H(x1, x2, t)
                df_dict["H1"].append(H_[0])
                df_dict["H2"].append(H_[1])
                df_dict["Hmag"].append(np.sqrt(H_[0]**2 + H_[1]**2))

    df = pd.DataFrame(df_dict).sort_values(by=["t", "x2", "x1"])

    df.to_csv("H.csv", index=False)

if __name__=="__main__":
    main()
