import time
import numpy as np
import scipy.sparse as sp
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

from qutip import basis, sigmaz, expect
from qutip.solver.heom import HEOMSolver, DrudeLorentzBath

# system
epsilon = 1.0
H = 0.5 * epsilon * sigmaz()
rho0 = basis(2, 0) * basis(2, 0).dag()

# bath
lam = 0.1
gamma = 0.5
T = 1.0
Nk = 8
Q = sigmaz()

tlist = np.linspace(0, 20, 500)
depth = 2
n = 2

# reference run using the standard qutip solver
bath = DrudeLorentzBath(Q=Q, lam=lam, gamma=gamma, T=T, Nk=Nk)
solver = HEOMSolver(H, bath, max_depth=depth, options={"progress_bar": False})

n_ados = solver._n_ados
n_total = n**2 * n_ados

print(f"ADOs: {n_ados}")
print(f"state vector length: {n_total}")

t0 = time.perf_counter()
result_qutip = solver.run(rho0, tlist)
t_qutip = time.perf_counter() - t0

print(f"qutip runtime: {t_qutip:.3f} s")

sz_qutip = expect(sigmaz(), result_qutip.states)

# extract the RHS matrix
# solver.rhs(0) returns the full HEOM RHS as a qutip Qobj
# .full() converts it to a numpy array
# for a time-independent system this matrix does not change, so rhs(0) is enough
RHS = sp.csr_matrix(solver.rhs(0).full())

print(f"RHS shape: {RHS.shape}")
print(f"non-zeros: {RHS.nnz} out of {n_total**2} ({100*RHS.nnz/n_total**2:.1f}%)")

# build the extended initial state vector
# HEOM state = [vec(rho_s), 0, 0, ...] where all ADOs start at zero
# vec uses column-major order (qutip's convention)
rho0_he = np.zeros(n_total, dtype=complex)
rho0_he[:n**2] = rho0.full().ravel("F")

print(f"first 4 entries: {rho0_he[:4].real}  (should be [1, 0, 0, 0])")

# solve the same ODE using an external integrator
# this is the part that will be replaced by petsc later
# the ODE is: d/dt rho_he = RHS @ rho_he

def rhs_fn(t, y):
    return RHS @ y

t0 = time.perf_counter()
sol = solve_ivp(
    rhs_fn,
    t_span=(tlist[0], tlist[-1]),
    y0=rho0_he,
    t_eval=tlist,
    method="RK45",
    rtol=1e-8,
    atol=1e-10,
)
t_bridge = time.perf_counter() - t0

print(f"bridge runtime: {t_bridge:.3f} s")
print(f"success: {sol.success}")

# extract rho_s(t) from the bridge solution
# first n^2 entries of each column are vec(rho_s), reshape back to n x n

def extract_rho_s(y):
    return y[:n**2].reshape((n, n), order="F")

sz_bridge = np.array([
    extract_rho_s(sol.y[:, i])[0, 0].real - extract_rho_s(sol.y[:, i])[1, 1].real
    for i in range(len(tlist))
])

# validation
max_err = np.max(np.abs(sz_qutip - sz_bridge))
print(f"max error vs qutip: {max_err:.2e}")

fig, axes = plt.subplots(1, 2, figsize=(12, 4))

axes[0].plot(tlist, sz_qutip, lw=2, label="QuTiP")
axes[0].plot(tlist, sz_bridge, lw=1.5, ls="--", label="bridge (scipy RK45)")
axes[0].set_xlabel("time")
axes[0].set_ylabel(r"$\langle \sigma_z \rangle$")
axes[0].legend()
axes[0].grid(alpha=0.3)

axes[1].semilogy(tlist, np.abs(sz_qutip - sz_bridge) + 1e-16)
axes[1].set_xlabel("time")
axes[1].set_ylabel("pointwise error")
axes[1].grid(alpha=0.3)

plt.tight_layout()
plt.savefig("bridge_validation.png", dpi=150, bbox_inches="tight")
plt.show()

# petsc bridge
# swap USE_PETSC to True after installing petsc4py
# install: conda install -c conda-forge petsc petsc4py mpi4py

USE_PETSC = False

if USE_PETSC:
    import petsc4py
    petsc4py.init()
    from petsc4py import PETSc

    # build petsc matrix from CSR data
    A = PETSc.Mat().createAIJ(
        size=(n_total, n_total),
        csr=(RHS.indptr, RHS.indices, RHS.data),
        comm=PETSc.COMM_WORLD,
    )
    A.assemblyBegin()
    A.assemblyEnd()

    x = A.createVecRight()
    f = A.createVecLeft()
    x.setValues(range(n_total), rho0_he)
    x.assemblyBegin()
    x.assemblyEnd()

    def rhs_petsc(ts, t, x, f):
        A.mult(x, f)

    ts = PETSc.TS().create(comm=PETSc.COMM_WORLD)
    ts.setType(PETSc.TS.Type.RK)
    ts.setRHSFunction(rhs_petsc, f)
    ts.setTime(tlist[0])
    ts.setMaxTime(tlist[-1])
    ts.setTimeStep(tlist[1] - tlist[0])
    ts.setTolerances(rtol=1e-8, atol=1e-10)
    ts.setFromOptions()

    t0 = time.perf_counter()
    ts.solve(x)
    t_petsc = time.perf_counter() - t0

    y_final = np.array(x.getArray())
    rho_final = extract_rho_s(y_final)
    sz_final = rho_final[0, 0].real - rho_final[1, 1].real

    print(f"petsc runtime: {t_petsc:.3f} s")
    print(f"<sigma_z> at t_final: {sz_final:.6f}")
    print(f"qutip reference:      {sz_qutip[-1]:.6f}")
else:
    print("petsc bridge skipped. set USE_PETSC = True to run.")