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.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)
            if code is not qubit:
                IS, num_IS = graph.independent_sets_qudit(self.code)
            else:
                # We have already solved for this information
                IS, num_IS = graph.independent_sets, graph.num_independent_sets
            self.transition = (0, 1)
            self._hamiltonian_rr = np.zeros((num_IS, num_IS))
            self._hamiltonian_gg = np.zeros((num_IS, num_IS))
            self._hamiltonian_cross_terms = np.zeros((num_IS, num_IS))
            for k in range(IS.shape[0]):
                self._hamiltonian_rr[k, k] = np.sum(IS[k][1] == self.transition[0])
                self._hamiltonian_gg[k, k] = np.sum(IS[k][1] == 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 range(graph.num_independent_sets):
                for j in range(graph.n):
                    if IS[i,j] == self.transition[1]:
                        # Flip spin at this location
                        # Get binary representation
                        temp = IS[i].copy()
                        temp[j] = self.transition[0]
                        where_matched = (np.argwhere(np.sum(np.abs(IS - temp), axis=1) == 0).flatten())
                        if len(where_matched) > 0:
                            # This is a valid spin flip
                            rows[num_terms] = where_matched[0]
                            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, num_IS = graph.independent_sets_qudit(self.code)
            else:
                # We have already solved for this information
                IS, num_IS = graph.independent_sets, 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 range(IS.shape[0]):
                    if IS[i,j] == self.transition[0]:
                        # Flip spin at this location
                        # Get binary representation
                        temp = IS[i].copy()
                        temp[j] = self.transition[1]
                        where_matched = (np.argwhere(np.sum(np.abs(IS - temp), axis=1) == 0).flatten())
                        if len(where_matched) > 0:
                            # This is a valid spin flip
                            rows_rg[num_terms_rg] = where_matched[0]
                            columns_rg[num_terms_rg] = i
                            entries_rg[num_terms_rg] = 1
                            num_terms_rg += 1
                    elif IS[i,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 k_alpha(delta_r_bar, graph: Graph, mode='hybrid', verbose=False, integrator = 'quad'):
    if not isinstance(delta_r_bar, float):
        delta_r_bar = delta_r_bar[0]

    def schedule_hybrid(t, tf, delta_r_bar=0):
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = (delta_r_bar * 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)

    def schedule_adiabatic(t, tf, delta_r_bar=0):
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = ((delta_r_bar+1) * np.sin(phi) ** 2-np.cos(phi) ** 2,)
        laser.omega_g = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        laser.omega_r = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        dissipation.omega_g = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        dissipation.omega_r = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))

    if mode == 'hybrid':
        schedule = schedule_hybrid
    elif mode=='adiabatic':
        schedule = schedule_adiabatic

    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, delta_r_bar=0):
        if verbose:
            print(t)
        schedule(t, 1, delta_r_bar=delta_r_bar)
        # Construct the first order transition matrix
        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
    if integrator == 'quad':
        k_a, err_k_a = scipy.integrate.quad(lambda t: k_alpha_rate(t, delta_r_bar=delta_r_bar), 0, .95, limit=200)
    else:
        times = np.linspace(.001, .99, 100)
        rates = np.zeros(100)
        i = 0
        for t in times:
            rates[i] = k_alpha_rate(t, delta_r_bar)
            i += 1
        return np.sum(rates)*(times[-1]-times[0])/len(times)
    if verbose:
        print(delta_r_bar, k_a)
    return k_a
from qsim.graph_algorithms.graph import ring_graph
for i in range(0):
    print(i, k_alpha([0], ring_graph(i), integrator='quad', verbose=False))

