#  =========================================================================  #
#   Copyright 2018 National Technology & Engineering Solutions of Sandia,
#   LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS,
#   the U.S. Government retains certain rights in this software.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#  =========================================================================  #

import numpy as np
from scipy.optimize import curve_fit
from scipy.optimize import brentq, newton
import matplotlib.pyplot as plt
from .. import circuit_runners
from ..qeccs import Surface4444
from ..decoders import MWPM2D
from ..error_gens import XGen
from .threshold_tools import codecapacity_logical_rate, codecapacity_logical_rate2, codecapacity_logical_rate3
from ..misc.threshold_curve import func


def pseudo_threshold_code_capacity(ps, distance, runs, qecc_class=None, error_gen=None, decoder_class=None,
                                   verbose=True, mode=1, deg=2, circuit_runner=None, plotting=False, basis=None):
    """
    Function that generates p_logical values given a list of physical errors (ps) and distance (ds).
    Args:
        ps:
        distance:
        runs:
        qecc_class:
        error_gen:
        decoder_class:
        verbose:
        mode:
        deg:
        circuit_runner:

    Returns:

    """

    if circuit_runner is None:
        circuit_runner = circuit_runners.Standard()

    if error_gen is None:
        error_gen = XGen(model_level='code_capacity')

    if qecc_class is None:
        qecc_class = Surface4444

    if decoder_class is None:
        decoder_class = MWPM2D

    if mode == 1:
        determine_rate = codecapacity_logical_rate
    elif mode == 2:
        determine_rate = codecapacity_logical_rate2
    elif mode == 3:
        determine_rate = codecapacity_logical_rate3
    else:
        raise Exception('Mode "%s" is not handled!' % mode)

    ps = np.array(ps)

    plog = []

    qecc = qecc_class(distance=distance)
    decoder = decoder_class(qecc)

    for p in ps:

        logical_error_rate, time = determine_rate(runs, qecc, distance, error_gen, error_params={'p': p},
                                                  decoder=decoder, verbose=verbose, circuit_runner=circuit_runner,
                                                  basis=basis)
        if verbose and time:
            print('Runtime: %s s' % time)

        if verbose:

            print('----')

        plog.append(logical_error_rate)

    plog = np.array(plog)

    if verbose:
        print('ps=', ps)
        print('plog=', plog)

    if plotting:

        find_polyfit(ps, plog, deg, verbose)

        plot(ps, plog, deg)

    # return plist, dlist, plog, results
    # return {'plist': plist, 'distance': distance, 'plog': plog, 'opt': popt, 'std': stdev, 'pseudo_threshold': pseudo_thr}
    return {'ps': ps, 'distance': distance, 'plog': plog}


def find_polyfit(ps, plog, deg, verbose=True):

    plist = np.array(ps)

    popt, pcov = np.polyfit(ps, plog, deg=deg, cov=True)

    var = np.diag(pcov)
    stdev = np.sqrt(var)

    if verbose:
        print('params=', popt)
        print('std=', stdev)

    pseudo_thr = find_pseudo(plist, plog, deg)

    if verbose:
        print('Pseudo-threshold: %s' % pseudo_thr)

    return pseudo_thr, popt, pcov


def find_uniscalefit(ps, plog, distance, p0=None, maxfev=1000000, verbose=True, **kwargs):
    plist = np.array(ps)
    dlist = ns2nsfit(distance, len(plist))

    popt, pcov = curve_fit(func, (plist, dlist), plog, p0, maxfev=maxfev, **kwargs)

    var = np.diag(pcov)
    stdev = np.sqrt(var)

    for v in var:
        if np.isnan(v):
            raise Exception('Was not able to find a good fit. Suggestion: Use `p0` to specify parameter guess.')

    # print('popt', pcov)
    # print('pcov', pcov)
    # print('var', var)
    # print('std', stdev)

    pseudo_thr = popt[0]
    v0 = popt[1]
    pseudo_thr_std = stdev[0]
    v0_std = stdev[1]

    if verbose:
        print('pseudo-threshold: %s +- %s (1 stdev)' % (pseudo_thr, pseudo_thr_std))
        print('v0: %s +- %s (1 stdev)' % (v0, v0_std))

    return pseudo_thr, pseudo_thr_std, v0, v0_std, popt, pcov


def ns2nsfit(ns, num):
    """
    Returns a list of distances or ps for performing fits.

    If ds == 5 and num == 3:
        -> [5, 5, 5]

    If ds == [3, 5, 7] and num == 3:
        -> [3, 3, 3, 5, 5, 5, 7, 7, 7]

    Likewise for ps.

    Args:
        ds:
        num:

    Returns:

    """
    if isinstance(ns, int):
        return [ns]*num

    else:
        new_list = []

        for i in ns:
            new_list.extend([i]*num)
        return new_list


def find_pseudo(plist, plog, deg):
    """
    Determines the pseudo threshold from list of ps and plogs.

    Args:
        plist:
        plog:
        deg:

    Returns:
        float: The value of the pseudo-threshold.

    """
    popt = np.polyfit(plist, plog, deg=deg)
    poly = np.poly1d(popt)

    def fnc(x):
        return poly(x) - x

    try:
        pseudo_thr = brentq(fnc, 0, 1)
    except ValueError:
        pseudo_thr = newton(fnc, 0.05)

    return pseudo_thr


def plot(plist, plog, deg=2, figsize=(10, 10), p_start=None, p_end=None):
    """

    Args:
        plist:
        plog:
        deg(int): Degree of polynomial fit.
        figsize(tuple of int):
        axis_start(float): Where the x and y axes begin.
        axis_end(float): Where the x and y axes end.

    Returns:

    """

    if p_start is None:
        p_start = min(plog) * 0.9

    if p_end is None:
        p_end = max(plog) * 1.1

    psuedo_thr = find_pseudo(plist, plog, deg)

    popt, pcov = np.polyfit(plist, plog, deg, cov=True)  # C_z is estimated covariance matrix

    axis_start = p_start
    axis_end = p_end

    x = np.linspace(axis_start, axis_end, 1000)

    poly = np.poly1d(popt)
    yi = poly(x)

    # Do the plotting:
    fg, ax = plt.subplots(1, 1, figsize=figsize)
    # ax.set_title("Polynomial Fit of Degree %s with Error of $\pm1\sigma$" % deg, size=20)
    ax.set_title("Pseudothreshold from Polynomial Fit of Degree %s" % deg, size=20)

    ax.plot(x, yi, '-', )
    ax.plot(plist, plog, 'ro')
    ax.axis('tight')

    y = x
    plt.plot(x, y, 'k-', alpha=.30)

    ax.set_ylim(axis_start, axis_end)
    ax.set_xlim(axis_start, axis_end)

    plt.xlabel('Physical error rate', size=18)
    plt.ylabel('Logical error rate', size=18)

    pth = psuedo_thr
    plt.axvline(pth, color='green', linewidth=2, linestyle='dashed', label='Pseudo-threshold (%s)' % pth)
    plt.legend(fontsize=16)

    fg.canvas.draw()
    plt.show()
