import numpy as np
from numba import njit
from scipy.optimize import minimize_scalar


@njit
def get_i(R, v_m):
    return v_m / R



def find_slope(v_m, i_m):
    @njit
    def objective_function(R, v_m, i_m):
        return np.sum(((v_m / R) - i_m)**2)

    res = minimize_scalar(objective_function, args=(v_m, i_m), bounds=(0, 10e6), method='bounded')
    return res.x, get_i(res.x, v_m)



def split_signals(v_m, i_m, i_s):
    mask = np.zeros_like(v_m, dtype=bool)
    for i, v in enumerate(v_m):
        mask[i] = i_s[i] < i_m[i] if v >= 0 else i_s[i] > i_m[i]
    # print(len(mask), len(v_m))
    
    return (v_m[mask], i_m[mask]), (v_m[~mask], i_m[~mask])


def find_rss(v_m, i_m, splits: int =2):
    R, i_s = find_slope(v_m, i_m)
    (v_u, i_u), (v_d, i_d) = split_signals(v_m, i_m, i_s)

    for i in range(splits):
        R_u, i_s_u = find_slope(v_u, i_u)
        (v_u, i_u), (_, _) = split_signals(v_u, i_u, i_s_u)
        R_d, i_s_d = find_slope(v_d, i_d)
        (_, _), (v_d, i_d) = split_signals(v_d, i_d, i_s_d)
    
    Ron = R_u
    Roff = R_d
    print( v_m[0] / i_m[0])
    x0 = np.clip((Roff - v_m[0] / i_m[0]) / (Roff - Ron), 0, 1)
    return Ron, Roff, x0