def k_beta(delta_r_bar, graph: Graph, verbose=False, mode='hybrid'):

    def schedule_hybrid(t, tf, delta_r_bar=0):
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = (delta_r_bar * 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)

    def schedule_adiabatic(t, tf, delta_r_bar=0):
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = ((delta_r_bar+1) * np.sin(phi) ** 2-np.cos(phi) ** 2,)
        laser.omega_g = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        laser.omega_r = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        dissipation.omega_g = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        dissipation.omega_r = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))


    if mode == 'hybrid':
        schedule = schedule_hybrid
    elif mode=='adiabatic':
        schedule = schedule_adiabatic

    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])

    dt = 0.001
    overlap = 0
    for time in [0, 1-2*dt]:
        def normalize_phase(eig):
            where_nonzero = np.argwhere(np.absolute(eig) > 1e-9)[0]
            eig = np.e**(-1j*np.angle(eig[where_nonzero[0], where_nonzero[1]])) * eig
            # We can take the eigenvalues to be real for this Hamiltonian
            eig = eig.real
            return eig / np.linalg.norm(eig)

        schedule(time, 1, delta_r_bar=delta_r_bar)
        # Construct the first order transition matrix
        ground_energy, ground_state = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state()
        ground_state = normalize_phase(ground_state)
        ham = scipy.sparse.csr_matrix((-ground_energy * np.ones(graph.num_independent_sets),
                                       (range(graph.num_independent_sets), range(graph.num_independent_sets))),
                                      shape=(graph.num_independent_sets, graph.num_independent_sets))
        for h in eq.hamiltonians:
            ham = ham + h.hamiltonian
        schedule(time+dt, 1, delta_r_bar=delta_r_bar)
        # Construct the first order transition matrix
        ground_energy_dt, ground_state_dt = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state()
        ground_state_dt = normalize_phase(ground_state_dt)
        d_ground_state_dt = (ground_state_dt - ground_state) / dt
        # Construct a projector out of the ground subspace
        proj = np.identity(graph.num_independent_sets)
        proj = proj - tools.outer_product(ground_state, ground_state)
        if graph.degeneracy>1 and time == 0:
            # Project out of the ground subspace
            energies, states = SchrodingerEquation(hamiltonians=eq.hamiltonians).eig(k=graph.degeneracy)

            for i in range(graph.degeneracy-1):
                proj = proj - tools.outer_product(states[i+1, np.newaxis].T, states[i+1, np.newaxis].T)
        d_ground_state_dt = proj @ d_ground_state_dt
        # Multiply by 1/(H-E_0)
        # If at the end, remove the rows corresponding to degeneracies
        if time != 0:
            d_ground_state_dt = d_ground_state_dt[graph.degeneracy:]
            ham = ham[graph.degeneracy:, graph.degeneracy:]
        else:
            d_ground_state_dt = d_ground_state_dt[:-1]
            ham = ham[:-1, :-1]
        res = scipy.sparse.linalg.spsolve(ham, d_ground_state_dt).real
        overlap += np.linalg.norm(res)**2
    return overlap


def find_optimal_schedules_old (graph: Graph, verbose=False, mode='hybrid'):

    def schedule_hybrid(t, tf, delta_r_bar=0):
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = (delta_r_bar * 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)

    def schedule_adiabatic(t, tf, delta_r_bar=0):
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = ((delta_r_bar+1) * np.sin(phi) ** 2-np.cos(phi) ** 2,)
        laser.omega_g = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        laser.omega_r = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        dissipation.omega_g = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        dissipation.omega_r = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))

    if mode == 'hybrid':
        schedule = schedule_hybrid
    elif mode=='adiabatic':
        schedule = schedule_adiabatic

    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, delta_r_bar=0):
        schedule(t, 1, delta_r_bar=delta_r_bar)
        # Construct the first order transition matrix
        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)
            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

    def k_alpha(delta_r_bar):
        k_a, err_k_a = scipy.integrate.quad(lambda t: k_alpha_rate(t, delta_r_bar=delta_r_bar), 0, 1, limit=200)
        return k_a

    def k_beta(delta_r_bar):
        dt = 0.001
        overlap = 0
        for time in [0, 1-2*dt]:
            def normalize_phase(eig):
                where_nonzero = np.argwhere(np.absolute(eig) > 1e-9)[0]
                eig = np.e**(-1j*np.angle(eig[where_nonzero[0], where_nonzero[1]])) * eig
                # We can take the eigenvalues to be real for this Hamiltonian
                eig = eig.real
                return eig / np.linalg.norm(eig)

            schedule(time, 1, delta_r_bar=delta_r_bar)
            # Construct the first order transition matrix
            ground_energy, ground_state = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state()
            ground_state = normalize_phase(ground_state)
            ham = scipy.sparse.csr_matrix((-ground_energy * np.ones(graph.num_independent_sets),
                                           (range(graph.num_independent_sets), range(graph.num_independent_sets))),
                                          shape=(graph.num_independent_sets, graph.num_independent_sets))
            for h in eq.hamiltonians:
                ham = ham + h.hamiltonian
            schedule(time+dt, 1, delta_r_bar=delta_r_bar)
            # Construct the first order transition matrix
            ground_energy_dt, ground_state_dt = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state()
            ground_state_dt = normalize_phase(ground_state_dt)
            d_ground_state_dt = (ground_state_dt - ground_state) / dt
            # Construct a projector out of the ground subspace
            proj = np.identity(graph.num_independent_sets)
            proj = proj - tools.outer_product(ground_state, ground_state)
            if graph.degeneracy>1 and time == 0:
                # Project out of the ground subspace
                energies, states = SchrodingerEquation(hamiltonians=eq.hamiltonians).eig(k=graph.degeneracy)

                for i in range(graph.degeneracy-1):
                    proj = proj - tools.outer_product(states[i+1, np.newaxis].T, states[i+1, np.newaxis].T)
            d_ground_state_dt = proj @ d_ground_state_dt
            # Multiply by 1/(H-E_0)
            # If at the end, remove the rows corresponding to degeneracies
            if time != 0:
                d_ground_state_dt = d_ground_state_dt[graph.degeneracy:]
                ham = ham[graph.degeneracy:, graph.degeneracy:]
            else:
                d_ground_state_dt = d_ground_state_dt[:-1]
                ham = ham[:-1, :-1]
            res = scipy.sparse.linalg.spsolve(ham, d_ground_state_dt).real
            overlap += np.linalg.norm(res)**2
        return overlap

    def first_order_correction(delta_r_bar):
        delta_r_bar = delta_r_bar[0]
        beta = k_beta(delta_r_bar)
        alpha = k_alpha(delta_r_bar)
        res = beta*alpha
        if res <= 0:
            if verbose:
                print(delta_r_bar, alpha, beta, res)
            return 0
        res = np.sqrt(res)
        if verbose:
            print(delta_r_bar, alpha, beta, res)
        return res


