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
from qsim.evolution.lindblad_operators import SpontaneousEmission
from matplotlib import rc

"""rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
rc('font',**{'family':'serif'})
## for Palatino and other serif fonts use:
#rc('font',**{'family':'serif','serif':['Palatino']})
rc('text', usetex=True)
"""
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 true_vs_naive(graph: Graph):
    max_omega_r = 160
    max_omega_g = 50
    amplitude = np.sqrt(max_omega_r**2+max_omega_g**2)
    max_omega_r /= amplitude
    max_omega_g /= amplitude
    k = 14
    a = .2
    b = .35

    def schedule_exp_fixed_true_bright(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

        ratio = max_omega_g / max_omega_r
        omega_g = np.sqrt(amplitude * ratio)
        omega_r = np.sqrt(amplitude / ratio)
        laser.omega_g = omega_g
        laser.omega_r = omega_r
        offset = max_omega_g * max_omega_r * np.cos(x * np.pi) - (omega_r ** 2 - omega_g ** 2)
        energy_shift.energies = (offset,)
        dissipation.omega_g = -omega_g
        dissipation.omega_r = omega_r

    def schedule_exp_fixed_true_dark(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, choose the opposite of the STIRAP sequence
        ratio = max_omega_g / max_omega_r
        omega_g = np.sqrt(amplitude * ratio)
        omega_r = np.sqrt(amplitude / ratio)
        offset = max_omega_g*max_omega_r*np.cos(x * np.pi)-(omega_r**2-omega_g**2)
        energy_shift.energies = (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_naive(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
        # Now, choose the opposite of the STIRAP sequence
        ratio = max_omega_g/max_omega_r
        omega_g = np.sqrt(amplitude*ratio)
        omega_r = np.sqrt(amplitude/ratio)
        offset = max_omega_g * max_omega_r * np.cos(x * np.pi) - (omega_r ** 2 - omega_g ** 2)
        energy_shift.energies = (offset,)
        laser.omega_g = omega_g
        laser.omega_r = omega_r
        emission_gg.rates = (omega_g**2,)
        emission_gr.rates = (omega_r**2,)

    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,))
    emission_gr = SpontaneousEmission(transition=(0, 1), IS_subspace=True, graph=graph)
    emission_gg = SpontaneousEmission(transition=(1, 1), IS_subspace=True, graph=graph)
    eq_true = LindbladMasterEquation(hamiltonians=[laser, energy_shift], jump_operators=[dissipation])
    eq_naive = LindbladMasterEquation(hamiltonians=[laser, energy_shift], jump_operators=[emission_gr, emission_gg])

    def k_alpha_rate():
        # Construct the first order transition matrix
        which = 'S'

        ground_energy, ground_state = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state(which=which)
        if eq == eq_true:
            nh_energy_true = 0
            overlap_true = 0
            #print(times[i], eq_true.hamiltonian.todense())
            for op in eq_true.jump_operators[0].jump_operators:
                nh_energy_true += (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0].real
                overlap_true = overlap_true + (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0].real
            return nh_energy_true, overlap_true
        else:

            nh_energy_naive = 0
            overlap_naive = 0
            for op in eq_naive.jump_operators[0].jump_operators:
                nh_energy_naive += (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0].real
                overlap_naive = overlap_naive + (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0].real
            for op in eq_naive.jump_operators[1].jump_operators:
                nh_energy_naive += (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0].real
                overlap_naive = overlap_naive + (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0].real
            return nh_energy_naive, overlap_naive

    l = 0
    colors = ['limegreen', 'navy', 'firebrick']
    fill_colors = ['palegreen', 'skyblue', 'peachpuff']
    labels = ['incoherent scattering', 'coherent scattering, darker ground state',
              'coherent scattering, brighter ground state']

    fig, ax = plt.subplots(1, 2, sharey=True)
    """plt.rcParams.update({
        "text.usetex": True,
        "font.family": "serif",
        "font.serif": ["Palatino"],
    })"""
    for schedule in [schedule_exp_fixed_true_dark, schedule_exp_fixed_true_bright]:
        if schedule == schedule_exp_fixed_naive:
            eq = eq_naive
        else:
            eq = eq_true
        laser.energies = (1,)
        times = np.linspace(.001, .999, 100)
        nh_energies = []
        overlaps = []
        for i in range(len(times)):
            schedule(times[i], 1)
            #print(times[i], eq.hamiltonian.todense())
            res = k_alpha_rate()
            overlaps.append(res[1])
            nh_energies.append(res[0])
        nh_energies = np.array(nh_energies)
        overlaps = np.array(overlaps)
        print(np.sum(nh_energies / graph.n)/len(times))
        print(np.sum((nh_energies-overlaps) / graph.n)/len(times))

        ax[l].plot(times, nh_energies / graph.n, label='Scattering', color=colors[l])
        ax[l].plot(times, overlaps / graph.n, linestyle=':', color=colors[l], label='Overlap')
        ax[l].fill_between(times, overlaps / graph.n , nh_energies / graph.n, color=fill_colors[l], label='True leakage')
        ax[l].legend(shadow=False, fancybox=True, loc='upper right')
        l += 1
    plt.ylim(0, 1.1)

    #plt.legend()
    ax[0].set_title('Dark protocol')
    #ax[1].set_title('Incoherent scattering')
    ax[1].set_title('Bright protocol')
    ax[0].set_ylabel('Dimensionless rate per atom')
    ax[1].set_xlabel(r'Time $(t/T)$')
    ax[0].set_xlabel(r'Time $(t/T)$')

    plt.tight_layout()
    plt.show()


def leakage_vs_dissipation_angle(graph:Graph):
    angles = np.linspace(0.01, np.pi/2-.01, 20)
    leakages_dark = []
    leakages_bright = []
    scattering_dark = []
    scattering_bright = []
    scattering_incoherent = []
    for _ in range(len(angles)):
        print(_)
        max_omega_r = np.cos(angles[_])
        max_omega_g = np.sin(angles[_])
        amplitude = np.sqrt(max_omega_r ** 2 + max_omega_g ** 2)
        max_omega_r /= amplitude
        max_omega_g /= amplitude
        k = 14
        a = .2
        b = .35

        def schedule_exp_fixed_true_bright(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

            ratio = max_omega_g / max_omega_r
            omega_g = np.sqrt(amplitude * ratio)
            omega_r = np.sqrt(amplitude / ratio)
            laser.omega_g = omega_g
            laser.omega_r = omega_r
            offset = max_omega_g * max_omega_r * np.cos(x * np.pi) - (omega_r ** 2 - omega_g ** 2)
            energy_shift.energies = (offset,)
            dissipation.omega_g = -omega_g
            dissipation.omega_r = omega_r

        def schedule_exp_fixed_true_dark(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, choose the opposite of the STIRAP sequence
            ratio = max_omega_g / max_omega_r
            omega_g = np.sqrt(amplitude * ratio)
            omega_r = np.sqrt(amplitude / ratio)
            offset = max_omega_g * max_omega_r * np.cos(x * np.pi) - (omega_r ** 2 - omega_g ** 2)
            energy_shift.energies = (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_naive(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
            # Now, choose the opposite of the STIRAP sequence
            ratio = max_omega_g / max_omega_r
            omega_g = np.sqrt(amplitude * ratio)
            omega_r = np.sqrt(amplitude / ratio)
            offset = max_omega_g * max_omega_r * np.cos(x * np.pi) - (omega_r ** 2 - omega_g ** 2)
            energy_shift.energies = (offset,)
            laser.omega_g = omega_g
            laser.omega_r = omega_r
            emission_gg.rates = (omega_g ** 2,)
            emission_gr.rates = (omega_r ** 2,)

        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,))
        emission_gr = SpontaneousEmission(transition=(0, 1), IS_subspace=True, graph=graph)
        emission_gg = SpontaneousEmission(transition=(1, 1), IS_subspace=True, graph=graph)
        eq_true = LindbladMasterEquation(hamiltonians=[laser, energy_shift], jump_operators=[dissipation])
        eq_naive = LindbladMasterEquation(hamiltonians=[laser, energy_shift], jump_operators=[emission_gr, emission_gg])

        def k_alpha_rate():
            # Construct the first order transition matrix
            which = 'S'

            ground_energy, ground_state = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state(which=which)
            if eq == eq_true:
                nh_energy_true = 0
                overlap_true = 0
                # print(times[i], eq_true.hamiltonian.todense())
                for op in eq_true.jump_operators[0].jump_operators:
                    nh_energy_true += (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0].real
                    overlap_true = overlap_true + (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0].real
                return nh_energy_true, overlap_true
            else:

                nh_energy_naive = 0
                overlap_naive = 0
                for op in eq_naive.jump_operators[0].jump_operators:
                    nh_energy_naive += (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0].real
                    overlap_naive = overlap_naive + (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0].real
                for op in eq_naive.jump_operators[1].jump_operators:
                    nh_energy_naive += (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0].real
                    overlap_naive = overlap_naive + (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0].real
                return nh_energy_naive, overlap_naive


        for schedule in [schedule_exp_fixed_true_dark, schedule_exp_fixed_true_bright, schedule_exp_fixed_naive]:
            if schedule == schedule_exp_fixed_naive:
                eq = eq_naive
            else:
                eq = eq_true
            laser.energies = (1,)
            times = np.linspace(.001, .999, 100)
            nh_energies = []
            overlaps = []
            for i in range(len(times)):
                schedule(times[i], 1)
                # print(times[i], eq.hamiltonian.todense())
                res = k_alpha_rate()
                overlaps.append(res[1])
                nh_energies.append(res[0])
            nh_energies = np.array(nh_energies)
            overlaps = np.array(overlaps)
            if schedule == schedule_exp_fixed_true_dark:
                leakages_dark.append(np.sum((nh_energies - overlaps) / graph.n) / len(times))
                scattering_dark.append(np.sum((nh_energies) / graph.n) / len(times))
            elif schedule == schedule_exp_fixed_true_bright:
                leakages_bright.append(np.sum((nh_energies - overlaps) / graph.n) / len(times))
                scattering_bright.append(np.sum((nh_energies) / graph.n) / len(times))
            else:
                scattering_incoherent.append(np.sum((nh_energies) / graph.n) / len(times))

    plt.scatter(angles, leakages_dark, color='green', zorder=100)
    plt.plot(angles, leakages_dark, color='green', zorder=100)
    plt.scatter(angles, scattering_incoherent, color='red', zorder=100)
    plt.plot(angles, scattering_incoherent, color='red', zorder=100)

    plt.plot(angles, scattering_dark, color='lightgreen')
    plt.scatter(angles, scattering_dark, color='lightgreen')

    plt.scatter(angles, scattering_bright, color='blue', zorder=100)
    plt.plot(angles, scattering_bright, color='blue', zorder=100)

    plt.plot(angles, leakages_bright, color='lightblue')
    plt.scatter(angles, leakages_bright, color='lightblue')
    print(scattering_dark)
    print(np.flip(scattering_bright))

    plt.tight_layout()
    plt.show()


"""from qsim.graph_algorithms.graph import degree_fails_graph, unit_disk_grid_graph
degeneracy = 0
while degeneracy != 1:
    arr = np.reshape(np.random.binomial(1, [.7]*30), (6, 5))
    graph = unit_disk_grid_graph(arr)
    degeneracy = graph.degeneracy
    print(degeneracy)"""

def low_energy_leakage(graph):
    max_omega_r = 1
    max_omega_g = 1
    num = min(graph.num_independent_sets, 100)
    amplitude = np.sqrt(max_omega_r ** 2 + max_omega_g ** 2)
    max_omega_r /= amplitude
    max_omega_g /= amplitude
    k = 14
    a = .2
    b = .35

    def schedule_exp_fixed_true_bright(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

        ratio = max_omega_g / max_omega_r
        omega_g = np.sqrt(amplitude * ratio)
        omega_r = np.sqrt(amplitude / ratio)
        laser.omega_g = omega_g
        laser.omega_r = omega_r
        offset = max_omega_g * max_omega_r * np.cos(x * np.pi) - (omega_r ** 2 - omega_g ** 2)
        energy_shift.energies = (offset,)
        dissipation.omega_g = omega_g
        dissipation.omega_r = omega_r

    def schedule_exp_fixed_true_dark(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, choose the opposite of the STIRAP sequence
        ratio = max_omega_g / max_omega_r
        omega_g = np.sqrt(amplitude * ratio)
        omega_r = np.sqrt(amplitude / ratio)
        offset = max_omega_g * max_omega_r * np.cos(x * np.pi) - (omega_r ** 2 - omega_g ** 2)
        energy_shift.energies = (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=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)

    laser = EffectiveOperatorHamiltonian(graph=graph, IS_subspace=True,
                                         energies=(1,),
                                         omega_g=np.cos(np.pi / 4),
                                         omega_r=np.sin(np.pi / 4))
    energy_shift = hamiltonian.HamiltonianEnergyShift(IS_subspace=True, graph=graph,
                                                      energies=(2.5,), 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():
        # Construct the first order transition matrix
        which = 'S'
        energies, states = SchrodingerEquation(hamiltonians=eq.hamiltonians).eig(k=num, which=which)
        states = states.T
        rates = np.zeros(num)
        for op in eq.jump_operators[0].jump_operators:
            jump_rates = np.zeros(energies.shape[0] ** 2)
            jump_rates = jump_rates + (
                    np.abs(states.conj().T @ op @ states) ** 2).flatten()
            jump_rates = np.reshape(jump_rates, (energies.shape[0], energies.shape[0]))

            rates = rates + jump_rates[:, 0].flatten().real
        return energies, rates
    n_times = 500
    emax = 2.5
    tmin = .15
    tmax = .9
    times = np.linspace(tmin, tmax, n_times)
    energies = np.zeros((n_times, num))
    rates = np.zeros((n_times, num))
    for i in range(n_times):
        print(i)
        schedule_exp_fixed(times[i], 1)
        energy, rate = k_alpha_rate()
        energies[i, :] = energy
        rates[i, :] = rate
    rates = np.log10(rates/graph.n)
    energies = energies.T - energies[:,0].flatten()
    energies = energies.T
    mask = energies<emax
    print(rates*mask)
    from matplotlib.collections import LineCollection
    fig, ax = plt.subplots()
    norm = plt.Normalize(np.min(rates*mask), np.max(rates*(1-mask)*100+rates))
    print(np.min(rates*mask), np.max(rates*(1-mask)*100+rates))
    for i in range(num):
        print(i)
        points = np.array([times, energies[:, i]]).T.reshape(-1, 1, 2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)
        lc = LineCollection(segments, cmap='coolwarm', norm=norm)
        # Set the values used for colormapping
        lc.set_array(rates[:-1, i])
        lc.set_linewidth(1)
        line = ax.add_collection(lc)
    cbar = fig.colorbar(line, ax=ax)
    cbar.ax.set_ylabel(r'$\log_{10}(\bar{r}_{0\rightarrow j})$')
    ax.set_xlim(tmin,tmax)
    ax.set_xlabel(r'Time ($t/T$)')
    ax.set_ylabel(r'Eigenenergy ($E_j-E_0$)')
    ax.set_ylim(np.min(energies)-.3, emax)
    #ax.annotate(r'$r_{0\rightarrow j} = \sum_\mu |\langle j |c_\mu |0\rangle |^2$', xy=(0.4, 0.1), xycoords='data')
    plt.show()
    return

arr = np.array([[0, 1, 0, 0, 1], [1, 1, 1, 1, 0], [0, 1, 1, 0, 1], [1, 1, 1, 1, 1], [1, 1, 0, 1, 1], [1, 0, 1, 1, 0]])
from qsim.graph_algorithms.graph import unit_disk_grid_graph
unit_disk_grid_graph(arr, visualize=True)
import networkx as nx
graph = nx.erdos_renyi_graph(20, .3)
nx.draw(graph)
plt.show()
low_energy_leakage(Graph(graph))
