#!/usr/bin/env python3

# Copyright 2022 Valentin Bruch <valentin.bruch@rwth-aachen.de>
# License: MIT
"""
Kondo FRTRG, generate plots for comparison to experiments
"""

import scipy.constants as sc
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.colors as mplcolors
from matplotlib.widgets import Slider
import argparse
import numpy as np
from scipy.interpolate import bisplrep, bisplev, splrep, BSpline, griddata
import settings
from data_management import DataManager, KondoImport

# In this program all energies are given in units of the RTRG Kondo
# temperature Tkrg, which is an integration constant of the E-flow RG
# equations. The more conventional definition of the Kondo temperature is
# G(V=Tk)=G(V=0)/2=e²/h. The ratio Tk/Tkrg is:
TK_VOLTAGE = 3.4425351 # for rtol=1e-9, atol=1e-11, voltage_branches=10

TK_GLAZMAN = (3/8 / 0.2)**0.5 # from  G/G0 = 1 - 0.2*(V/Tkrg)**2 = 1 - 3/8*(V/Tkglazman)**2  for small V



"""
Kogan 04
* omega=5.9 when following G/G0 = 1 - 3/2*(V/Tk)**2
* omega=7.4 when following G(V=Tk) = G0/2

Bruhat 18
From fit:
f=19GHz: omega=3.57
f=12GHz: omega=2.26
"""


def main():
    """
    Parse command line arguments and call other functions
    """
    parser = argparse.ArgumentParser(description=main.__doc__)
    valid_functions = {f.__name__:f for f in globals().values() if type(f) == type(main)  and getattr(f, '__module__', '') == '__main__' and f.__name__[0] != '_'}
    parser.add_argument("functions", type=str, nargs='+', choices=valid_functions.keys(), help="functions to be called")
    parser.add_argument("--omega", type=float, help="Frequency, units of Tk")
    parser.add_argument("--method", type=str, choices=('J', 'mu'), help="method: J or mu")
    parser.add_argument("--nmax", metavar='int', type=int, help="Floquet matrix size")
    parser.add_argument("--padding", metavar='int', type=int, help="Floquet matrix ppadding")
    parser.add_argument("--voltage_branches", metavar='int', type=int, help="Voltage branches")
    parser.add_argument("--resonant_dc_shift", metavar='int', type=int, help="resonant DC shift")
    parser.add_argument("--vdc", metavar='float', type=float, help="Vdc, units of Tk")
    fourier_coef_group = parser.add_mutually_exclusive_group()
    fourier_coef_group.add_argument("--vac", metavar='float', type=float, help="Vac, units of Tk")
    fourier_coef_group.add_argument("--fourier_coef", metavar='tuple', type=float, nargs='*', help="Voltage Fourier arguments, units of omega")
    parser.add_argument("--d", metavar='float', type=float, help="D (UV cutoff), units of Tk")
    parser.add_argument("--xL", metavar='float', type=float, nargs='+', default=0.5, help="Asymmetry, 0 < xL < 1")
    parser.add_argument("--compact", metavar='int', type=int, help="compact FRTRG implementation (0,1, or 2)")
    parser.add_argument("--solver_tol_rel", metavar="float", type=float, help="Solver relative tolerance")
    parser.add_argument("--solver_tol_abs", metavar="float", type=float, help="Solver relative tolerance")
    args = parser.parse_args()

    dm = DataManager()
    options = args.__dict__
    results = []
    for name in options.pop("functions"):
        results.append(valid_functions[name](dm=dm, **options))
    plt.show()


