import scipy as sci
import scipy.optimize
import numpy as np
import math
from matplotlib import pyplot as plt

def amp2(v):
    return np.real(v)**2 + np.imag(v)**2

def pp_fit_against_data(f, args=[]):
    data = d = np.loadtxt("fit_input.txt", delimiter=",")
    data_s, data_t, data_amp2 = d[:,0].astype(np.complex128), d[:,1].astype(np.complex128), d[:,2]

    err_lambda = lambda args: \
        math.fsum((data_amp2 - amp2(f(data_s, data_t, *args)))**2)
    
    return sci.optimize.minimize(err_lambda, args)

def pp_plot_with_f(f, args=[], text=""):
    data = d = np.loadtxt("fit_input.txt", delimiter=",")
    data_s, data_t, data_amp2 = d[:,0], d[:,1], d[:,2]

    fig, ax = plt.subplots(figsize=(7, 7))
        
    for sel_s in reversed(np.sort(np.unique(data_s))):
        sel = (data_s == sel_s)

        t_range = np.arange(np.min(data_t[sel])-0.03, np.max(data_t[sel])+0.03, 0.001)
        s_range = np.zeros_like(t_range) + sel_s

        ax.plot(t_range, amp2(f(s_range.astype(np.complex128), t_range.astype(np.complex128), *args))) # #f(s_range, fit_a, fit_b)**2)
        ax.scatter(np.real(data_t[sel]), data_amp2[sel], label=f"$s={sel_s:.02f}" + "\\ \mathrm{GeV}^2$")

    ax.set_xlabel("$\\frac{t}{\mathrm{GeV}^2}$", fontsize=16)
    ax.set_ylabel("$A^2$", fontsize=16)

    ax.text(0.2, -0.1, text, horizontalalignment='center', verticalalignment='center',
            transform=ax.transAxes, fontsize=10)

    ax.legend()
    plt.show()

