import figsaver as fs
from hiro_models.one_qubit_model import QubitModel, StocProcTolerances
import hiro_models.model_auxiliary as aux
from hops.util.utilities import relative_entropy, relative_entropy_single, entropy, trace_distance
from hopsflow.util import EnsembleValue
import numpy as np
import qutip as qt
import scipy
import utilities as ut
import matplotlib
matplotlib.rcParams.update(fs.MPL_RC_POSTER)

import ray
ray.shutdown()
ray.init(address="auto")

from hops.util.logging_setup import logging_setup
import logging
logging_setup(logging.INFO, show_stocproc=False)

from tabulate import tabulate
import hiro_models.model_auxiliary as aux
import hiro_models.model_base as Model

def org_tbl(dct: dict[str, Model]) -> str:
    return tabulate(dct, headers=dct.keys(), tablefmt="orgtbl")


def print_diff_tbl(models):
    print(
        org_tbl(
            aux.model_diff_dict(
                models,
                extra_fields=models[0].std_extra_fields
                | {"Samples": lambda model: aux.get_data(model).samples},
            )
        )
    )

def ergo(T, ω=1):
    return (T * np.log(1 + np.exp(-ω / T)))

Δ = 2
model, strobe_t, strobe_indices = energy_shovel(Δ=Δ, γ=.5, periods=6)

from hops.util.dynamic_matrix import SmoothStep, Periodic, Harmonic, ConstantMatrix

Δ = 5
γ = 1 / Δ * 1 / 10

# , np.linspace(1.9, 2.3, 25)
ωs = np.unique(np.sort(np.concatenate([[.5], np.linspace(1, 1.8, 6)])))
ω_models = []
 # ω_models = [QubitModel(
 #     δ=1/(ω_c),
 #     ω_c=ω_c,
 #     ω_s=1+Δ-ω_c,
 #     t=ut.linspace_with_strobe(0, t_max, 500, Δ),
 #     ψ_0=(0.5 * 0 * qt.basis([2], [0]) + qt.basis([2], [1])),
 #     description=f"Testing the time dependent coupling with smooth step.",
 #     k_max=5,
 #     bcf_terms=7,
 #     truncation_scheme="simplex",
 #     driving_process_tolerance=StocProcTolerances(1e-3, 1e-3),
 #     thermal_process_tolerance=StocProcTolerances(1e-3, 1e-3),
 #     T=3,
 #     L=L,
 #     H=H,
 #     therm_method="fft",
 # ) for ω_c in ωs]

for ω_c in ωs:
     proto, strobe_t, strobe_indices = energy_shovel(Δ, periods=30, modulate_system=False)
     proto.δ = 2 / ω_c
     proto.ω_c = ω_c
     proto.ω_s = 1 + Δ - ω_c
     proto.k_max = 7
     proto.therm_method = "fft"

     ω_models.append(proto)

fs.tex_value(
    ω_models[1].bcf_scale * ω_models[1].bcf(0).real,
    prec=1,
    prefix="α(0)=",
    save="omega_alpha",
), fs.tex_value(
    Δ,
    prec=1,
    prefix="Δ=",
    save="omega_delta",
), fs.tex_value(
    model.T,
    prec=1,
    prefix="T=",
    save="omega_delta",
)

aux.integrate_multi(ω_models, 10000)

from pprint import pprint
print_diff_tbl(ω_models)

fig, ax = plt.subplots()
fig.set_size_inches(fs.get_figsize("poster", 0.49))
for model, data in aux.model_data_iterator(ω_models):
    _, _, bar = fs.plot_with_σ(
        model.t,
        model.total_energy(data) * (1 / (model.T * np.log(1 + np.exp(-1 / model.T)))),
        ax=ax,
        label=fr"$ω_c={model.ω_c:.2f}$",
        strobe_frequency=Δ,
        markersize=3,
    )

    fs.plot_with_σ(
        model.t,
        model.total_energy(data) * (1 / (model.T * np.log(1 + np.exp(-1 / model.T)))),
        ax=ax,
        # label=fr"$ω_c={model.ω_c:.2f}$",
        # strobe_frequency=Δ,
        linewidth=0.5,
        alpha=0.2,
        color=bar.lines[0].get_color(),
    )

    τ_off = 1 / (model.bcf_coefficients()[1][0].real.min())
    a0 = abs(model.bcf(0))
    τ_off = scipy.optimize.fsolve(lambda τ: abs(model.bcf(τ)) - (a0/300), 10)
    ax.axvline(
        τ_off,
        color=bar.lines[0].get_color(),
        linestyle="dotted",
        linewidth=1,
        zorder=-1,
    )

ax.set_xlabel(r"$\tau$")
ax.set_ylabel(r"$(\langle H\rangle_\tau -\langle H\rangle_0)/\Delta E_\mathrm{max}$")
# ax.axhline(, color="grey", linewidth=1, linestyle="dashed")
# ax.plot(model.t, model.L.operator_norm(model.t), linewidth=1)
ax.legend()
fs.export_fig("omegas_total")

from hops.util.dynamic_matrix import SmoothStep, Periodic, Harmonic, ConstantMatrix


Δs = np.linspace(1,10,10)#np.sort(np.concatenate((np.linspace(1, 5, 20), np.linspace(5, 7, int(20/5 * 2)))))
Δ_models = []
strobe_ts = []
strobe_indices_s = []
for Δ in Δs:
    proto, strobe_t, strobe_indices = energy_shovel(Δ, periods=20, k_max=7, modulate_system=False)
    strobe_ts.append(strobe_t)
    strobe_indices_s.append(strobe_indices)
    Δ_models.append(proto)

aux.integrate_multi(Δ_models, 1000)

final_e = []
final_e_error = []
ensemble_arg = dict(overwrite_cache=False)
fig, ax = plt.subplots()
fig.set_size_inches(fs.get_figsize("poster", .49))
for (model, data), Δ, strobe_t, strobe_indices in zip(
    aux.model_data_iterator(Δ_models), Δs, strobe_ts, strobe_indices_s
):
    energies = model.energy_change_from_interaction_power(data, **ensemble_arg).for_bath(0)
    idx = energies.value[strobe_indices].argmin()
    energies = energies * (1 / strobe_t[idx])
    final_e.append(energies.value[strobe_indices[idx]])
    final_e_error.append(energies.σ[strobe_indices[idx]])
    # fs.plot_energy_overview(model, ensemble_args=ensemble_arg)
    # fig, ax = plt.subplots()
    # fs.plot_with_σ(model.t, 1 *model.energy_change_from_interaction_power(data, **ensemble_arg).for_bath(0), ax=ax)
    # fs.plot_with_σ(model.t,  model.interaction_power(data, **ensemble_arg).for_bath(0), ax=ax)
    # fs.plot_with_σ(model.t, model.total_energy(data, **ensemble_arg), ax=ax)
    # ax.plot(model.t, model.L.derivative()(model.t)[:,0,1])
    # ax.plot(model.t, model.L(model.t)[:,0,1])
    # print(strobe_t[1])
    # break

#  plt.ylabel(r"$\frac{\Delta E}{T}$")
ax.set_xlabel(r"$\Delta$")
ax.set_ylabel(r"$P_\mathrm{max}$")
ax.errorbar(Δs, final_e, final_e_error)