def bruhat18(dm, **parameters):
    """
    PRB 98.075121
    Tk=28.2μeV
    f=19GHz or f=12GHz
    vac_mueV = [20, 40, 60, ..., 140, 180, 220, 300]

    some data accidentally lie in frtrg-omega10-vdc24-vac22-03.h5 instead
    of frtrg-bruhat18-f?-??.h5
    """
    #tk_mueV = 32.7 # Kondo temperature in μeV, defined by  G/G0 = 1 - 3/8 (V/Tk)^2  for small V
    tkrg_mueV = 19 # Kondo temperature in μeV, estimated based on fit
    f1_ghz = 19 # Frequency, GHz
    f2_ghz = 12 # Frequency, GHz

    omega1_mueV = f1_ghz * 1e15*sc.h / sc.eV
    omega2_mueV = f2_ghz * 1e15*sc.h / sc.eV

    #tkrg_mueV = tk_mueV / TK_GLAZMAN # Kondo temperature in μeV defined as integration constant in RTRG
    print(f"Tkrg = {tkrg_mueV} μeV")

    omega1 = omega1_mueV / tkrg_mueV # Frequency, in RTRG Tk units
    omega2 = omega2_mueV / tkrg_mueV # Frequency, in RTRG Tk units
    print(f"fake precision: Omega1 = {omega1:.6g}, Omega2 = {omega2:.6g}")
    omega1 = round(omega1, 2)
    omega2 = round(omega2, 2)
    print(f"Omega1 = {omega1}\nOmega2 = {omega2}")

    #omega1_old = 9.2159791 # Frequency, in RTRG Tk units
    #omega2_old = 5.8206184 # Frequency, in RTRG Tk units

    voltage_branches = 4
    vac_mueV = np.array([20, 40, 60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280, 300])
    vac = vac_mueV / tkrg_mueV # Vac, in RTRG Tk units
    print("Vac =", vac.round(6))
    print("f=19GHz: Vac/omega =", (vac_mueV/omega1_mueV).round(6))
    print("f=12GHz: Vac/omega =", (vac_mueV/omega2_mueV).round(6))
    print("f=19GHz with correction factor 0.89: Vac/omega =\n\t", (0.89*vac_mueV/omega1_mueV).round(6))
    print("f=12GHz with correction factor 1.07: Vac/omega =\n\t", (1.07*vac_mueV/omega2_mueV).round(6))
    print("Vac max = ", 330 / tkrg_mueV)
    print("19GHz: Vac max = ", 410.5 * omega1/omega1_mueV)
    print("12GHz: Vac max = ", 410.5 * omega2/omega2_mueV)

    print(f"Frequency 19GHz = {omega1_mueV:.6g} μeV")
    print(f"Frequency 12GHz = {omega2_mueV:.6g} μeV")
    print(f"1 meV = {1e3/omega1_mueV:.6g} * 19GHz")
    print(f"1 meV = {1e3/omega2_mueV:.6g} * 12GHz")

    # TODO: generate data, implement plot, check convergence, ...

def bruhat18_fig2a(dm, **kwargs):
    for name in ("omega", "vdc", "vac"):
        kwargs.pop(name, None)
    return bruhat18_fig2ab(dm, omega=5.8206184, **kwargs)

def bruhat18_fig2b(dm, **kwargs):
    for name in ("omega", "vdc", "vac"):
        kwargs.pop(name, None)
    return bruhat18_fig2ab(dm, omega=9.2159791, **kwargs)

