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

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,
    ]
)

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():
    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)
    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.axvline(first_lt1, color='C0', linestyle='dashed')
    plt.axvline(second_lt1, color='C1', 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()