import numpy as np
from copy import deepcopy


def golden_section(A, C, x, delta, tau_decimal, function, tol):
    iter_num = 0
    score_dict = {}
    diff = 1e9
    opt_score = None
    opt_tau = None
    B = A + delta * np.abs(C - A)
    D = C - delta * np.abs(C - A)
    while np.abs(diff) > tol and iter_num < 200:
        iter_num += 1
        B = np.round(B, tau_decimal)
        D = np.round(D, tau_decimal)
        if B in score_dict:
            score_B = score_dict[B]
        else:
            score_B = function(x, B)
            score_dict[B] = score_B

        if D in score_dict:
            score_D = score_dict[D]
        else:
            score_D = function(x, D)
            score_dict[D] = score_D
        if score_B <= score_D:
            opt_score = score_B
            opt_tau = B
            C = D
            D = B
            B = A + delta * np.abs(C - A)
        else:
            opt_score = score_D
            opt_tau = D
            A = B
            B = D
            D = C - delta * np.abs(C - A)
        diff = score_B - score_D
    return opt_tau, opt_score


def two_d_golden_section(a, c, A, C, delta, function,
                         tol, bw_decimal, tau_decimal, verbose=False):
    b = a + delta * np.abs(c - a)
    d = c - delta * np.abs(c - a)
    opt_bw = None
    opt_tau = None
    diff = 1e9
    score_dict = {}
    iter_num = 0
    while np.abs(diff) > tol and iter_num < 200:
        iter_num += 1
        b = np.round(b, bw_decimal)
        d = np.round(d, bw_decimal)
        if b in score_dict:
            tau_b, score_b = score_dict[b]
        else:
            tau_b, score_b = golden_section(A, C, b, delta, tau_decimal, function, tol)
            score_dict[b] = [tau_b, score_b]
        if d in score_dict:
            tau_d, score_d = score_dict[d]
        else:
            tau_d, score_d = golden_section(A, C, d, delta, tau_decimal, function, tol)
            score_dict[d] = [tau_d, score_d]
        if score_b <= score_d:
            opt_score = score_b
            opt_bw = b
            opt_tau = tau_b
            c = d
            d = b
            b = a + delta * np.abs(c - a)
        else:
            opt_score = score_d
            opt_bw = d
            opt_tau = tau_d
            a = b
            b = d
            d = c - delta * np.abs(c - a)
        diff = score_b - score_d
        if verbose:
            print('----------\n bw:%s, tau:%s, score:%s' % (opt_bw, opt_tau, opt_score))
    return opt_bw, opt_tau


def multi_bws(init_bw, init_tau, X, y, n, k, tol, rss_score,
              gtwr_func, bw_func, sel_func, multi_bw_min, multi_bw_max,
              multi_tau_min, multi_tau_max, verbose=False):
    """
    Multiscale GTWR bandwidth search procedure using iterative GAM back fitting
    """
    if init_bw or init_tau is None:
        bw, tau = sel_func(bw_func(X, y))
    else:
        bw, tau = init_bw, init_tau
    opt_model = gtwr_func(X, y, bw, tau)
    bw_gtwr = bw
    tau_gtwr = tau
    err = opt_model.reside
    betas = opt_model.betas

    XB = np.multiply(betas, X)
    rss = np.sum(err ** 2) if rss_score else None
    scores = []
    bws = np.empty(k)
    taus = np.empty(k)
    BWs = []
    Taus = []
    Betas = None

    for iter_num in range(1, 201):
        new_XB = np.zeros_like(X)
        Betas = np.zeros_like(X)

        for j in range(k):
            temp_y = XB[:, j].reshape((-1, 1))
            temp_y = temp_y + err
            temp_X = X[:, j].reshape((-1, 1))
            bw_class = bw_func(temp_X, temp_y)

            bw, tau = sel_func(bw_class, multi_bw_min[j], multi_bw_max[j],
                               multi_tau_min[j], multi_tau_max[j])

            opt_model = gtwr_func(temp_X, temp_y, bw, tau)
            err = opt_model.reside
            betas = opt_model.betas
            new_XB[:, j] = (betas * temp_X).reshape(-1)
            Betas[:, j] = betas.reshape(-1)
            bws[j] = bw
            taus[j] = tau

        num = np.sum((new_XB - XB) ** 2) / n
        den = np.sum(np.sum(new_XB, axis=1) ** 2)
        score = (num / den) ** 0.5
        XB = new_XB

        if rss_score:
            predy = np.sum(np.multiply(betas, X), axis=1).reshape((-1, 1))
            new_rss = np.sum((y - predy) ** 2)
            score = np.abs((new_rss - rss) / new_rss)
            rss = new_rss
        scores.append(deepcopy(score))
        delta = score
        BWs.append(deepcopy(bws))
        Taus.append(deepcopy(taus))

        if verbose:
            print("Current iteration:", iter_num, ",SOC:", np.round(score, 7))
            print("Bandwidths:", ', '.join([str(bw) for bw in bws]))
            print("taus:", ','.join([str(tau) for tau in taus]))

        if delta < tol:
            break
    opt_bws = BWs[-1]
    opt_tau = Taus[-1]
    return (opt_bws, opt_tau, np.array(BWs), np.array(Taus), np.array(scores),
            Betas, err, bw_gtwr, tau_gtwr)
