import time

import numpy as np
import matplotlib.pyplot as plt
from cycler import cycler
from scipy.integrate import quad

from plot_helpers import CO, bold_ticks


ifun = 0
ifun2 = 0


def fig6_16():
    # compute integral using adaptive simpson's rule
    # and plot points used for final computation
    # also: a comparison is made with MATLAB's quad fucntion
    #
    # Inputs:
    #       a = left endpoint of interval
    #       b = right endpoint of interval
    #       n = starting number of subintervals
    #

    global ifun, ifun2

    #### y = 1+(gg/pi)/((x-x0)^2+gg^2);
    #### tolerances selected to get total error of 1e-6
    a = 0
    b = 4
    exact = (np.arctan(6 * np.pi) + 4 * np.pi + np.arctan(2 * np.pi)) / np.pi
    tol = 5e-5
    tolQ = 6e-6  ### error for quad command

    #### y = cos(x.^3).^200;
    #### tolerances selected to get total error of 1e-8
    # a = 0
    # b = 3
    # exact = 0.531594451912990
    # tol = 3e-8
    # tolQ = 2e-9

    #### y = exp(2*x)*sin(6*x);
    #### tolerances selected to get total error of 1e-6
    # a = 0
    # b = 2
    # exact = -8.225724190775976
    # tol = 1e-4
    # tolQ = 1e-5

    # L = max number of allowed levels
    L = 25

    ### ifun = number of function evaluations
    ifun = 0

    tic = time.perf_counter()

    # F = nf x 7 matrix, where rows contain info on subintervals that failed
    # at previous level
    # A = rows are subintervals that are done, with last column containing
    # the integral
    # nf = number of subintervals that failed at previous level
    nf = 1
    m = 0.5 * (a + b)
    fm = f(m)
    fa = f(a)
    fb = f(b)
    s = simp(a, m, b, fa, fm, fb)
    F = np.array([[a, m, b, fa, fm, fb, s]], dtype=float)
    A = np.empty((0, 3), dtype=float)
    Lmax = L

    # this is the loop for the levels
    for il in range(2, L + 1):
        FF = np.empty((0, 7), dtype=float)

        ### use simpson on each remaining subinterval
        for is_ in range(nf):
            mL = 0.5 * (F[is_, 0] + F[is_, 1])
            fmL = f(mL)
            mR = 0.5 * (F[is_, 1] + F[is_, 2])
            fmR = f(mR)
            s2L = simp(F[is_, 0], mL, F[is_, 1], F[is_, 3], fmL, F[is_, 4])
            s2R = simp(F[is_, 1], mR, F[is_, 2], F[is_, 4], fmR, F[is_, 5])
            s2 = s2L + s2R
            Err = abs(s2 - F[is_, 6]) / 15
            tolA = tol * (F[is_, 2] - F[is_, 0]) / (b - a)

            ### determine if the accuracy is attained
            if Err < tolA:
                A = np.vstack(
                    (A, [F[is_, 0], F[is_, 2], (16 * s2 - F[is_, 6]) / 15])
                )
            else:
                FF = np.vstack(
                    (
                        FF,
                        [
                            F[is_, 0],
                            mL,
                            F[is_, 1],
                            F[is_, 3],
                            fmL,
                            F[is_, 4],
                            s2L,
                        ],
                    )
                )
                FF = np.vstack(
                    (
                        FF,
                        [
                            F[is_, 1],
                            mR,
                            F[is_, 2],
                            F[is_, 4],
                            fmR,
                            F[is_, 5],
                            s2R,
                        ],
                    )
                )

        nn = FF.shape

        ### if no remaining subintervals, then leave the for-loop
        nf = nn[0]
        if nf == 0:
            Lmax = il
            break
        F = FF

    ### if max number of levels exceeded, just use what was calcuated
    if nf > 0:
        for is_ in range(nf):
            A = np.vstack((A, [FF[is_, 0], FF[is_, 2], FF[is_, 6]]))

    ### add up values from each subinterval
    ss = A.shape
    print(f"ss = {ss}")
    sum = 0.0
    for is_ in range(ss[0]):
        sum = sum + A[is_, 2]
    print(f"Elapsed time: {time.perf_counter() - tic:.6f} seconds")

    # number of levels and subintevals used
    levels_used = Lmax
    subintervals = ss[0]
    print(f"levels_used = {levels_used}")
    print(f"subintervals = {subintervals}")

    # value for integral, error, and total number of function evaluations
    integral = sum
    error = sum - exact
    total_fns = ifun
    print(f"integral = {integral}")
    print(f"error = {error}")
    print(f"total_fns = {total_fns}")

    # compare with MATLAB's quad routine
    tic = time.perf_counter()
    ifun2 = 0
    quad_value, _ = quad(ff, a, b, epsabs=tolQ, epsrel=tolQ)
    quad_error = quad_value - exact
    quad_fns = ifun2
    print(f"quad_error = {quad_error}")
    # quad(ff, a, b, epsabs=1e-8)
    print(f"quad_fns = {quad_fns}")
    print(f"Elapsed time: {time.perf_counter() - tic:.6f} seconds")

    ### plot function and subintervals
    plt.close("all")
    # get(gcf)
    fig, ax = plt.subplots(figsize=(8.8, 3.3))
    fig.subplots_adjust(left=0.08, right=0.98, bottom=0.22, top=0.98)
    ax.set_prop_cycle(cycler(color=CO))

    nx = 200
    x = np.linspace(a - 0.1, b + 0.1, nx)
    y = np.zeros(nx)
    for ix in range(nx):
        y[ix] = f(x[ix])

    ax.plot(x, y, "r", linewidth=2)

    for is_ in range(ss[0]):
        xx = A[is_, 0]
        ax.plot([xx, xx], [0, f(xx)], "b", linewidth=1.3)
        xx = 0.5 * (A[is_, 0] + A[is_, 1])
        ax.plot([xx, xx], [0, f(xx)], "--b", linewidth=1.1)
    xx = b
    ax.plot([xx, xx], [0, f(xx)], "b", linewidth=1.3)

    ax.set_xlabel("x-axis")
    ax.set_ylabel("y-axis")
    ax.set_xlim(x[0], x[nx - 1])
    ax.set_ylim(0, 3)
    bold_ticks(ax, 14)
    plt.show()

    # plt.savefig('/Users/mark/Desktop/aq8.eps')

    return {
        "integral": integral,
        "error": error,
        "total_fns": total_fns,
        "levels_used": levels_used,
        "subintervals": subintervals,
        "quad_error": quad_error,
        "quad_fns": quad_fns,
    }


def simp(a, m, b, fa, fm, fb):
    g = (b - a) * (fa + 4 * fm + fb) / 6
    return g


def f(x):
    global ifun
    ifun = ifun + 1

    #### 1 + Lorentzian Function
    g = 0.5 * 2 / np.pi
    x0 = 3
    gg = g / 2
    y = 1 + (gg / np.pi) / ((x - x0) ** 2 + gg ** 2)
    # y = np.cos(x**3)**200
    # y = 50/(np.pi*(1+2500*x**2))
    # y = np.exp(2*x)*np.sin(6*x)
    return y


# this is used for the MATLAB quad command
def ff(x):
    global ifun2
    ifun2 = ifun2 + np.size(x)

    #### 1 + Lorentzian Function
    g = 0.5 * 2 / np.pi
    x0 = 3
    gg = g / 2
    y = 1 + (gg / np.pi) / ((x - x0) ** 2 + gg ** 2)
    # y = np.cos(x**3)**200
    # y = 50/(np.pi*(1+2500*x**2))
    # y = np.exp(2*x)*np.sin(6*x)
    return y


if __name__ == "__main__":
    fig6_16()