def bruhat18_fig2ab(dm, omega, dc_res=100, ac_res=100, vdc_min=0, vac_min=0, vdc_max=50, vac_max=40, **parameters):
    """
    Plot overview of dc and ac current and dc conductance for harmonic driving
    at fixed frequency as function of Vdc and Vac.
    """
    results_all = dm.list(omega=omega, **parameters)
    results = results_all.loc[(results_all["solver_flags"] & DataManager.SOLVER_FLAGS["simplified_initial_conditions"]) == 0]
    j = results.method == "J"
    mu = results.method == "mu"
    vac_max = min(vac_max, results.vac.max())
    vdc_max = min(vdc_max, results.vdc.max())

    tk_mueV = 28.2  # Kondo temperature in μeV, defined by G(V=Tk)=e²/h
    tkrg_mueV = tk_mueV / TK_VOLTAGE # Kondo temperature in μeV defined as integration constant in RTRG

    # Interpolate
    #gdc_J_tck   = bisplrep(results.vac[j], results.vdc[j], results.dc_conductance[j], s=2e-6, kx=3, ky=3)
    #idc_J_tck   = bisplrep(results.vac[j], results.vdc[j], results.dc_current[j], s=2e-6, kx=3, ky=3)
    #iac_J_tck   = bisplrep(results.vac[j], results.vdc[j], results.ac_current_abs[j], s=2e-6, kx=3, ky=3)
    #phase_J_tck = bisplrep(results.vac[j], results.vdc[j], results.ac_current_phase[j], s=2e-6, kx=3, ky=3)
    gdc_mu_tck   = bisplrep(results.vac[mu], results.vdc[mu], results.dc_conductance[mu], s=1e-6, kx=3, ky=3)
    #idc_mu_tck   = bisplrep(results.vac[mu], results.vdc[mu], results.dc_current[mu], s=1e-6, kx=3, ky=3)
    #iac_mu_tck   = bisplrep(results.vac[mu], results.vdc[mu], results.ac_current_abs[mu], s=1e-6, kx=3, ky=3)
    #phase_mu_tck = bisplrep(results.vac[mu], results.vdc[mu], results.ac_current_phase[mu], s=1e-6, kx=3, ky=3)

    vac_arr = np.linspace(vac_max/(2*ac_res), vac_max*(1 - 0.5/ac_res), ac_res)
    vdc_arr = np.linspace(vdc_max/(2*dc_res), vdc_max*(1 - 0.5/dc_res), dc_res)
    #gdc_g_J_interp = bisplev(vac_arr, vdc_arr, gdc_J_tck)
    #gdc_i_J_interp = bisplev(vac_arr, vdc_arr, idc_J_tck, dy=1)
    #idc_J_interp   = bisplev(vac_arr, vdc_arr, idc_J_tck)
    #iac_J_interp   = bisplev(vac_arr, vdc_arr, iac_J_tck)
    #phase_J_interp = bisplev(vac_arr, vdc_arr, phase_J_tck)
    gdc_g_mu_interp = bisplev(vac_arr, vdc_arr, gdc_mu_tck)
    #gdc_i_mu_interp = bisplev(vac_arr, vdc_arr, idc_mu_tck, dy=1)
    #idc_mu_interp   = bisplev(vac_arr, vdc_arr, idc_mu_tck)
    #iac_mu_interp   = bisplev(vac_arr, vdc_arr, iac_mu_tck)
    #phase_mu_interp = bisplev(vac_arr, vdc_arr, phase_mu_tck)

    # Create figure
    fig, ax1 = plt.subplots(1, 1, sharex=True, sharey=True)
    ax1.set_ylabel("Vac (μV)")
    ax1.set_xlabel("Vdc (μV)")

    # DC conductance
    gnorm = plt.Normalize(np.pi*results.dc_conductance.min(), np.pi*results.dc_conductance.max())
    cmap = plt.cm.Oranges
    ax1.set_title('DC conductance')
    #ax1.scatter(tkrg_mueV*results.vdc[j], tkrg_mueV*results.vac[j], c=np.pi*results.dc_conductance[j], marker='x', norm=gnorm, cmap=plt.cm.viridis)
    #ax1.scatter(tkrg_mueV*results.vdc[mu], tkrg_mueV*results.vac[mu], c=np.pi*results.dc_conductance[mu], marker='+', norm=gnorm, cmap=plt.cm.viridis)
    img = ax1.imshow(np.pi*gdc_g_mu_interp, extent=(0, vdc_max*tkrg_mueV, 0, vac_max*tkrg_mueV), aspect='auto', cmap=cmap, norm=gnorm, origin='lower')
    ax1.imshow(np.pi*gdc_g_mu_interp, extent=(0, -vdc_max*tkrg_mueV, 0, vac_max*tkrg_mueV), aspect='auto', cmap=cmap, norm=gnorm, origin='lower')
    ax1.set_xlim(-tkrg_mueV*vdc_max, tkrg_mueV*vdc_max)
    fig.colorbar(img, ax=ax1)

def bruhat18_fig2c(dm, **kwargs):
    for name in ("omega", "vdc", "vac"):
        kwargs.pop(name, None)
    return bruhat18_fig2cd(dm, omega=5.8206184, **kwargs)

def bruhat18_fig2d(dm, **kwargs):
    for name in ("omega", "vdc", "vac"):
        kwargs.pop(name, None)
    return bruhat18_fig2cd(dm, omega=9.2159791, **kwargs)

def bruhat18_fig2cd(dm, omega, **parameters):
    tk_mueV = 28.2  # Kondo temperature in μeV, defined by G(V=Tk)=e²/h
    tkrg_mueV = tk_mueV / TK_VOLTAGE # Kondo temperature in μeV defined as integration constant in RTRG
    xscale = 1e-3 * tkrg_mueV
    yscale = np.pi

    vac_mueV_arr = np.array([20, 40, 60, 80, 100, 120, 140, 160, 180, 220, 300])
    colors = ['#000000', '#7f7f7f', '#a64f00', '#dc2121', '#f07d2e', '#ffa600', '#54a800', '#00a8ff', '#0000ff', '#7f00ff', '#cf00f8']
    vac_arr = vac_mueV_arr / tkrg_mueV # Vac, in RTRG Tk units

    fig, ax = plt.subplots()
    ax.set_xlabel("Vdc (mV)")
    ax.set_ylabel("G (2e²/h)")

    for vac, color in zip(vac_arr.round(6), colors):
        table = dm.list(omega=omega, vac=vac, **parameters)
        table.sort_values('vdc', inplace=True)
        ax.plot(table.vdc*xscale, table.dc_conductance*yscale, '.-', color=color)
        ax.plot(-table.vdc*xscale, table.dc_conductance*yscale, '.-', color=color)

