import qiskit_tools as qt
import numpy as np
import matplotlib.pyplot as plt
import matplotlib

matplotlib.rcParams['mathtext.fontset'] = 'stix'
matplotlib.rcParams['font.family'] = 'STIXGeneral'
width=0.75
color='black'
fontsize=28
ticksize=22
figsize=(10,8)

def f_x(x):
    return np.arctan(x)

phase = True
nx = 5
nintx = 3
npresx = nx - nintx
if phase:
    npresx -= 1

xmin = -2.**(nintx)
xmax = 2.**(nintx) - 2.**(-npresx)

xs = np.linspace(xmin,xmax,2**nx)

print(xmin, xmax)

print(npresx,qt.get_npres(xs))

m = 3
bounds = np.linspace(xmin, xmax, (2**m))

norder = 1
coeffs = qt.get_bound_coeffs(f_x, bounds, norder, reterr=False).T
ys = qt.piecewise_poly(xs, coeffs.T, bounds)

fig = plt.figure(figsize=figsize)
ax = fig.gca()

for bound in bounds:
    ax.axvline(bound, ls=':', color='black')

ax.plot(xs, f_x(xs), color='black');
ax.plot(xs, ys, color='black', ls='--', lw=3);
ax.set_xlabel(r'$x$', fontsize=fontsize);
ax.set_ylabel(r'$f(x)$', fontsize=fontsize);
ax.set_xlim(xmin,xmax);
ax.tick_params(axis='both', labelsize=ticksize);

# The number of integer bits of A1, A0 and the result
nint0 = qt.get_nint(coeffs[0])
nint1 = qt.get_nint(coeffs[1])
nint = nintx + nint0

# The number of precision bits of A1, A0 and the result
npres0 = qt.get_npres(coeffs[0])
npres1 = qt.get_npres(coeffs[1])
npres = npresx + npres0

# The total number of bits of A1, A0 and the result
n0 = npres0 + nint0 + 1
n1 = npres1 + nint1 + 1
n = npres + nint + 1

print('A0 bits:', n0)
print('A1 bits:', n1)
print('Result bits:', n)

nsig0, nsig1 = 5, 5
coeffs[0] = qt.round_sig(coeffs[0], nsig0)
coeffs[1] = qt.round_sig(coeffs[1], nsig1)

# The number of integer bits of A1, A0 and the result
nint0 = qt.get_nint(coeffs[0])
nint1 = qt.get_nint(coeffs[1])
nint = nintx + nint0

# The number of precision bits of A1, A0 and the result
npres0 = qt.get_npres(coeffs[0])
npres1 = qt.get_npres(coeffs[1])
npres = npresx + npres0

# The total number of bits of A1, A0 and the result
n0 = npres0 + nint0 + 1
n1 = npres1 + nint1 + 1
n = npres + nint + 1

print('A0 bits:', n0)
print('A1 bits:', n1)
print('Result bits:', n)

npres0, npres1 = 4, 4

n0 = npres0 + nint0 + 1
n1 = npres1 + nint1 + 1

rcoeffs = []
for coeff in coeffs[0]:
    # Convert coefficient to binary string of new length n0, then calculate the corresponding decimal value
    rcoeff = qt.bin_to_dec(qt.my_binary_repr(coeff, n0, nint=nint0, phase=phase), nint=nint0, phase=phase)
    rcoeffs.append(rcoeff)
    
coeffs[0] = np.array(rcoeffs)

# Calculate the differences between the f(x) and A1x over all x
fdifs = f_x(xs) - qt.piecewise_poly(xs, np.array([coeffs[0],np.zeros(len(coeffs[1]))]).T, bounds)
coeffs_ = []
bounds__ = bounds
bounds__[-1] = np.inf
for i in np.arange(len(bounds__))[:-1]:
    # Calculate the mean differences in each domain to be the new bias A0
    coeffs_.append(np.mean(fdifs[np.greater_equal(xs,bounds__[i])&np.greater(bounds__[i+1],xs)]))

coeffs[1] = np.array(coeffs_)    

nint1 = qt.get_nint(coeffs[1])

print(nint1)

rcoeffs = []
for coeff in coeffs[1]:
    # Convert coefficient to binary string of new length n1, then calculate the corresponding decimal value
    rcoeff = qt.bin_to_dec(qt.my_binary_repr(coeff, n1, nint=nint1, phase=phase), nint=nint1, phase=phase)
    rcoeffs.append(rcoeff)

coeffs[1] = np.array(rcoeffs)

# The number of integer bits of A1, A0 and the result
nint0 = qt.get_nint(coeffs[0])
nint1 = qt.get_nint(coeffs[1])
nint = nintx + nint0

# The number of precision bits of the result
npres = npresx + npres0

# The total number of bits of A1, A0 and the result
n0 = npres0 + nint0 + 1
n1 = npres1 + nint1 + 1
n = npres + nint + 1

print('A0 bits:', n0)
print('A1 bits:', n1)
print('Result bits:', n)

ys_rnd = qt.piecewise_poly(xs, coeffs.T, bounds)

fig = plt.figure(figsize=figsize)
ax = fig.gca()

for bound in bounds:
    ax.axvline(bound, ls=':', color='black')
    
ax.plot(xs, f_x(xs), color='black');
ax.plot(xs, ys, color='black', ls='--', lw=3);
ax.plot(xs, ys_rnd, color='black', ls=':', lw=3);
ax.set_xlabel(r'$x$', fontsize=fontsize);
ax.set_ylabel(r'$f(x)$', fontsize=fontsize);
ax.set_xlim(xmin,xmax);
ax.tick_params(axis='both', labelsize=ticksize);

A1x = qt.piecewise_poly(xs, np.array([coeffs[0],np.zeros(len(coeffs[1]))]).T, bounds)

nint = qt.get_nint([A1x, ys_rnd])
n = npres + nint + 1

print(nint, npres, n)

n, n0, nint, nintcs, coeffs, bounds = qt.optimize_coeffs_qubits(f_x, xs, m, npres0, npres1, norder=norder, phase=phase)