def find_optimal_schedules(graph: Graph, verbose=False, mode='hybrid'):

    def schedule_hybrid(t, tf, delta_r_bar=0):
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = (delta_r_bar * 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)

    def schedule_adiabatic(t, tf, delta_r_bar=0):
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = ((delta_r_bar+1) * np.sin(phi) ** 2-np.cos(phi) ** 2,)
        laser.omega_g = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        laser.omega_r = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        dissipation.omega_g = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))
        dissipation.omega_r = np.sqrt(np.abs(np.cos(phi)*np.sin(phi)))

    if mode == 'hybrid':
        schedule = schedule_hybrid
    elif mode=='adiabatic':
        schedule = schedule_adiabatic

    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, delta_r_bar=0):
        schedule(t, 1, delta_r_bar=delta_r_bar)
        # Construct the first order transition matrix
        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)
            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

    def k_alpha(delta_r_bar):
        k_a, err_k_a = scipy.integrate.quad(lambda t: k_alpha_rate(t, delta_r_bar=delta_r_bar), 0, 1, limit=200)
        return k_a

    def k_beta(delta_r_bar):
        dt = 0.001
        overlap = 0
        for time in [0, 1-2*dt]:
            def normalize_phase(eig):
                where_nonzero = np.argwhere(np.absolute(eig) > 1e-9)[0]
                eig = np.e**(-1j*np.angle(eig[where_nonzero[0], where_nonzero[1]])) * eig
                # We can take the eigenvalues to be real for this Hamiltonian
                eig = eig.real
                return eig / np.linalg.norm(eig)

            schedule(time, 1, delta_r_bar=delta_r_bar)
            # Construct the first order transition matrix
            ground_energy, ground_state = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state()
            ground_state = normalize_phase(ground_state)
            ham = scipy.sparse.csr_matrix((-ground_energy * np.ones(graph.num_independent_sets),
                                           (range(graph.num_independent_sets), range(graph.num_independent_sets))),
                                          shape=(graph.num_independent_sets, graph.num_independent_sets))
            for h in eq.hamiltonians:
                ham = ham + h.hamiltonian
            schedule(time+dt, 1, delta_r_bar=delta_r_bar)
            # Construct the first order transition matrix
            ground_energy_dt, ground_state_dt = SchrodingerEquation(hamiltonians=eq.hamiltonians).ground_state()
            ground_state_dt = normalize_phase(ground_state_dt)
            d_ground_state_dt = (ground_state_dt - ground_state) / dt
            # Construct a projector out of the ground subspace
            proj = np.identity(graph.num_independent_sets)
            proj = proj - tools.outer_product(ground_state, ground_state)
            if graph.degeneracy>1 and time == 0:
                # Project out of the ground subspace
                energies, states = SchrodingerEquation(hamiltonians=eq.hamiltonians).eig(k=graph.degeneracy)

                for i in range(graph.degeneracy-1):
                    proj = proj - tools.outer_product(states[i+1, np.newaxis].T, states[i+1, np.newaxis].T)
            d_ground_state_dt = proj @ d_ground_state_dt
            # Multiply by 1/(H-E_0)
            # If at the end, remove the rows corresponding to degeneracies
            if time != 0:
                d_ground_state_dt = d_ground_state_dt[graph.degeneracy:]
                ham = ham[graph.degeneracy:, graph.degeneracy:]
            else:
                d_ground_state_dt = d_ground_state_dt[:-1]
                ham = ham[:-1, :-1]
            res = scipy.sparse.linalg.spsolve(ham, d_ground_state_dt).real
            overlap += np.linalg.norm(res)**2
        return overlap

    def first_order_correction(delta_r_bar):
        delta_r_bar = delta_r_bar[0]
        beta = k_beta(delta_r_bar)
        alpha = k_alpha(delta_r_bar)
        res = beta*alpha
        if res <= 0:
            if verbose:
                print(delta_r_bar, alpha, beta, res)
            return 0
        res = np.sqrt(res)
        if verbose:
            print(delta_r_bar, alpha, beta, res)
        return res