def bruhat18_fig2c_interpolate(dm, **kwargs):
    """
    Plot lines of G vs Vdc for different Vac as in PRB 98.075121 Fig. 2 c/d.
    A slider can be used to manipulate the calibration of Vac.
    """
    for name in ("vdc", "vac", "omega"):
        kwargs.pop(name, None)
    tk_mueV = 28.2  # WRONG! Kondo temperature in μeV, defined by G(V=Tk)=e²/h
    tkrg_mueV = tk_mueV / TK_VOLTAGE # Kondo temperature in μeV defined as integration constant in RTRG
    return plot_calibration_lines(
            dm,
            omega = 5.8206184,
            vac = np.array([20, 40, 60, 80, 100, 120, 140, 160, 180, 220, 300]) / tkrg_mueV,
            colors = ('#000000', '#7f7f7f', '#a64f00', '#dc2121', '#f07d2e', '#ffa600', '#54a800', '#00a8ff', '#0000ff', '#7f00ff', '#cf00f8'),
            xscale = 1e-3 * tkrg_mueV,
            yscale = np.pi,
            **kwargs,
            )

def bruhat18_fig2d_interpolate(dm, **kwargs):
    """
    Plot lines of G vs Vdc for different Vac as in PRB 98.075121 Fig. 2 c/d.
    A slider can be used to manipulate the calibration of Vac.
    """
    for name in ("vdc", "vac", "omega"):
        kwargs.pop(name, None)
    tk_mueV = 28.2  # WRONG! Kondo temperature in μeV, defined by G(V=Tk)=e²/h
    tkrg_mueV = tk_mueV / TK_VOLTAGE # Kondo temperature in μeV defined as integration constant in RTRG
    return plot_calibration_lines(
            dm,
            omega = 9.2159791,
            vac = np.array([20, 40, 60, 80, 100, 120, 140, 160, 180, 220, 300]) / tkrg_mueV,
            colors = ('#000000', '#7f7f7f', '#a64f00', '#dc2121', '#f07d2e', '#ffa600', '#54a800', '#00a8ff', '#0000ff', '#7f00ff', '#cf00f8'),
            xscale = 1e-3 * tkrg_mueV,
            yscale = np.pi,
            **kwargs,
            )

