import numpy as np
# import pygpufit.gpufit as gf
import pycpufit.cpufit as gf


def test(n_fits, dtype=np.float32):
    params = np.random.rand(n_fits, 2).astype(dtype)

    x = np.random.rand(n_fits, 300).astype(dtype) * 100
    y = params[:, :1] + params[:, 1:] * x

    init_params = np.array([0.1, 0.1], dtype=dtype)
    init_params = np.tile(init_params, (n_fits, 1))

    parameters, states, chi_squares, number_iterations, execution_time = gf.fit(
        data=y,
        weights=None,
        model_id=gf.ModelID.LINEAR_1D,
        initial_parameters=init_params,
        tolerance=1e-8,
        user_info=x
        )
    # Check mean absolute error
    error = np.mean(abs(params - parameters))
    print(f"\n{n_fits:2.2e} fits with dtype: {dtype} -> mean abs error = {error:2.2e}")
    testindex = np.random.randint(0, n_fits)
    error_i = abs(params[testindex] - parameters[testindex])
    print(f"\tchi_squares: {chi_squares[testindex]:2.2e}"
          + f"\titerations: {number_iterations[testindex]:3d}"
          + f"\ttime: {execution_time:3.3f} ms")
    print(f"\tCheck fit result for random index: #{testindex}: "
          + f"\n\t\ttrue:\t{params[testindex][0]:2.3f}, {params[testindex][1]:2.3f}"
          + f"\n\t\tpreds:\t{parameters[testindex][0]:2.3f}, {parameters[testindex][1]:2.3f}"
          + f"\n\t\terror:\t{error_i[0]:2.3e}, {error_i[1]:2.3e}")


for n in range(1, 3):
    test(int(n*1e6))