import numpy as np
from scipy import optimize
import matplotlib.pyplot as plt
import yaml


x0 = np.array(
    [
        102,
        128,
        143,
        161,
        176,
        191,
        207,
        222,
        237,
        255,
        270,
    ]
)
y0 = np.array(
    [
        1.5,
        1.3,
        1.1,
        1.3,
        1.5,
        1.6,
        2.3,
        2.9,
        4.2,
        5.9,
        7.7,
    ]
)
x1 = np.array(
    [
        100,
        120,
        140,
        160,
        180,
        200,
        220,
        240,
        260,
    ]
)
y1 = np.array(
    [
        1.4,
        1,
        1.4,
        1.5,
        1.7,
        2.9,
        2.8,
        3,
        6.2,
    ]
)

x2 = np.array(
    [
        154,
        172,
        191,
        209,
        228,
        246,
        265,
        283,
        302,
    ]
)

y2 = np.array(
    [
        1.3,
        1.4,
        1.4,
        1.8,
        2,
        2.6,
        3.5,
        4.1,
        6.6,
    ]
)


def get_poly_fit(x, y, remove_first_n=1, remove_last_n=1, n_points=100):
    coeffs = np.polyfit(x, y, 3)
    poly = np.poly1d(coeffs)
    x_start = x[remove_first_n]
    x_end = x[-1 - remove_last_n]
    x_poly = np.linspace(x_start, x_end, n_points)
    y_poly = poly(x_poly)
    return x_poly, y_poly


def get_log_log_threshold(x, y):
    x_log = np.log(x)
    y_log = np.log(y)
    p, e = optimize.curve_fit(
        piecewise_linear, x_log, y_log, p0=[np.mean(x_log), np.mean(y_log), 0.5, 4.0]
    )
    return np.exp(p[0])


def piecewise_linear(x, x0, y0, k1, k2):
    return np.piecewise(
        x, [x < x0], [lambda x: k1 * x + y0 - k1 * x0, lambda x: k2 * x + y0 - k2 * x0]
    )


def main():
    with open('/home/lennartalff/Nextcloud/Private/HRV4Training/lactate/ramp.yaml') as f:
        data = yaml.safe_load(f)
    first_x, first_y = get_poly_fit(x0, y0)
    second_x, second_y = get_poly_fit(x1, y1)
    first_lt1 = get_log_log_threshold(first_x, first_y)
    second_lt1 = get_log_log_threshold(second_x, second_y)
    third_x, third_y = get_poly_fit(x2, y2)
    third_lt1 = get_log_log_threshold(third_x, third_y)
    plt.plot(first_x, first_y, color="C0")
    plt.scatter(x0, y0, color="C0")
    plt.plot(second_x, second_y, color="C1")
    plt.scatter(x1, y1, color="C1")
    plt.plot(third_x, third_y, color="C2")
    plt.scatter(x2, y2, color="C2")
    plt.axvline(first_lt1, color="C0", linestyle="dashed")
    plt.axvline(second_lt1, color="C1", linestyle="dashed")
    plt.axvline(third_lt1, color="C2", linestyle="dashed")
    plt.grid(True, which="major")
    major_ticks = np.arange(0, 8, 1)
    minor_ticks = np.arange(0, 8, 0.1)
    plt.yticks(major_ticks)
    plt.yticks(minor_ticks, minor=True)
    plt.show()


if __name__ == "__main__":
    main()