def kogan04_calibrate(dm, **kwargs):
    """
    Plot lines of G vs Vdc for different Vac as in PRB 98.075121 Fig. 2 c/d.
    A slider can be used to manipulate the calibration of Vac.

    One good parameter choice:
        Vdc shift = 0.0177
        G scale = 10^{-0.947} = 0.113
        G shift = 0.0268
        Vac scale = 1.4
    """
    for name in ("vdc", "vac", "omega"):
        kwargs.pop(name, None)
    tk_mK = 300
    tk_meV = (tk_mK * sc.k / sc.eV)
    tkrg_mueV = 1e3 * (tk_mK * sc.k / sc.eV) / TK_VOLTAGE
    #f_ghz = 13.47
    #omega = (sc.h * f_ghz * 1e9) / (1e-3*sc.k*tk_mK) * TK_VOLTAGE
    #omega = round(omega, 4)
    colors = ("red", "green", "blue", "orange", "violet")
    #omega = 7.1271
    #omega = 3.5
    omega = 3.62
    fig1, ax1 = adjust_background(
            dm,
            omega = omega,
            #vac_arr = np.array([0.72881, 1.13091, 1.50788, 1.6838, 3.61891])*omega,
            vac_arr = np.array([0.73922, 1.14706, 1.52942, 1.70785, 3.6706])*omega,
            xscale = 1/(omega * sc.eV/(13.47e12*sc.h)),
            colors = colors,
            include_Ga = True,
            integral_method = -15,
            **kwargs,
            )
    fig2, ax2 = plot_calibration_lines(
            dm,
            omega = omega,
            #omega = 7.418,
            vac = np.array([29, 45, 60, 67, 144]) * omega * sc.eV/(13.47e15*sc.h),
            colors = colors,
            xscale = 1/(omega * sc.eV/(13.47e12*sc.h)),
            yscale = np.pi,
            calibrate_min = 1,
            calibrate_max = 1.6,
            include_Ga = True,
            integral_method = -15,
            **kwargs,
            )

    shift_Vdc = 0.017
    for i, trace in enumerate((0, 10, 21, 26, 40)):
        data = np.genfromtxt(settings.BASEPATH + '/../exp_data/d764n766_didv_trace%d++.txt'%trace, skip_header=1)
        vdc_exp = data[:,1] # in mV
        g_exp = data[:,0] / 2 # in 2e²/h = 1/π
        # Plot original data
        ax1.plot(vdc_exp, g_exp, '.', markersize=1, color=colors[i])
        ax2.plot(vdc_exp, g_exp, '.', markersize=1, color=colors[i])
        # Plot symmetrized data
        sorting = np.argsort(vdc_exp)
        vdc_exp = vdc_exp[sorting]
        g_exp = g_exp[sorting]
        split = np.searchsorted(vdc_exp, shift_Vdc)
        truncate = 2*split - vdc_exp.size
        assert truncate > 0
        g_mean = (g_exp[truncate:] + g_exp[:truncate-1:-1]) / 2
        spline = BSpline(*splrep(vdc_exp[truncate:], g_mean, k=3, s=2e-4 if trace==40 else 6e-5), extrapolate=False)
        ax1.plot(vdc_exp[truncate:], spline(vdc_exp[truncate:]), color=colors[i], alpha=0.5, zorder=10)
        ax2.plot(vdc_exp[truncate:], spline(vdc_exp[truncate:]), color=colors[i], alpha=0.5, zorder=10)

    return (fig1, fig2), (ax1, ax2)

def bruhat18_calibrate(dm, frequency_GHz=12, **kwargs):
    """
    TODO
    """
    for name in ("vdc", "vac", "omega"):
        kwargs.pop(name, None)
    #tk_mueV = 32.7
    #tkrg_mueV = tk_mueV / TK_GLAZMAN
    tkrg_mueV = 19
    #tkrg_GHz = tk_mueV * 1e-15*sc.eV/sc.h / TK_GLAZMAN
    tkrg_GHz = tkrg_mueV * 1e-15*sc.eV/sc.h
    omega = round(frequency_GHz / tkrg_GHz, 2)
    print(omega)
    #f_ghz = 13.47
    #omega = (sc.h * f_ghz * 1e9) / (1e-3*sc.k*tk_mK) * TK_VOLTAGE
    #omega = round(omega, 4)
    vac_muV = np.array([20, 40, 60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 300])
    colors = ("black", "gray", "brown", "red", "#f08000", "#c0a000", "green", "cyan", "blue", "violet", "magenta", "black", "gray")

    fig, ax = plot_calibration_lines(
            dm,
            omega = omega,
            #vac = vac_muV / tkrg_mueV,
            vac = vac_muV * omega * sc.eV/(frequency_GHz*1e15*sc.h),
            colors = colors,
            xscale = 1e-3 * tkrg_mueV,
            yscale = np.pi,
            calibrate_min = 0.667,
            calibrate_max = 1.5,
            include_Ga = True,
            integral_method = -15,
            **kwargs,
            )

    shift_Vdc = 0
    data = np.genfromtxt(settings.BASEPATH + f"/../exp_data/KondoAC_Freq{frequency_GHz:d}GHz.dat", names=True)
    vdc_exp = data["Vsd_mV"]

    for i, vac in enumerate(vac_muV):
        try:
            g_exp = data[f"Vac{vac:d}"]
        except:
            continue
        # Plot original data
        ax.plot(vdc_exp, g_exp, '.', markersize=1, color=colors[i])
        # Plot symmetrized data
        sorting = np.argsort(vdc_exp)
        vdc_exp = vdc_exp[sorting]
        g_exp = g_exp[sorting]
        split = np.searchsorted(vdc_exp, shift_Vdc)
        truncate = max(2*split - vdc_exp.size, 0)
        g_mean = (g_exp[truncate:] + g_exp[:truncate-1:-1]) / 2 if truncate > 0 else g_exp
        spline = BSpline(*splrep(vdc_exp[truncate:], g_mean, k=3, s=1e-3), extrapolate=False)
        ax.plot(vdc_exp[truncate:], spline(vdc_exp[truncate:]), color=colors[i], alpha=0.5, zorder=10)

    return fig, ax

