import numpy as np
from sklearn.linear_model import LinearRegression


def get_signal_lr(f_window):
    lr = LinearRegression()
    x = np.array(list(range(len(f_window)))).reshape(-1, 1)
    lr.fit(x, f_window)
    return x, lr.coef_, lr.intercept_


# Buy = 2, Hold = 1, Sell = 0
def signal_code_to_str(signal, coef):
    if signal == 2:
        return f'Buy:{coef:.3f}'
    elif signal == 0:
        return f'Sell:{coef:.3f}'
    else:
        return f'Hold:{coef:.3f}'


def draw_signal_lr(b_window, f_window, signal, axis):
    x, coef, intercept = get_signal_lr(f_window)
    b_window_length = len(b_window)
    axis.scatter(x + b_window_length, f_window, marker='*', color='r')
    axis.plot(x + b_window_length, x * coef + intercept, color='r')
    ymin, ymax = axis.get_ylim()
    axis.text(b_window_length - .38,
              ymax - (ymax - ymin) * .12,
              signal_code_to_str(signal, coef[0]),
              fontsize='x-large',
              color='black')