# First question - does raising delta_r_bar coincide with the optimal dissipative correction?
# Let's look at graphs on a 4x4 grid and select some nodes at random
rates_ad = []
deg = []
rates_hybrid = []
bob = []


for z in range(1):
    from qsim.graph_algorithms.graph import grid_graph, unit_disk_grid_graph
    arr = np.reshape(np.random.binomial(1, [.65]*20), (4, 5))
    graph = unit_disk_grid_graph(arr)
    #graph = line_graph(17)
    graph = ring_graph(14)
    #graph = line_graph(8)
    deg.append(graph.degeneracy)    #res = scipy.optimize.minimize(lambda dr: k_alpha(dr, graph=graph, verbose=True), [1.], bounds=[(0, 3)])
    #print(res)

    def schedule_hybrid(t, tf, delta_r_bar=0):
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = (delta_r_bar * 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)

    from qsim.graph_algorithms.adiabatic import SimulateAdiabatic
    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 = SimulateAdiabatic(graph=graph, hamiltonian=[laser, energy_shift], noise_model='continuous', noise=[dissipation],
                           cost_hamiltonian=energy_shift)

    #eq.spectrum_vs_time(1, schedule_hybrid, num=500, plot=True)
    times = np.linspace(.01, .95, 500)
    eigvecs, indices = eq.groundstate_ordering_vs_time(times, schedule_hybrid, verbose=True)
    #rates_better_reit = np.zeros(len(times))
    rates_self_reit = np.zeros(len(times))
    rates_out_reit = np.zeros(len(times))
    for i in range(len(times)):
        schedule_hybrid(times[i], 1)
        overlap = 0
        overlap_self = 0
        for op in dissipation.jump_operators:
            ground_state = eigvecs[i, 0, :, np.newaxis]
            overlap_self = overlap_self + (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0]
            overlap = overlap +  (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 = eigvecs[i, :].T

            rates_into_degenerate = np.zeros(states.shape[-1] ** 2)
            for op in dissipation.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, (states.shape[-1], states.shape[-1]))

            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
        rates_out_reit[i] = overlap
        rates_self_reit[i] = overlap_self
    plt.scatter(times, rates_out_reit.real, label=r'$\sum_u\langle c_u^\dagger c_u \rangle $', c='blue')
    plt.scatter(times, rates_self_reit.real, label=r'$\sum_u|\langle  c_u \rangle|^2$ ', c='purple')
    #plt.scatter(times, rates_out_reit.real/rates_self_reit.real, label=r'$\sum_u\langle c_u^\dagger c_u \rangle/\sum_u|\langle  c_u \rangle|^2 $', c='m')
    #rates_hybrid.append(np.sum(rates_better)/len(times))
    rates_self_reit = np.zeros(len(times))
    rates_out_reit = np.zeros(len(times))
    def schedule_adiabatic(t, tf, delta_r_bar=0):
        phi = (tf - t) / tf * np.pi / 2
        energy_shift.energies = ((delta_r_bar + 1) * np.sin(phi) ** 2 - np.cos(phi) ** 2,)
        laser.omega_g = np.sqrt(np.abs(np.cos(phi) * np.sin(phi)))
        laser.omega_r = np.sqrt(np.abs(np.cos(phi) * np.sin(phi)))
        dissipation.omega_g = np.sqrt(np.abs(np.cos(phi) * np.sin(phi)))
        dissipation.omega_r = np.sqrt(np.abs(np.cos(phi) * np.sin(phi)))
    for i in range(len(times)):
        schedule_adiabatic(times[i], 1)
        overlap = 0
        overlap_self = 0
        for op in dissipation.jump_operators:
            ground_state = eigvecs[i, 0, :, np.newaxis]
            overlap_self = overlap_self + (np.abs(ground_state.conj().T @ op @ ground_state) ** 2)[0, 0]
            overlap = overlap + (ground_state.conj().T @ op.conj().T @ op @ ground_state)[0, 0]
        if False:
            # Solve for the k lowest eigenvalues, where k=degeneracy
            states = eigvecs[i, :].T
            rates_into_degenerate = np.zeros(states.shape[-1] ** 2)
            for op in dissipation.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, (states.shape[-1], states.shape[-1]))

            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
        rates_out_reit[i] = overlap
        rates_self_reit[i] = overlap_self
    print(rates_self_reit)

    #plt.scatter(times, rates_out_reit.real, label=r'$\sum_u\langle c_u^\dagger c_u \rangle exp$')
    #plt.scatter(times, rates_self_reit.real, label=r'$\sum_u|\langle  c_u \rangle|^2 exp$')
    #plt.hlines(9, times[0], times[-1], label='naive theory estimate for STIRAP', colors='k')
    #plt.plot(times, 9*np.sin((1-times)*np.pi/2)*np.cos((1-times)*np.pi/2), label='naive theory estimate for exp', c='k', linestyle=':')
    rates_better_ad = np.zeros(len(times))


    for i in range(len(times)):
        schedule_adiabatic(times[i], 1)
        overlap = 0
        for op in dissipation.jump_operators:
            ground_state = eigvecs[i, 0, :, np.newaxis]
            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 = eigvecs[i, :].T
            rates_into_degenerate = np.zeros(states.shape[-1] ** 2)
            for op in dissipation.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, (states.shape[-1], states.shape[-1]))

            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
        rates_better_ad[i] = overlap
    #print(np.sum(rates_better)/len(times))
    #plt.scatter(times, rates_better, label='naive', color='purple')
    plt.xlabel(r"$t/T$")
    plt.legend()
    #plt.ylabel(r"rate exiting ground subspace")
    plt.show()
    """rates_ad.append(np.sum(rates_better_ad)/len(times))
    rates_bob = np.min(np.array([rates_better, rates_better_ad]), axis=0)

    bob.append(np.sum(rates_bob)/len(times))
    print(deg)
    print(rates_ad)
    print(rates_hybrid)
    print(bob)"""
    """
    [[0 0 0 1]
     [1 1 1 0]
     [1 1 0 1]
     [1 1 0 1]] 
     
     bad for reit 
     [[0 1 1 0]
     [1 1 1 1]
     [1 1 0 1]
     [0 1 1 0]] 
    times, rates = k_alpha([3.], graph=graph, verbose=True, integrator='other', mode='hybrid')
    print(np.sum(rates)/len(rates))
    plt.scatter(times, rates)
    times, rates = k_alpha([3.], graph=graph, verbose=True, integrator='other', mode='adiabatic')
    print(np.sum(rates)/len(rates))
    plt.scatter(times, rates)"""

    """plt.legend()
    plt.show()"""

    """for n in []:#[3, 5, 7, 9, 11, 13]:
        graph = grid_graph(n, 2, periodic=True)
        print(graph.degeneracy)
        res = k_alpha(graph, verbose=True, mode='hybrid')
        print(res)
        # print(res.fun, res.x[0])
    plt.show()
    ns_even = np.arange(4, 18, 2)
    rates_grid = [0.012950903366653766, 0.03974034001327809, 0.07070722697744833, 0.10743881803101309, 0.13721670912770378,
                  0.17074450566345314, 0.22870897289690184, 0.24391634600437007]"""