def bruhat18_fig3a(dm, omega_res=250, vac_res=250, **kwargs):
    """
    Best parameters:
    --solver_tol_rel=1e-8
    --solver_tol_abs=1e-10
    --d=1e9
    --voltage_branches=0
    """
    for name in ("omega", "vdc", "vac"):
        kwargs.pop(name, None)
    results = dm.list(vdc=0, **kwargs)
    results = results.loc[(results.vac < 15) & np.isfinite(results.dc_conductance) & (results.method != "mu")]

    omega_arr = np.linspace(0.8, 10, omega_res) # units of Tkrg
    vac_arr = np.linspace(0.1, 12, vac_res) # units of Tkrg

    #tk_mueV = 28.2  # Kondo temperature in μeV, defined by G(V=Tk)=e²/h
    #tkrg_mueV = tk_mueV / TK_VOLTAGE # Kondo temperature in μeV defined as integration constant in RTRG

    fig, ax = plt.subplots()
    ax.set_xlabel("Ω (Tk)")
    ax.set_ylabel("Vac (Tk)")
    #gdc_tck = bisplrep(results.omega, results.vac, results.dc_conductance, s=1e-4, kx=3, ky=3)
    #gdc_interp = bisplev(omega_arr, vac_arr, gdc_tck).T

    omega_mesh, vac_mesh = np.meshgrid(omega_arr, vac_arr)
    gdc_interp = griddata(
            (results.omega, results.vac),
            results.dc_conductance,
            (omega_mesh, vac_mesh),
            method="cubic")

    img = ax.imshow(
            np.pi*gdc_interp,
            extent = ((0.8-4.6/omega_res)/TK_VOLTAGE, (10+4.6/omega_res)/TK_VOLTAGE, (0.1-5.95/vac_res)/TK_VOLTAGE, (12+5.95/vac_res)/TK_VOLTAGE),
            aspect = 'auto',
            cmap = plt.cm.jet,
            origin='lower')
    ax.scatter(results.omega/TK_VOLTAGE, results.vac/TK_VOLTAGE, c=np.pi*results.dc_conductance, cmap=img.cmap, norm=img.norm, s=10)
    fig.colorbar(img, ax=ax)
    return fig, ax

def adjust_background(
        dm,
        omega,
        xscale = 1,
        vac_arr = [],
        colors = [],
        **kwargs):
    """
    Plot lines of G vs Vdc for different Vac. A slider can be used to
    manipulate the calibration of Vac.
    Arguments:
        dm: DataManager instance
        omega: frequency, in units of Tkrg
        vac: list or 1d array of AC voltage, units of Tkrg
        colors: list of colors, must have at least same length as vac
        xscale: Tkrg to mV, used to scale X axis (Vdc)
        vdc_max: max. value of Vdc (units of Tkrg)
        vdc_res: resolution of Vdc
        **kwargs: parameters for selecting data points
            (e.g. voltage_branches, solver_tol_rel, ...)
    """
    fig, ax = plt.subplots()
    gshift_slider_ax = fig.add_axes((0.1, 0.04, 0.8, 0.02))
    gscale_slider_ax = fig.add_axes((0.1, 0.065, 0.8, 0.02))
    vdc_shift_slider_ax = fig.add_axes((0.1, 0.09, 0.8, 0.02))
    gshift_slider = Slider(gshift_slider_ax, "G shift", 0., 0.2, 0.)
    gscale_slider = Slider(gscale_slider_ax, "log10(G scale)", -1.5, 0., 0.)
    vdc_shift_slider = Slider(vdc_shift_slider_ax, "Vdc shift", -0.03, 0.03, 0.)

    data = dm.list(omega=omega, **kwargs)

    # Interpolate
    g_arrs = []
    vdc_arrs = []
    for vac in vac_arr:
        data_sel = data[np.abs(data.vac-vac)<1e-6]
        g_arrs.append(np.pi*pd.concat((data_sel.dc_conductance, data_sel.dc_conductance)))
        vdc_arrs.append(xscale*pd.concat((data_sel.vdc, -data_sel.vdc)))

    # Create figure
    ax.set_ylabel("G (2e²/h)")
    ax.set_xlabel("Vdc")

    if colors == []:
        colors = len(vac_arr) * ['black']
    plot_g = [
            ax.plot(vdc, g, 'o', color=color)[0]
            for vdc, g, color in zip(vdc_arrs, g_arrs, colors)
        ]

    def update(trash):
        gscale = 10**gscale_slider.val
        gshift = gshift_slider.val
        for i, (vdc, g) in enumerate(zip(vdc_arrs, g_arrs)):
            plot_g[i].set_data(vdc+vdc_shift_slider.val, gshift+gscale*g)

    gshift_slider.on_changed(update)
    gscale_slider.on_changed(update)
    vdc_shift_slider.on_changed(update)

    fig.sliders = (gshift_slider, gscale_slider, vdc_shift_slider)
    return fig, ax

