import numpy as np
import matplotlib.pyplot as plt
import itertools

import numpy as np
import scipy.integrate
import scipy.optimize
from odeintw import odeintw

import scipy.sparse as sparse
from scipy.linalg import expm
from scipy.sparse.linalg import expm_multiply

from qsim.codes import qubit
from qsim.codes.quantum_state import State
from qsim.evolution import lindblad_operators, hamiltonian
from qsim.graph_algorithms.graph import Graph
from qsim.graph_algorithms.graph import line_graph, degree_fails_graph
from qsim.lindblad_master_equation import LindbladMasterEquation
from qsim.schrodinger_equation import SchrodingerEquation
from qsim.graph_algorithms.adiabatic import SimulateAdiabatic
from qsim.tools import tools

"""
Plan:
We want to minimize k_alpha * k_beta, given a fixed functional form for the two photon detuning which has one or two
free parameters that scale and shift the curve, respectively.
"""


class EffectiveOperatorHamiltonian(object):
    def __init__(self, omega_g, omega_r, energies=(1,), graph: Graph = None, IS_subspace=True, code=qubit):
        # Just need to define self.hamiltonian
        assert IS_subspace
        self.energies = energies
        self.IS_subspace = IS_subspace
        self.graph = graph
        self.omega_r = omega_r
        self.omega_g = omega_g
        self.code = code
        assert self.code is qubit

        if self.IS_subspace:
            # Generate sparse mixing Hamiltonian
            assert graph is not None
            assert isinstance(graph, Graph)
            # We have already solved for this information
            IS, nary_to_index, num_IS = graph.independent_sets, graph.binary_to_index, graph.num_independent_sets
            self.transition = (0, 1)
            try:
                self._hamiltonian_rr = np.zeros((num_IS, num_IS))
                self._hamiltonian_gg = np.zeros((num_IS, num_IS))
            except:
                self._hamiltonian_rr = sparse.csc_matrix((num_IS, num_IS))
                self._hamiltonian_gg = sparse.csc_matrix((num_IS, num_IS))
            self._hamiltonian_cross_terms = sparse.csc_matrix((num_IS, num_IS))
            for k in IS:
                self._hamiltonian_rr[k, k] = np.sum(IS[k][2] == self.transition[0])
                self._hamiltonian_gg[k, k] = np.sum(IS[k][2] == self.transition[1])
            self._csc_hamiltonian_rr = sparse.csc_matrix(self._hamiltonian_rr)
            self._csc_hamiltonian_gg = sparse.csc_matrix(self._hamiltonian_gg)
            # For each IS, look at spin flips generated by the laser
            # Over-allocate space
            rows = np.zeros(graph.n * num_IS, dtype=int)
            columns = np.zeros(graph.n * num_IS, dtype=int)
            entries = np.zeros(graph.n * num_IS, dtype=float)
            num_terms = 0
            for i in IS:
                for j in range(len(IS[i][2])):
                    if IS[i][2][j] == self.transition[1]:
                        # Flip spin at this location
                        # Get binary representation
                        temp = IS[i][2].copy()
                        temp[j] = self.transition[0]
                        flipped_temp = tools.nary_to_int(temp, base=code.d)
                        if flipped_temp in nary_to_index:
                            # This is a valid spin flip
                            rows[num_terms] = nary_to_index[flipped_temp]
                            columns[num_terms] = i
                            entries[num_terms] = 1
                            num_terms += 1
            # Cut off the excess in the arrays
            columns = columns[:2 * num_terms]
            rows = rows[:2 * num_terms]
            entries = entries[:2 * num_terms]
            # Populate the second half of the entries according to self.pauli
            columns[num_terms:2 * num_terms] = rows[:num_terms]
            rows[num_terms:2 * num_terms] = columns[:num_terms]
            entries[num_terms:2 * num_terms] = entries[:num_terms]
            # Now, construct the Hamiltonian
            self._csc_hamiltonian_cross_terms = sparse.csc_matrix((entries, (rows, columns)), shape=(num_IS, num_IS))
            self._hamiltonian_cross_terms = self._csc_hamiltonian_cross_terms

        else:
            # We are not in the IS subspace
            pass

    @property
    def hamiltonian(self):
        return self.energies[0] * (self.omega_g * self.omega_r * self._hamiltonian_cross_terms +
                                   self.omega_g ** 2 * self._csc_hamiltonian_gg +
                                   self.omega_r ** 2 * self._csc_hamiltonian_rr)

    def left_multiply(self, state: State):
        return self.hamiltonian @ state

    def right_multiply(self, state: State):
        return state @ self.hamiltonian

    def evolve(self, state: State, time):
        if state.is_ket:
            return State(expm_multiply(-1j * time * self.hamiltonian, state),
                         is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

        else:
            exp_hamiltonian = expm(-1j * time * self.hamiltonian)
            return State(exp_hamiltonian @ state @ exp_hamiltonian.conj().T,
                         is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)


class EffectiveOperatorDissipation(lindblad_operators.LindbladJumpOperator):
    def __init__(self, omega_g, omega_r, rates=(1,), graph: Graph = None, IS_subspace=True, code=qubit):
        self.omega_g = omega_g
        self.omega_r = omega_r

        self.IS_subspace = IS_subspace
        self.transition = (0, 1)
        self.graph = graph
        # Construct jump operators
        if self.IS_subspace:
            # Generate sparse mixing Hamiltonian
            assert graph is not None
            assert isinstance(graph, Graph)
            if code is not qubit:
                IS, nary_to_index, num_IS = graph.independent_sets_qudit(self.code)
            else:
                # We have already solved for this information
                IS, nary_to_index, num_IS = graph.independent_sets, graph.binary_to_index, graph.num_independent_sets
            self._jump_operators_rg = []
            self._jump_operators_gg = []
            # For each atom, consider the states spontaneous emission can generate transitions between
            # Over-allocate space
            for j in range(graph.n):
                rows_rg = np.zeros(num_IS, dtype=int)
                columns_rg = np.zeros(num_IS, dtype=int)
                entries_rg = np.zeros(num_IS, dtype=int)
                rows_gg = np.zeros(num_IS, dtype=int)
                columns_gg = np.zeros(num_IS, dtype=int)
                entries_gg = np.zeros(num_IS, dtype=int)
                num_terms_gg = 0
                num_terms_rg = 0
                for i in IS:
                    if IS[i][2][j] == self.transition[0]:
                        # Flip spin at this location
                        # Get binary representation
                        temp = IS[i][2].copy()
                        temp[j] = self.transition[1]
                        flipped_temp = tools.nary_to_int(temp, base=code.d)
                        if flipped_temp in nary_to_index:
                            # This is a valid spin flip
                            rows_rg[num_terms_rg] = nary_to_index[flipped_temp]
                            columns_rg[num_terms_rg] = i
                            entries_rg[num_terms_rg] = 1
                            num_terms_rg += 1
                    elif IS[i][2][j] == self.transition[1]:
                        rows_gg[num_terms_gg] = i
                        columns_gg[num_terms_gg] = i
                        entries_gg[num_terms_gg] = 1
                        num_terms_gg += 1

                # Cut off the excess in the arrays
                columns_rg = columns_rg[:num_terms_rg]
                rows_rg = rows_rg[:num_terms_rg]
                entries_rg = entries_rg[:num_terms_rg]
                columns_gg = columns_gg[:num_terms_gg]
                rows_gg = rows_gg[:num_terms_gg]
                entries_gg = entries_gg[:num_terms_gg]
                # Now, append the jump operator
                jump_operator_rg = sparse.csc_matrix((entries_rg, (rows_rg, columns_rg)), shape=(num_IS, num_IS))
                jump_operator_gg = sparse.csc_matrix((entries_gg, (rows_gg, columns_gg)), shape=(num_IS, num_IS))

                self._jump_operators_rg.append(jump_operator_rg)
                self._jump_operators_gg.append(jump_operator_gg)
            self._jump_operators_rg = np.asarray(self._jump_operators_rg)
            self._jump_operators_gg = np.asarray(self._jump_operators_gg)
        else:
            # self._jump_operators_rg = []
            # self._jump_operators_gg = []
            op_rg = np.array([[[0, 0], [1, 0]]])
            op_gg = np.array([[[0, 0], [0, 1]]])
            self._jump_operators_rg = op_rg
            self._jump_operators_gg = op_gg

        super().__init__(None, rates=rates, graph=graph, IS_subspace=IS_subspace, code=code)

    @property
    def jump_operators(self):
        return np.sqrt(self.rates[0]) * (self.omega_g * self._jump_operators_gg +
                                         self.omega_r * self._jump_operators_rg)

    @property
    def liouville_evolution_operator(self):
        if self._evolution_operator is None and self.IS_subspace:
            num_IS = self.graph.num_independent_sets
            self._evolution_operator = sparse.csr_matrix((num_IS ** 2, num_IS ** 2))
            for jump_operator in self.jump_operators:
                # Jump operator is real, so we don't need to conjugate
                self._evolution_operator = self._evolution_operator + sparse.kron(jump_operator,
                                                                                  jump_operator) - 1 / 2 * \
                                           sparse.kron(jump_operator.T @ jump_operator, sparse.identity(num_IS)) - \
                                           1 / 2 * sparse.kron(sparse.identity(num_IS), jump_operator.T @ jump_operator)

        elif self._evolution_operator is None:
            # TODO: generate the evolution operator for non-IS subspace states
            raise NotImplementedError
        return self.rates[0] * self._evolution_operator


def angle_vs_time(graph: Graph, verbose=False, mode='reit'):
    max_omega_g = 1
    max_omega_r = 1
    k = 50
    a = .16
    b = .35

    def schedule(t, tf, angle):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                    -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                                (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        omega_g = np.sqrt(amplitude / np.tan(angle))
        omega_r = omega_g * np.tan(angle)
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = (np.sin(phi) ** 2 - np.cos(phi) ** 2 - (omega_r ** 2 - omega_g ** 2),)
        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def is_allowed_angle(t, tf, angle):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                    -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / ((
                                                                                1 + np.e ** (
                                                                                    -k * (1 / 2 - (tf - a)))) ** b) + 1)
        omega_g = np.sqrt(amplitude / np.tan(angle))
        omega_r = omega_g * np.tan(angle)
        if omega_g <= max_omega_g and omega_r <= max_omega_r:
            if mode == 'reit':
                if (omega_g >= omega_r and t / tf >= .5) or (omega_g <= omega_r and t / tf <= .5):
                    return True
            return True
        return False

    laser = EffectiveOperatorHamiltonian(graph=graph, IS_subspace=True, energies=(1,), omega_g=1, omega_r=1)
    energy_shift = hamiltonian.HamiltonianEnergyShift(IS_subspace=True, graph=graph, index=0)
    dissipation = EffectiveOperatorDissipation(graph=graph, omega_r=1, omega_g=1, rates=(1,))
    if mode == 'reit':
        eq = LindbladMasterEquation(hamiltonians=[laser], jump_operators=[dissipation])
    else:
        eq = LindbladMasterEquation(hamiltonians=[laser, energy_shift], jump_operators=[dissipation])

    def k_alpha_rate(t, ratio):
        # Construct the first order transition matrix
        if is_allowed_angle(t, 1, ratio):
            schedule(t, 1, ratio)
            ground_energy, ground_state = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state()
            overlap = 0
            for op in eq.jump_operators[0].jump_operators:
                overlap = overlap - (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0] + \
                          (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0]
            if graph.degeneracy > 1:
                # Solve for the k lowest eigenvalues, where k=degeneracy
                energies, states = SchrodingerEquation(hamiltonians=eq.hamiltonians).eig(k=graph.degeneracy + 1)
                states = states.T
                rates_into_degenerate = np.zeros(energies.shape[0] ** 2)
                for op in eq.jump_operators[0].jump_operators:
                    rates_into_degenerate = rates_into_degenerate + (
                            np.abs(states.conj().T @ op @ states) ** 2).flatten()
                rates_into_degenerate = np.reshape(rates_into_degenerate, (energies.shape[0], energies.shape[0]))
                rates_into_degenerate = rates_into_degenerate[:, 0].flatten().real

                rates_into_degenerate = rates_into_degenerate[1:graph.degeneracy]
                rates_into_degenerate = np.sum(rates_into_degenerate)

                overlap = overlap.real - rates_into_degenerate
            return overlap.real
        else:
            return np.nan

    n_times = 50
    n_ratios = 200
    times = np.linspace(.001, .99, n_times)
    angles = np.linspace(.001, .99, n_ratios) * np.pi / 2
    rates = np.zeros((n_times, n_ratios))
    for t in range(n_times):
        if verbose:
            print(times[t])
        for r in range(n_ratios):
            rates[t, r] = k_alpha_rate(times[t], angles[r])
    return rates.T


import networkx as nx


def dissipation_angle_comparison(graph: Graph, visualize=False):
    # TODO: correct the method for handling degeneracy to actually track the eigenstates
    max_omega_g = 1/np.sqrt(2)
    max_omega_r = 1/np.sqrt(2)
    k = 14
    a = .2
    b = .35

    def schedule_STIRAP(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                    -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                                (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be
        if x < 1/2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        energy_shift.energies = (0,)

        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_fixed_angle(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                    -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                                (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1/2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r**2-omega_g**2
        energy_shift.energies = (offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    def schedule_opposite_STIRAP(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                    -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                                (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1/2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r**2-omega_g**2
        # Now, choose the opposite of the STIRAP sequence
        if x < 1/2:
            # Pick omega_g to be small
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        else:
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        energy_shift.energies = (offset-(omega_r**2-omega_g**2),)

        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_opposite_energy_shift(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                    -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                                (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1/2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r**2-omega_g**2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (-2*offset,)
        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_exp_fixed(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (-offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)


    laser = EffectiveOperatorHamiltonian(graph=graph, IS_subspace=True, energies=(1,), omega_g=1, omega_r=1)
    energy_shift = hamiltonian.HamiltonianEnergyShift(IS_subspace=True, graph=graph, index=0)
    dissipation = EffectiveOperatorDissipation(graph=graph, omega_r=1, omega_g=1, rates=(1,))
    eq = LindbladMasterEquation(hamiltonians=[laser, energy_shift], jump_operators=[dissipation])

    def k_alpha_rate(t):
        # Construct the first order transition matrix
        schedule(t, 1)
        if schedule == schedule_opposite_energy_shift:
            which = 'L'
        else:
            which = 'S'

        ground_energy, ground_state = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state(which=which)
        overlap = 0
        for op in eq.jump_operators[0].jump_operators:
            overlap = overlap - (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0] + \
                      (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0]
        if graph.degeneracy > 1:
            # Solve for the k lowest eigenvalues, where k=degeneracy
            energies, states = SchrodingerEquation(hamiltonians=eq.hamiltonians).eig(k=graph.degeneracy + 1, which=which)
            states = states.T
            rates_into_degenerate = np.zeros(energies.shape[0] ** 2)
            for op in eq.jump_operators[0].jump_operators:
                rates_into_degenerate = rates_into_degenerate + (
                        np.abs(states.conj().T @ op @ states) ** 2).flatten()
            rates_into_degenerate = np.reshape(rates_into_degenerate, (energies.shape[0], energies.shape[0]))
            rates_into_degenerate = rates_into_degenerate[:, 0].flatten().real

            rates_into_degenerate = rates_into_degenerate[1:graph.degeneracy]
            rates_into_degenerate = np.sum(rates_into_degenerate)

            overlap = overlap.real - rates_into_degenerate
        return overlap.real

    labels = ['Option 1', 'Option 2', 'Option 3']
    l = 0
    for schedule in [schedule_exp_fixed]:
        if True:
            if not visualize:
                print(scipy.integrate.quad(k_alpha_rate, .001, .999)[0])
            if visualize:
                times = np.linspace(.001, .999, 100)
                rates = []
                for i in range(len(times)):
                    rates.append(k_alpha_rate(times[i]))
                print(np.sum(rates)/len(times))
                plt.scatter(times, rates, label=labels[l])
        if False:
            times = np.linspace(.001, .999, 800)
            rates = []
            ad = SimulateAdiabatic(graph, noise_model='continuous', hamiltonian=[laser, energy_shift],
                                   noise=[dissipation], cost_hamiltonian=energy_shift)
            all_states, all_indices = ad.groundstate_ordering_vs_time(times, schedule=schedule)
            for i in range(len(times)):
                overlap = 0
                ground_state = all_states[i, 0, :, np.newaxis]
                for op in eq.jump_operators[0].jump_operators:
                    if i == 0:
                        print((np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0],
                        (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0])
                    overlap = overlap - (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0] + \
                              (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0]
                if graph.degeneracy > 1:
                    # Solve for the k lowest eigenvalues, where k=degeneracy
                    states = all_states[i].T
                    rates_into_degenerate = np.zeros(graph.degeneracy ** 2)
                    for op in eq.jump_operators[0].jump_operators:
                        rates_into_degenerate = rates_into_degenerate + (
                                np.abs(states.conj().T @ op @ states) ** 2).flatten()
                    rates_into_degenerate = np.reshape(rates_into_degenerate, (graph.degeneracy, graph.degeneracy))
                    rates_into_degenerate = rates_into_degenerate[:, 0].flatten().real

                    rates_into_degenerate = rates_into_degenerate[1:graph.degeneracy]
                    rates_into_degenerate = np.sum(rates_into_degenerate)

                    overlap = overlap.real - rates_into_degenerate
                    if i == 0:
                        print(rates_into_degenerate, overlap)
                rates.append(overlap)
            print(np.sum(rates) / len(times))
            plt.scatter(times, rates, label=labels[l]+' new')
        l+=1

    if visualize:
        plt.legend()
        plt.xlabel(r'$t/T$')
        plt.ylabel(r'Leakage rate')
        plt.show()


def leakage_rate_plot(graph: Graph):
    max_omega_g = np.sqrt(80 ** 2 / (80 ** 2 + 25 ** 2))
    max_omega_r = np.sqrt(25 ** 2 / (80 ** 2 + 25 ** 2))
    """max_omega_g = 1 / np.sqrt(2)
    max_omega_r = 1 / np.sqrt(2)"""
    k = 14
    a = .2
    b = .35

    def schedule_STIRAP(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        energy_shift.energies = (0,)

        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_fixed_angle(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        energy_shift.energies = (offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    def schedule_opposite_STIRAP(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        else:
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        energy_shift.energies = (offset - (omega_r ** 2 - omega_g ** 2),)

        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_opposite_energy_shift(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (-2 * offset,)
        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_exp_fixed(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (-offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    def schedule_exp_opp_fixed(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (-offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)


    laser = EffectiveOperatorHamiltonian(graph=graph, IS_subspace=True, energies=(1,), omega_g=1, omega_r=1)
    energy_shift = hamiltonian.HamiltonianEnergyShift(IS_subspace=True, graph=graph, index=0)
    dissipation = EffectiveOperatorDissipation(graph=graph, omega_r=1, omega_g=1, rates=(1,))
    eq = LindbladMasterEquation(hamiltonians=[laser, energy_shift], jump_operators=[dissipation])

    def k_alpha_rate(t):
        # Construct the first order transition matrix
        schedule(t, 1)
        if schedule == schedule_opposite_energy_shift or schedule == schedule_exp_opp_fixed:
            which = 'L'
        else:
            which = 'S'

        ground_energy, ground_state = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state(which=which)
        #print(t, SchrodingerEquation(hamiltonians=eq.hamiltonians).hamiltonian)
        nh_energy = 0
        overlap = 0
        for op in eq.jump_operators[0].jump_operators:
            nh_energy += (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0]
            overlap = overlap + (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0]

        if graph.degeneracy > 1:
            # Solve for the k lowest eigenvalues, where k=degeneracy
            energies, states = SchrodingerEquation(hamiltonians=eq.hamiltonians).eig(k=graph.degeneracy + 1,
                                                                                     which=which)
            states = states.T
            rates_into_degenerate = np.zeros(energies.shape[0] ** 2)
            for op in eq.jump_operators[0].jump_operators:
                rates_into_degenerate = rates_into_degenerate + (
                        np.abs(states.conj().T @ op @ states) ** 2).flatten()
            rates_into_degenerate = np.reshape(rates_into_degenerate, (energies.shape[0], energies.shape[0]))
            rates_into_degenerate = rates_into_degenerate[:, 0].flatten().real

            rates_into_degenerate = rates_into_degenerate[1:graph.degeneracy]
            rates_into_degenerate = np.sum(rates_into_degenerate)

            overlap = overlap.real + rates_into_degenerate
        return overlap.real, nh_energy.real

    l = 0
    colors = ['limegreen', 'navy']
    labels = ['STIRAP', r'fixed bright state $|B\rangle \propto |g\rangle +|r\rangle $']
    labels_jump = ['ground state algorithm jump prob.', 'highest excited state algorithm jump prob.']
    labels_leakage = ['ground state algorithm leakage prob.', 'highest excited state algorithm leakage prob.']

    def schedule_exp_fixed_new(t, tf=1):
        x = t / tf*np.pi/2
        amplitude = np.abs(np.cos(x)*np.sin(x))
        # Now we need to figure out what the driver strengths should be for STIRAP
        omega_g = np.sin(x)
        omega_r = np.cos(x)
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    def schedule_reit_new(t, tf=1):
        x = t / tf*np.pi/2
        amplitude = np.abs(np.cos(x)*np.sin(x))
        # Now we need to figure out what the driver strengths should be for STIRAP
        omega_g = np.sin(x)
        omega_r = np.cos(x)
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (0,)
        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    math_labels = []
    for schedule in [schedule_exp_fixed_new, schedule_reit_new]:
        laser.energies = (-1,)
        if True:
            times = np.linspace(.001, .999, 100)
            nh_energies = []
            overlaps = []
            for i in range(len(times)):
                res = k_alpha_rate(times[i])
                overlaps.append(res[0])
                nh_energies.append(res[1])
            nh_energies = np.array(nh_energies)
            overlaps = np.array(overlaps)
            #plt.plot(times, nh_energies/graph.n, label=r'jump prob. per atom', color='navy')
            #plt.plot(times, overlaps/graph.n, label=r'self jump prob. per atom', color='green')
            #plt.fill_between(times, nh_energies/graph.n, overlaps/graph.n, color='palegreen', label='leakage prob. per atom')
            #plt.plot(times, (nh_energies-overlaps)/graph.n, label=r'leakage rate per atom', color='cornflowerblue')
            plt.plot(times, nh_energies/graph.n, label=labels_jump[l], color=colors[l])
            plt.plot(times, (nh_energies-overlaps)/graph.n, color=colors[l], label=labels_leakage[l], linestyle=':')
            #print((np.sum(nh_energies-overlaps)/len(times)))
            print((np.sum((nh_energies-overlaps)))/len(times))
            #print(overlaps[len(times)//2]/nh_energies[len(times)//2])
            #print(np.max(nh_energies)/graph.n)
            #print(np.max(nh_energies-overlaps) / graph.n)
        if False:
            times = np.linspace(.001, .999, 800)
            rates = []
            ad = SimulateAdiabatic(graph, noise_model='continuous', hamiltonian=[laser, energy_shift],
                                   noise=[dissipation], cost_hamiltonian=energy_shift)
            all_states, all_indices = ad.groundstate_ordering_vs_time(times, schedule=schedule)
            for i in range(len(times)):
                overlap = 0
                ground_state = all_states[i, 0, :, np.newaxis]
                for op in eq.jump_operators[0].jump_operators:
                    if i == 0:
                        print((np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0],
                              (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0])
                    overlap = overlap - (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0] + \
                              (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0]
                if graph.degeneracy > 1:
                    # Solve for the k lowest eigenvalues, where k=degeneracy
                    states = all_states[i].T
                    rates_into_degenerate = np.zeros(graph.degeneracy ** 2)
                    for op in eq.jump_operators[0].jump_operators:
                        rates_into_degenerate = rates_into_degenerate + (
                                np.abs(states.conj().T @ op @ states) ** 2).flatten()
                    rates_into_degenerate = np.reshape(rates_into_degenerate, (graph.degeneracy, graph.degeneracy))
                    rates_into_degenerate = rates_into_degenerate[:, 0].flatten().real

                    rates_into_degenerate = rates_into_degenerate[1:graph.degeneracy]
                    rates_into_degenerate = np.sum(rates_into_degenerate)

                    overlap = overlap.real - rates_into_degenerate
                    if i == 0:
                        print(rates_into_degenerate, overlap)
                rates.append(overlap)
            print(np.sum(rates) / len(times))
            plt.scatter(times, rates, label=labels[l] + ' new')
        l += 1
    times = np.linspace(.001, .999, 100)
    omegas = []
    #plt.plot([], [], linestyle=':', color='black', label='leakage rate per atom')
    #plt.plot([], [], color='black', label='jump rate per atom')

    #plt.plot(times, omegas, color='black', label=r'$\frac{\gamma\Omega^2(t)}{\delta_e^2}\cdot T\cdot |V|$')
    plt.legend(loc='upper left')
    plt.xlabel(r'$t/T$')
    plt.ylabel(r'probability')

    plt.show()

#leakage_rate_plot(line_graph(1))
def STIRAP_vs_fixed(graph: Graph):
    max_omega_g = np.sqrt(80**2/(80**2+25**2))
    max_omega_r = np.sqrt(25**2/(80**2+25**2))
    k = 14
    a = .2
    b = .35

    def schedule_STIRAP(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        energy_shift.energies = (0,)

        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_fixed_angle(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        energy_shift.energies = (offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    def schedule_opposite_STIRAP(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        else:
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        energy_shift.energies = (offset - (omega_r ** 2 - omega_g ** 2),)

        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_opposite_energy_shift(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (-2 * offset,)
        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_exp_fixed(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    def schedule_exp_opp_fixed(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (-offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    laser = EffectiveOperatorHamiltonian(graph=graph, IS_subspace=True, energies=(1,), omega_g=1, omega_r=1)
    energy_shift = hamiltonian.HamiltonianEnergyShift(IS_subspace=True, graph=graph, index=0)
    dissipation = EffectiveOperatorDissipation(graph=graph, omega_r=1, omega_g=1, rates=(1,))
    eq = LindbladMasterEquation(hamiltonians=[laser, energy_shift], jump_operators=[dissipation])

    def k_alpha_rate(t):
        # Construct the first order transition matrix
        schedule(t, 1)
        if schedule == schedule_opposite_energy_shift or schedule == schedule_exp_opp_fixed:
            which = 'L'
        else:
            which = 'S'

        ground_energy, ground_state = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state(which=which)
        nh_energy = 0
        overlap = 0
        for op in eq.jump_operators[0].jump_operators:
            nh_energy += (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0]
            overlap = overlap + (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0]

        if graph.degeneracy > 1:
            # Solve for the k lowest eigenvalues, where k=degeneracy
            energies, states = SchrodingerEquation(hamiltonians=eq.hamiltonians).eig(k=graph.degeneracy + 1,
                                                                                     which=which)
            states = states.T
            rates_into_degenerate = np.zeros(energies.shape[0] ** 2)
            for op in eq.jump_operators[0].jump_operators:
                rates_into_degenerate = rates_into_degenerate + (
                        np.abs(states.conj().T @ op @ states) ** 2).flatten()
            rates_into_degenerate = np.reshape(rates_into_degenerate, (energies.shape[0], energies.shape[0]))
            rates_into_degenerate = rates_into_degenerate[:, 0].flatten().real

            rates_into_degenerate = rates_into_degenerate[1:graph.degeneracy]
            rates_into_degenerate = np.sum(rates_into_degenerate)

            overlap = overlap.real + rates_into_degenerate
        return overlap.real, nh_energy.real

    l = 0
    colors = ['limegreen', 'navy']
    labels = ['STIRAP jump operator', r'fixed jump operator $\propto |g\rangle (\langle g| +\langle r|)$']
    labels_leakage = ['ground state, leakage rate', 'highest excited, leakage rate']

    math_labels = []
    for schedule in [schedule_STIRAP, schedule_exp_fixed]:
        if True:
            times = np.linspace(.001, .999, 100)
            nh_energies = []
            overlaps = []
            for i in range(len(times)):
                res = k_alpha_rate(times[i])
                overlaps.append(res[0])
                nh_energies.append(res[1])
            nh_energies = np.array(nh_energies)
            overlaps = np.array(overlaps)
            #plt.plot(times, nh_energies/graph.n, label=r'jump rate per atom', color='navy')
            #plt.plot(times, overlaps/graph.n, label=r'self jump per atom', color='limegreen')
            #plt.plot(times, (nh_energies-overlaps)/graph.n, label=r'leakage rate per atom', color='cornflowerblue')
            #plt.plot(times, nh_energies/graph.n, label=labels[l], color=colors[l])
            plt.plot(times, (nh_energies-overlaps)/graph.n, color=colors[l], label=labels[l])
            print(np.sum(nh_energies-overlaps)/len(times))
        if False:
            times = np.linspace(.001, .999, 800)
            rates = []
            ad = SimulateAdiabatic(graph, noise_model='continuous', hamiltonian=[laser, energy_shift],
                                   noise=[dissipation], cost_hamiltonian=energy_shift)
            all_states, all_indices = ad.groundstate_ordering_vs_time(times, schedule=schedule)
            for i in range(len(times)):
                overlap = 0
                ground_state = all_states[i, 0, :, np.newaxis]
                for op in eq.jump_operators[0].jump_operators:
                    if i == 0:
                        print((np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0],
                              (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0])
                    overlap = overlap - (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0] + \
                              (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0]
                if graph.degeneracy > 1:
                    # Solve for the k lowest eigenvalues, where k=degeneracy
                    states = all_states[i].T
                    rates_into_degenerate = np.zeros(graph.degeneracy ** 2)
                    for op in eq.jump_operators[0].jump_operators:
                        rates_into_degenerate = rates_into_degenerate + (
                                np.abs(states.conj().T @ op @ states) ** 2).flatten()
                    rates_into_degenerate = np.reshape(rates_into_degenerate, (graph.degeneracy, graph.degeneracy))
                    rates_into_degenerate = rates_into_degenerate[:, 0].flatten().real

                    rates_into_degenerate = rates_into_degenerate[1:graph.degeneracy]
                    rates_into_degenerate = np.sum(rates_into_degenerate)

                    overlap = overlap.real - rates_into_degenerate
                    if i == 0:
                        print(rates_into_degenerate, overlap)
                rates.append(overlap)
            print(np.sum(rates) / len(times))
            plt.scatter(times, rates, label=labels[l] + ' new')
        l += 1
    times = np.linspace(.001, .999, 100)
    omegas = []
    #plt.plot([], [], linestyle=':', color='black', label='leakage rate per atom')
    #plt.plot([], [], color='black', label='jump rate per atom')

    #plt.plot(times, omegas, color='black', label=r'$\frac{\gamma\Omega^2(t)}{\delta_e^2}\cdot T\cdot |V|$')
    plt.legend(loc='upper left')
    plt.xlabel(r'$t/T$')
    plt.ylabel(r'leakage probability per atom')

    plt.show()


def low_spectrum(graph: Graph):
    max_omega_g = 1 / np.sqrt(2)
    max_omega_r = 1 / np.sqrt(2)
    k = 14
    a = .2
    b = .35

    def schedule_STIRAP(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        energy_shift.energies = (0,)

        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_fixed_angle(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        energy_shift.energies = (offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    def schedule_opposite_STIRAP(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        else:
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        energy_shift.energies = (offset - (omega_r ** 2 - omega_g ** 2),)

        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_opposite_energy_shift(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (-2 * offset,)
        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_exp_fixed(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    def schedule_exp_opp_fixed(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (-offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    def schedule_hybrid(t, tf):
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = (1/2*np.sin(phi)**2,)
        laser.omega_g = np.cos(phi)
        laser.omega_r = np.sin(phi)
        dissipation.omega_g = np.cos(phi)
        dissipation.omega_r = np.sin(phi)

    laser = EffectiveOperatorHamiltonian(graph=graph, IS_subspace=True, energies=(1,), omega_g=1, omega_r=1)
    energy_shift = hamiltonian.HamiltonianEnergyShift(IS_subspace=True, graph=graph, index=0)
    dissipation = EffectiveOperatorDissipation(graph=graph, omega_r=1, omega_g=1, rates=(1,))
    eq = LindbladMasterEquation(hamiltonians=[laser, energy_shift], jump_operators=[dissipation])

    def jump_rate(states, energies, t):
        # Construct the first order transition matrix
        schedule(t, 1)
        states = states.T
        rates_into_degenerate = np.zeros(energies.shape[0] ** 2)
        for op in eq.jump_operators[0].jump_operators:
            rates_into_degenerate = rates_into_degenerate + (
                    np.abs(states.conj().T @ op @ states) ** 2).flatten()
        rates_into_degenerate = np.reshape(rates_into_degenerate, (energies.shape[0], energies.shape[0]))
        rates_into_degenerate = rates_into_degenerate[:, 0].flatten().real

        return rates_into_degenerate
    from matplotlib.cm import get_cmap
    for schedule in [schedule_hybrid]:
        times = np.linspace(.001, .999, 100)
        ad = SimulateAdiabatic(graph, hamiltonian=[laser, energy_shift], noise=[dissipation], cost_hamiltonian=energy_shift)
        all_states, all_indices = ad.groundstate_ordering_vs_time(times, schedule=schedule)
        res = ad.spectrum_vs_time(1, schedule, k=5, num=100)
        for k in range(res.shape[-1]):
            if k == 0:
                plt.plot(np.linspace(0, 1, res.shape[0]), res[:,k].flatten(), color='cornflowerblue')
            elif k < 2:
                plt.plot(np.linspace(0, 1, res.shape[0]), res[:,k].flatten(), color='cornflowerblue')
            else:
                plt.plot(np.linspace(0, 1, res.shape[0]), res[:, k].flatten(), color='black')
        plt.xlabel(r'$t/T$')
        plt.yticks(range(0, 4))
        plt.ylabel(r'energy')
        plt.show()
        ad = SchrodingerEquation(hamiltonians=[laser, energy_shift])
        max_index = int(np.max(all_indices))+0
        all_rates = []
        all_energies = []
        times = np.linspace(.1, .999, 100)
        for i in range(len(times)):
            schedule(times[i], 1)
            energies, states = ad.eig(k=max_index+1)
            rates = jump_rate(states, energies, times[i])
            all_rates.append(rates)
            all_energies.append(energies)
    all_rates = np.log10(np.array(all_rates))
    all_energies = np.array(all_energies).T
    from matplotlib.collections import LineCollection
    fig, ax = plt.subplots()
    norm = plt.Normalize(all_rates.min(), all_rates.max())
    for i in range(max_index+1):
        if True:
            time_segments = np.concatenate([times[:-1], times[1:]]).reshape((2, len(times)-1))
            time_segments = time_segments.flatten(order='F')
            energies = all_energies[i,:]
            energies = np.concatenate([energies[:-1], energies[1:]]).reshape((2, len(energies) - 1))
            energies = energies.flatten(order='F')
            segments = np.reshape(np.concatenate([time_segments, energies]), (2, len(time_segments))).T
            segments = np.reshape(segments, (int(len(time_segments)/2), 2, 2))
        else:
            energies = all_energies[i,:]
            segments = np.reshape(np.concatenate([times, energies]), (2, len(times))).T
            segments = np.reshape(segments, (int(len(times) / 2), 2, 2))
        lc = LineCollection(segments, cmap='viridis', norm=norm)
        lc.set_array(all_rates[:-1, i])
        lc.set_linewidth(2)

        line = ax.add_collection(lc)
    #cbar = fig.colorbar(line, ax=ax)
    #cbar.set_label(r'$\log_{10}$(jump probability)')
    plt.ylim(np.min(all_energies), np.max(all_energies))
    #plt.legend(loc='upper left')
    plt.xlabel(r'$t/T$')
    plt.ylabel(r'eigenenergies')
    #plt.plot([], [], color='navy', label='optimal solution')
    #plt.plot([], [], linestyle='dashed', color='navy', label='non-optimal solution')
    plt.legend(loc='upper right')
    plt.show()


def leakage_rate_degenerate(graph: Graph):
    max_omega_g = 1 / np.sqrt(2)
    max_omega_r = 1 / np.sqrt(2)
    k = 14
    a = .2
    b = .35

    def schedule_STIRAP(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        energy_shift.energies = (0,)

        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_fixed_angle(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        energy_shift.energies = (offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    def schedule_opposite_STIRAP(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        else:
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        energy_shift.energies = (offset - (omega_r ** 2 - omega_g ** 2),)

        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_opposite_energy_shift(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (-2 * offset,)
        laser.omega_g = omega_g
        laser.omega_r = omega_r
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_exp_fixed(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    def schedule_exp_opp_fixed(t, tf):
        x = t / tf
        amplitude = max_omega_g * max_omega_r * (
                -1 / (1 + np.e ** (k * (x - a))) ** b - 1 / (1 + np.e ** (-k * (x - (tf - a)))) ** b + 1) / \
                    (-1 / ((1 + np.e ** (k * (1 / 2 - a))) ** b) - 1 / (
                            (1 + np.e ** (-k * (1 / 2 - (tf - a)))) ** b) + 1)
        # Now we need to figure out what the driver strengths should be for STIRAP
        if x < 1 / 2:
            # Pick omega_g to be small
            omega_g = amplitude / max_omega_r
            omega_r = max_omega_r
        else:
            omega_r = amplitude / max_omega_g
            omega_g = max_omega_g
        offset = omega_r ** 2 - omega_g ** 2
        # Now, choose the opposite of the STIRAP sequence
        energy_shift.energies = (-offset,)
        laser.omega_g = np.sqrt(amplitude)
        laser.omega_r = np.sqrt(amplitude)
        dissipation.omega_g = np.sqrt(amplitude)
        dissipation.omega_r = np.sqrt(amplitude)

    laser = EffectiveOperatorHamiltonian(graph=graph, IS_subspace=True, energies=(1,), omega_g=1, omega_r=1)
    energy_shift = hamiltonian.HamiltonianEnergyShift(IS_subspace=True, graph=graph, index=0)
    dissipation = EffectiveOperatorDissipation(graph=graph, omega_r=1, omega_g=1, rates=(1,))
    eq = LindbladMasterEquation(hamiltonians=[laser, energy_shift], jump_operators=[dissipation])

    def jump_rate(states, t):
        # Construct the first order transition matrix
        schedule(t, 1)
        size = states.shape[0]
        states = states.T
        leakage = 0
        for op in eq.jump_operators[0].jump_operators:
            leakage = leakage + (np.abs(states.conj().T @ op.conj().T @ op @ states)).flatten()[0]
        rates_into_degenerate = np.zeros(size ** 2)
        for op in eq.jump_operators[0].jump_operators:
            rates_into_degenerate = rates_into_degenerate + (
                    np.abs(states.conj().T @ op @ states) ** 2).flatten()
        rates_into_degenerate = np.reshape(rates_into_degenerate, (size, size))
        rates_into_degenerate = rates_into_degenerate[:, 0].flatten().real

        return leakage-np.sum(rates_into_degenerate)

    for schedule in [schedule_exp_fixed]:
        times = np.linspace(.75, .999, 300)
        if not isinstance(graph, Graph):
            ad = SimulateAdiabatic(graph, hamiltonian=[laser, energy_shift], noise=[dissipation],
                                   cost_hamiltonian=energy_shift)
            all_states, all_indices = ad.groundstate_ordering_vs_time(times, schedule=schedule)
            all_rates = []
            for i in range(len(times)):
                states = all_states[i]
                schedule(times[i], 1)
                rates = jump_rate(states, times[i])
                all_rates.append(rates)
        else:
            ad = SchrodingerEquation(hamiltonians=[laser, energy_shift])
            all_rates = []
            for i in range(len(times)):
                schedule(times[i], 1)
                energies, states = ad.eig(k=graph.degeneracy)
                rates = jump_rate(states, times[i])
                all_rates.append(rates)
    return np.sum(all_rates)/len(times)


from qsim.graph_algorithms.graph import unit_disk_grid_graph, ring_graph

"""
Works: 
[[0 1 1]
 [0 1 1]
 [1 0 1]]
"""

#arr = np.reshape(np.random.binomial(1, [.5]*9), (3, 3))
#print(arr)
#low_spectrum(unit_disk_grid_graph(arr))

"""degeneracy = 2
while degeneracy != 1:
    arr = np.reshape(np.random.binomial(1, [.8]*30), (6, 5))
    graph = unit_disk_grid_graph(arr)
    degeneracy = graph.degeneracy
print(repr(arr), np.sum(arr))"""
"""arr = np.array([[1, 1, 0, 0, 1],
 [1, 1, 1, 1, 0],
 [1, 1, 1, 1, 0],
 [1, 1, 1, 1, 1],
 [1, 1, 0, 1, 1],
 [0, 0, 1, 1, 1]])"""
arr = np.reshape(np.random.binomial(1, [.8]*25), (5, 5))
print(len(arr))
from matplotlib import rc
rc('text', usetex=True)
rc('font', **{'family': 'serif'})
#leakage_rate_plot(unit_disk_grid_graph(arr))
data = np.array([0.1724948655910386, 0.3034094275495633, 0.4409289424024603, 0.5064584681298279, 0.5801610644825725])
data_STIRAP = np.array([0.12888770689657505, 0.2420344464508475, 0.3624534403491076, 0.4202028578783863, 0.4849868979133825])
# Last part of ramp
data_STIRAP = np.array([0.0034982512792422816, 0.006810576042416697, 0.010375622892122944, 0.01405023978879786])
data = np.array([0.14903346566668588, 0.22347824867186522, 0.2980234122539539, 0.3726195422204543])
ns = np.array([9, 15, 21, 27])
plt.scatter(ns, data_STIRAP, color='midnightblue', zorder=100)
plt.scatter(ns, data, color='steelblue', zorder=100)
#plt.scatter([9, 15, 21, 27], data-data_STIRAP)
plt.xlabel(r'System Size ($|V|$)')
plt.xticks(ns)
plt.legend()
plt.ylabel(r'Dimensionless leakage rate, last quarter of ramp')
m, b = np.polyfit(ns, data, deg=1)
print(m)
plt.plot(ns, m*ns+b, color='lightsteelblue', label=r'Dark protocol, leakage$=1\times 10^{-4}|V|$')
m, b = np.polyfit(ns, data_STIRAP, deg=1)
print(m)
plt.plot(ns, m*ns+b, color='royalblue', label=r'STIRAP, leakage$=1\times 10^{-2}|V|$')
plt.legend(loc='lower right')
plt.hlines(0, 0, 30, linestyles=':', color='black')
plt.show()
plt.ylim(-.1, np.max(data))
#leakage_rate_plot(unit_disk_grid_graph(arr))
rates = []
for i in [(3, 3), (3, 5), (3, 7), (3, 9)]:
    arr = np.ones(i)
    rates.append(leakage_rate_degenerate(unit_disk_grid_graph(arr)))
    print(rates)
#[0.1724948655910386, 0.08298716167325303] [(3, 3), (3, 4), (4, 4), (4, 5), (5, 5)]
"""rates = []
degeneracies = []
ns = []
while True:
    try:
        arr = np.reshape(np.random.binomial(1, [.8]*30), (6, 5))
        graph = unit_disk_grid_graph(arr)
        degeneracy = graph.degeneracy
        print(degeneracy, graph.n)
        rate = leakage_rate_degenerate(graph)/graph.n
        degeneracies.append(graph.degeneracy)
        rates.append(rate)
        ns.append(graph.n)
        print(rates, degeneracies, ns)
    except Exception as e:
        print(e)"""

#graph = unit_disk_grid_graph(arr, visualize=False)
#graph = Graph(nx.star_graph(5))
graph = unit_disk_grid_graph(np.ones((2,1)), visualize=False)
"""try:
    optimal_rates = np.nanargmin(rates, axis=0)
    worst_rates = np.nanargmax(rates, axis=0)
except ValueError:
    

print(optimal_rates, worst_rates)
best_rate = 0
worst_rate = 0
for i in range(len(optimal_rates)):
    if rates[optimal_rates[i], i] != 10:
        best_rate += rates[optimal_rates[i], i]
        worst_rate += rates[worst_rates[i], i]
    rates[optimal_rates[i], i] = 0
    rates[worst_rates[i], i] = 0
print(best_rate/rates.shape[1])
print(worst_rate/rates.shape[1])
plt.imshow(np.log10(rates), extent=[0, np.pi/2, 0, 1], aspect='auto')
plt.colorbar()
plt.xlabel(r'$t/T$')
plt.ylabel(r'$\theta$')

plt.show()"""