def plot_calibration_lines(
        dm,
        omega,
        vac,
        colors = [],
        xscale = 1,
        yscale = np.pi,
        calibrate_min = 0.8,
        calibrate_max = 1.2,
        include_i = False,
        vdc_max = 25,
        vdc_res = 200,
        **kwargs):
    """
    Plot lines of G vs Vdc for different Vac. A slider can be used to
    manipulate the calibration of Vac.
    Arguments:
        dm: DataManager instance
        omega: frequency, in units of Tkrg
        vac: list or 1d array of AC voltage, units of Tkrg
        colors: list of colors, must have at least same length as vac
        xscale: Tkrg to mV, used to scale X axis (Vdc)
        yscale: G/(e²/2) = π, used to scale Y axis (G)
        calibrate_min: start value of the calibration slider
        calibrate_max: end value of the calibration slider
        include_i: include G computed from current
        vdc_max: max. value of Vdc (units of Tkrg)
        vdc_res: resolution of Vdc
        **kwargs: parameters for selecting data points
            (e.g. voltage_branches, solver_tol_rel, ...)
    """
    results = dm.list(omega=omega, **kwargs)
    results = results.loc[(results["solver_flags"] & DataManager.SOLVER_FLAGS["simplified_initial_conditions"]) == 0]
    j = results.method == "J"
    mu = results.method == "mu"

    fig, ax = plt.subplots()

    vac_slider_ax = fig.add_axes((0.1, 0.015, 0.8, 0.02))
    gshift_slider_ax = fig.add_axes((0.1, 0.04, 0.8, 0.02))
    gscale_slider_ax = fig.add_axes((0.1, 0.065, 0.8, 0.02))
    vdc_shift_slider_ax = fig.add_axes((0.1, 0.09, 0.8, 0.02))
    vac_slider = Slider(vac_slider_ax, "vac scale", calibrate_min, calibrate_max, 1.)
    gshift_slider = Slider(gshift_slider_ax, "G shift", 0., 0.2, 0.)
    gscale_slider = Slider(gscale_slider_ax, "log10(G scale)", -1.5, 0., 0.)
    vdc_shift_slider = Slider(vdc_shift_slider_ax, "Vdc shift", -0.03, 0.03, 0.)

    # Interpolate
    vdc_arr = np.linspace(0, vdc_max, vdc_res)
    dummy_interp = lambda vac: np.nan*np.zeros_like(vdc_arr)
    try:
        gdc_mu_tck = bisplrep(results.vac[mu], results.vdc[mu], results.dc_conductance[mu], s=1e-6, kx=3, ky=3)
        gdc_g_mu_interp = lambda vac: bisplev(vac, vdc_arr, gdc_mu_tck)
    except:
        gdc_g_mu_interp = dummy_interp

    try:
        gdc_J_tck = bisplrep(results.vac[j], results.vdc[j], results.dc_conductance[j], s=2e-6, kx=3, ky=3)
        gdc_g_J_interp = lambda vac: bisplev(vac, vdc_arr, gdc_J_tck)
    except:
        gdc_g_J_interp = dummy_interp

    if include_i:
        try:
            idc_mu_tck = bisplrep(
                    np.concatenate((results.vac[mu], results.vac[mu])),
                    np.concatenate((-results.vdc[mu], results.vdc[mu])),
                    np.concatenate((-results.dc_current[mu], results.dc_current[mu])),
                    s=1e-5, kx=3, ky=3)
            gdc_i_mu_interp = lambda vac: bisplev(vac, vdc_arr, idc_mu_tck, dy=1)
        except:
            gdc_i_mu_interp = dummy_interp
        try:
            assert False
            idc_J_tck = bisplrep(
                    np.concatenate((results.vac[j], results.vac[j])),
                    np.concatenate((-results.vdc[j], results.vdc[j])),
                    np.concatenate((-results.dc_current[j], results.dc_current[j])),
                    s=1e-5, kx=3, ky=3)
            gdc_i_J_interp = lambda vac: bisplev(vac, vdc_arr, idc_J_tck, dy=1)
        except:
            gdc_i_J_interp = dummy_interp

    mirror = lambda x: np.concatenate((x[::-1], x))
    vdc_full_scaled = xscale * np.concatenate((-vdc_arr[::-1], vdc_arr))

    # Create figure
    ax.set_ylabel("G (2e²/h)")
    ax.set_xlabel("Vdc (mV)")

    vac_arr = vac
    if colors == []:
        colors = len(vac_arr) * ['black']
    plot_g_mu = [
            ax.plot(
                vdc_full_scaled,
                yscale*mirror(gdc_g_mu_interp(vac)),
                color = color,
                )[0]
            for vac, color in zip(vac_arr, colors)
        ]
    plot_g_J = [
            ax.plot(
                vdc_full_scaled,
                yscale*mirror(gdc_g_J_interp(vac)),
                color = color,
                )[0]
            for vac, color in zip(vac_arr, colors)
        ]
    if include_i:
        plot_i_mu = [
                ax.plot(
                    vdc_full_scaled,
                    yscale*mirror(gdc_i_mu_interp(vac)),
                    color = color,
                    )[0]
                for vac, color in zip(vac_arr, colors)
            ]
        plot_i_J = [
                ax.plot(
                    vdc_full_scaled,
                    yscale*mirror(gdc_i_J_interp(vac)),
                    color = color,
                    )[0]
                for vac, color in zip(vac_arr, colors)
            ]

    def update(trash):
        x = vdc_full_scaled + vdc_shift_slider.val
        gscale = yscale * 10**gscale_slider.val
        gshift = gshift_slider.val
        for i, vac in enumerate(vac_arr):
            vac *= vac_slider.val
            plot_g_mu[i].set_data(x, gshift + gscale*mirror(gdc_g_mu_interp(vac)))
            plot_g_J[i].set_data(x, gshift + gscale*mirror(gdc_g_J_interp(vac)))
            if include_i:
                plot_i_mu[i].set_data(x, gshift + gscale*mirror(gdc_i_mu_interp(vac)))
                plot_i_J[i].set_data(x, gshift + gscale*mirror(gdc_i_J_interp(vac)))
    vac_slider.on_changed(update)
    gshift_slider.on_changed(update)
    gscale_slider.on_changed(update)
    vdc_shift_slider.on_changed(update)

    fig.sliders = (vac_slider, gshift_slider, gscale_slider, vdc_shift_slider)
    return fig, ax

def kogan04(dm, **parameters):
    tk_mK = 300
    f_ghz = 13.47
    omega_mueV = f_ghz * 1e15*sc.h/sc.eV
    tkrg_mueV = 1e3 * (tk_mK * sc.k / sc.eV) / TK_GLAZMAN
    # based on Tk=300mK:
    #omega = (sc.h * f_ghz * 1e9) / (1e-3*sc.k*tk_mK) * TK_GLAZMAN
    # based on Ω=2Tk:
    #omega = 2 * TK_GLAZMAN
    #omega = round(omega, 5)
    #omega = 7.1271
    #omega = 7
    omega = 3.62
    correction = 1.42
    vac_mueV_arr = np.array([29, 45, 60, 67, 144])
    vac_omega_arr = vac_mueV_arr / omega_mueV
    vac_arr = vac_mueV_arr / tkrg_mueV
    print(f"Tkrg = {tkrg_mueV} μeV")
    print(f"omega = {omega}")
    print(f"Vac = {vac_arr}")
    print(f"Vac/omega = {vac_omega_arr}")
    print(f"corrected Vac = {correction*vac_arr}")
    print(f"corrected Vac = {(correction*vac_omega_arr).round(5)} Ω")
    vdc_mueV_max = 400
    vdc_max = vdc_mueV_max / omega_mueV * omega
    print(f"Vdc max. = {vdc_max}")


if __name__ == '__main__':
    main()