import h5py
import matplotlib.pyplot as plt
import numpy as np
import multifile
import os
import math
from matplotlib import cm
plt.rcParams['text.usetex'] = True
plt.rcParams['text.latex.preamble'] = r'\usepackage{amssymb}'

SMALL_SIZE = 8
MEDIUM_SIZE = 12
BIGGER_SIZE = 16

plt.rc('font', size=BIGGER_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=BIGGER_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=BIGGER_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=BIGGER_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=BIGGER_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=MEDIUM_SIZE)    # legend fontsize
plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title
plt.rcParams["font.family"] = "serif"
plt.rcParams["text.usetex"] = True
plt.rcParams["font.family"] = "cmr10",  # Computer Modern Roman

       
def load_objs(dat_names, Beta=20):
    objs = []
    omega0_vals = set([])
    g0_vals = set([])

    mf = multifile.MultiFile(dat_names[0] + '/')
    f = mf.open_final_file()

    momgrid = np.array(f["w_func/momgrid"])
    ffgrid = np.array(f["Sig/fgrid"])
    f.close()
    mf.close()
    
    #Beta = 10.0 #obj['Omega0'] = mf.params_file["Model"].attrs["Beta"]
    freq = lambda n: (2*n) * np.pi/Beta

    dat_names
    
    for dat_dir in dat_names:
        mf = multifile.MultiFile(dat_dir + '/')
        obj = {}

        obj['dir'] = dat_dir
        obj['U'] = mf.params_file["Model"].attrs["U"]
        obj['Vh'] = mf.params_file["Model"].attrs["g0"]
        obj['Omega0'] = mf.params_file["Model"].attrs["OMEGA0"]
        obj['g0'] = np.sqrt(obj['Vh'] * obj['Omega0']/2.0)
        obj['mu'] = mf.params_file["Model"].attrs["Mu"]
        obj['BZ_path'] = np.array(mf.params_file["Model/Special_paths/path_Gamma_X_M"])
  
        jj = 0
        for ii in range(len(obj['BZ_path'])-1):
            if obj['BZ_path'][ii-jj] == obj['BZ_path'][ii+1-jj]:
                obj['BZ_path'] = np.delete(obj['BZ_path'], ii)
                jj += 1

        
        #print(obj['BZ_path'])
            
        omega0_vals.add(obj['Omega0'])
        g0_vals.add(obj['g0'])
        obj['Vh'] = 2*obj['g0']*obj['g0']/obj['Omega0']
        U_Omega0 = lambda n: obj['Vh'] * freq(n)*freq(n) * (obj['Omega0']*obj['Omega0'] + freq(n)*freq(n)) 
        D0 = lambda n: -2*obj['Omega0']/(freq(n)*freq(n) + obj['Omega0']*obj['Omega0'])
        #print(obj['Vh'])
        if not mf.is_finished:
            f = mf.open_file_number(int(mf.scales_count))
            B = obj['U'] - obj['Vh'] 
            print( (B - np.min(-np.array(f["w_func/RE_D"]))/(4*np.pi*np.pi))/B/B  )
            #print(dat_dir)
            continue
        f = mf.open_final_file()

        obj['Sig'] = np.array(f["Sig/RE"]) + 1j * np.array(f["Sig/IM"])
        obj['delta_mu'] = f["Flow_obs"].attrs["delta_mu"]
        
        #print(f["Flow_obs/S_Wave_Susc_info"].attrs.keys())
        #print(np.array(f["Flow_obs/S_Wave_Susc_info"].keys)[0, 0])
        obj['leading_susc'] = f["Flow_obs/S_Wave_Susc_info"].attrs['Leading_name']
        obj['max_sc'] = f["Flow_obs/S_Wave_Susc_info"].attrs['RE_Max_sc']
        obj['max_d'] = f["Flow_obs/S_Wave_Susc_info"].attrs['RE_Max_d']
        obj['max_m'] = f["Flow_obs/S_Wave_Susc_info"].attrs['RE_Max_m']

        obj['max_sc_pp'] = np.max(np.array(f["Flow_obs/Postprocessing_Susc_info/RE_Susc_sc"])[0,:, 0, 0, 0, 0, 0, 0])
        obj['max_d_pp'] = np.max(np.array(f["Flow_obs/Postprocessing_Susc_info/RE_Susc_d"])[0,:, 0, 0, 0, 0, 0, 0]) #136
        obj['max_m_pp'] = np.max(np.array(f["Flow_obs/Postprocessing_Susc_info/RE_Susc_m"])[0,:, 0, 0, 0, 0, 0, 0])
        if np.array(f["Flow_obs/Postprocessing_Susc_info/RE_Susc_sc"]).shape[3] > 1:
            print("dwave!")
            obj['max_dsc_pp'] = np.max(np.array(f["Flow_obs/Postprocessing_Susc_info/RE_Susc_sc"])[0,:, 0, 1, 0, 1, 0, 0])

        obj['sc_pp'] =np.max(np.array(f["Flow_obs/Postprocessing_Susc_info/RE_Susc_sc"])[0,:, 0, 0, 0, 0, 0, 0])
        obj['d_pp'] = np.max(np.array(f["Flow_obs/Postprocessing_Susc_info/RE_Susc_d"])[0,:, 0, 0, 0, 0, 0, 0]) #136
        obj['m_pp'] = np.max(np.array(f["Flow_obs/Postprocessing_Susc_info/RE_Susc_m"])[0,:, 0, 0, 0, 0, 0, 0])
        
        obj['leading_w'] = f["Flow_obs/w_info"].attrs['Leading_name']
        obj['max_w_sc'] = f["Flow_obs/w_info"].attrs['RE_Max_sc']
        obj['max_w_d'] = f["Flow_obs/w_info"].attrs['RE_Max_d']
        obj['max_w_m'] = f["Flow_obs/w_info"].attrs['RE_Max_m']
        old = np.array(f["Flow_obs/S_Wave_Susc_info/RE_Susc_d"])[0, :]
        obj["diverged"] = mf.is_diverged
        wd = np.array(f["w_func/RE_D"])[:, :]/(4*np.pi*np.pi)
        wd = wd[int((wd.shape[0]-1)/2), :]
        obj['susc_d_pp'] = np.array(f["Flow_obs/Postprocessing_Susc_info/RE_Susc_d"])[0,:, 0, 0, 0, 0, 0, 0]
        bose_bare_V = np.zeros(wd.shape)
        D0_vals =  np.zeros(wd.shape)
        #print(bare_V.shape)
        
        for Q_ in range(bose_bare_V.shape[0]):
            bose_bare_V[Q_] = obj['U'] - obj['Vh'] + 2.0*U_Omega0(0)
            D0_vals[Q_] = D0(0)
                
        bose_bare_V_static = bose_bare_V[ :]
                
        obj['susc_d'] = ( bose_bare_V_static+wd)/bose_bare_V_static/bose_bare_V_static
        wd_pp = obj['susc_d_pp'] * bose_bare_V_static  * bose_bare_V_static - bose_bare_V_static
        obj['phonon_selfenergy'] = -obj['g0']*obj['g0']*obj['susc_d_pp']/(1 - (-2/obj['Omega0'])*obj['g0']*obj['g0']*obj['susc_d_pp'])
        
                                                                           
        #print(obj['susc_d'][int((bose_bare_V.shape[0]-1)/2), 136] - old[136])#obj['susc_d'])
        #print(wd[int((bose_bare_V.shape[0]-1)/2), 136], bose_bare_V[int((bose_bare_V.shape[0]-1)/2), 136] )#obj['susc_d'])
        f.close()
        objs.append(obj)
        
    return (objs, omega0_vals, g0_vals, momgrid, ffgrid)
        
        

from scipy.optimize import curve_fit

def linear(x, c):
    return  -x + c

def plot_d_susc_vs_g0(objs, title_ = "", label_ = ""):
    objs_ = sorted(objs, key=lambda o: o["Vh"] )
    plt.grid(True)
    #plt.ylim([2, 5])
    #plt.xlim([0.0, 3.5])
    plt.title(r"$\Omega = "+str(objs[0]['Omega0'])+r"$")
    plt.xlabel(r"$2g_0^2/\Omega_0$")
    plt.ylabel(r"$\chi_D^{-1}(\pi, \pi)$")
    #plt.xscale('log')
    chi_vals = []
    Vh_vals = []
    for o in objs_:
        if abs(o["Vh"]- o["U"]) < 0.001:
            continue
        if o["max_d"] != o["max_d"] or math.isinf(o["max_d"]):
            continue
        if o["max_d"] == 0:
            continue
        else:
            chi_vals.append(1.0/o["max_d"])
        Vh_vals.append(o["Vh"])
    #print(chi_vals)
    popt, pcov = curve_fit(linear, Vh_vals, chi_vals, p0=[0])
    print(round(pcov[0][0], 4))
    
    fitted_y_vals = [linear(x, popt[0]) for x in Vh_vals]
    
    plt.plot(Vh_vals, chi_vals, "*-", label=r"$U$ = "+str(objs[0]['U']) + label_)
    plt.plot(Vh_vals, fitted_y_vals, "-.")
    

import matplotlib.pyplot as plt
import matplotlib.tri as tri
import numpy as np
import numpy.random as rnd
import matplotlib.mlab as mlab
import time
from scipy.interpolate import griddata
import math

def plot_leading_fluctuations_all(objs, title_ = ""):
    fig, axs = plt.subplots(1, 3, figsize=(13, 8))

    axs[0].set_ylabel('U')
    axs[1].set_title(title_, pad=20, fontsize=20)
    for i in [0, 1, 2]:
        axs[i].set_xlim(0, 3)
        axs[i].set_ylim(0, 3)
        axs[i].set_aspect('equal')
        axs[i].set_xlabel('Vh')
        
    XXsc = []
    YYsc = []
    ZZsc = []


    XXd = []
    YYd = []
    ZZd = []
    

    XXm = []
    YYm = []
    ZZm = []

        

    ZZwsc = []
    ZZwd = []
    ZZwm = []

    for obj in objs:    
        
        if ( (not math.isnan(obj["max_sc"])) and math.isfinite(obj["max_sc"])):
            XXsc.append(obj["Vh"])
            YYsc.append(obj["U"])
            ZZsc.append(round(obj['max_sc'], 2))
            ZZwsc.append(obj['max_w_sc'])

        if ( (not math.isnan(obj["max_d"])) and math.isfinite(obj["max_d"])):
            XXd.append(obj["Vh"])
            YYd.append(obj["U"])
            ZZd.append(round(obj['max_d'], 2))
            ZZwd.append(obj['max_w_d'])

        
        if ( (not math.isnan(obj["max_m"])) and math.isfinite(obj["max_m"])):
            XXm.append(obj["Vh"])
            YYm.append(obj["U"])
            ZZm.append(round(obj['max_m'], 2))
            ZZwm.append(obj['max_w_m'])
        
    
            
    the_levels=  sorted(list(set(ZZsc))) #[0, 0.8, 1.0, 1.4, 1.9, 2.5, 2.7, 3.5, 4.4, 5.6, 7.5, 8.9, 14, 17, 20, 30.4, 55, 64 ,100, 210, 500]

    the_levelsd=  sorted(list(set(ZZd))) #[0, 0.8, 1.0, 1.4, 1.9, 2.5, 2.7, 3.5, 4.4, 5.6, 7.5, 8.9, 14, 17, 20, 30.4, 55, 64 ,100, 210, 500]
    the_levelsm=  sorted(list(set(ZZm))) #[0, 0.8, 1.0, 1.4, 1.9, 2.5, 2.7, 3.5, 4.4, 5.6, 7.5, 8.9, 14, 17, 20, 30.4, 55, 64 ,100, 210, 500]


    
    colorsr = [(.9 +(0.9-lam)/9.0, .9 + (0.9-lam)/9.0 - lam*lam, 0.9+(0.9-lam)/9.0 - lam*lam) for lam in np.linspace(0, 0.9, len(the_levels))]
    colorsg = [(.9+(0.9-lam)/9.0 - lam*lam, .9+(0.9-lam)/9.0, 0.9+(0.9-lam)/9.0 - lam*lam) for lam in np.linspace(0, 0.9, len(the_levelsd))]
    colorsb = [(.9+(0.9-lam)/9.0 - lam*lam, .9+(0.9-lam)/9.0 - lam*lam, .9+(0.9-lam)/9.0) for lam in np.linspace(0, 0.9, len(the_levelsm))]

    #ZZsc = [min(100, x) for x in ZZsc]
    #ZZd = [min(100, x) for x in ZZd]
    #ZZm = [min(100, x) for x in ZZm]

    if len(XXsc) > 1:
        XXsc_, YYsc_ = np.meshgrid(np.linspace(min(XXsc), max(XXsc), 100), np.linspace(min(YYsc), max(YYsc), 100))
        ZZsc_ = griddata((XXsc, YYsc), ZZsc, (XXsc_, YYsc_), method='nearest')
        contourfr = axs[0].contourf(XXsc_, YYsc_, ZZsc_,levels = the_levels, colors = colorsr, extend='max', alpha=0.5)
        #contourfr = axs[0].tricontourf(XXsc, YYsc, ZZsc, levels = the_levels, colors = colorsr, extend='max')
    if len(XXd) > 1:
        XXd_, YYd_ = np.meshgrid(np.linspace(min(XXd), max(XXd), 100), np.linspace(min(YYd), max(YYd), 100))
        ZZd_ = griddata((XXd, YYd), ZZd, (XXd_, YYd_), method='nearest')
        contourfg = axs[1].contourf(XXd_, YYd_, ZZd_, levels = the_levelsd, colors = colorsg, extend='max', alpha=0.5)
    if len(XXm) > 1:
            # Generate grid points for contour plot                                                                                                                                                                                           
        XXm_, YYm_ = np.meshgrid(np.linspace(min(XXm), max(XXm), 100), np.linspace(min(YYm), max(YYm), 100))
        ZZm_ = griddata((XXm, YYm), ZZm, (XXm_, YYm_), method='nearest')
        contourfb = axs[2].contourf(XXm_, YYm_, ZZm_, levels = the_levelsm, colors = colorsb, extend='max', alpha=0.5 )
        
        

    #cbar = fig.colorbar(contourfg, shrink=0.7, pad=-0.1, ticks = [])
    #cbar = fig.colorbar(contourfb, shrink=0.7,pad=-0.1, ticks = [])
    #cbar = fig.colorbar(contourfr, shrink=0.7, pad=0.02, ticks = [])
    
    if len(XXd) > 1:
        cbar = fig.colorbar(contourfg, shrink=0.4, ax=axs[1])
    if len(XXm) > 1:
        cbar = fig.colorbar(contourfb, shrink=0.4, ax=axs[2])
    if len(XXsc) > 1:
        cbar = fig.colorbar(contourfr, shrink=0.4, ax=axs[0])

    
    plt.show()

def get_renormalised_omegas(objs):
    objs_ = sorted(objs, key=lambda o: o["Vh"] )
    pure_chi_d_electronic_vals = []
    chi_d_vals = []
    Vh_vals = []
    U_vals = []
    Omega0_vals = []
    phonon_selfenergy_vals = []
    for o in objs_:
        
        #if o["diverged"]:
            #print("susc_d")
        #    continue
        #if abs(o["U"] - o["Vh"]) < 0.1:
        #    continue
        if abs(o["Vh"]) < 0.01:
            phonon_selfenergy_vals.append(np.zeros(o["phonon_selfenergy"].shape))
        else:
            phonon_selfenergy_vals.append(o["phonon_selfenergy"][:])
        if "pure_susc_d" in o.keys():
            pure_chi_d_electronic_vals.append(o["pure_susc_d"][:])
        chi_d_vals.append(o["susc_d"][:])
        Vh_vals.append(o["Vh"])
        U_vals.append(o["U"])
        Omega0_vals.append(o["Omega0"])
        
        #if (1.0 - 2.0*(phonon_selfenergy_vals[-1][136]/Omega0_vals[-1]))<0:
        #    break
    
    Vh_vals = np.array(Vh_vals)
    Omega0_vals = np.array(Omega0_vals)
    pure_chi_d_electronic_vals = np.array(pure_chi_d_electronic_vals)
    phonon_selfenergy_vals = np.array(phonon_selfenergy_vals)
    
    #print(phonon_selfenergy_vals.shape)
    #sqrt_vals = -Vh_vals*chi_d_vals+ 1
    sqrt_vals = 1.0 + 2.0*(phonon_selfenergy_vals[:, 136]/Omega0_vals)
    dispersion_vals = np.array([(0 if e < 0 else e) for e in sqrt_vals])
    
    
    phonon_selfenergy_vals_over_BZ = np.zeros((phonon_selfenergy_vals.shape[0], len(objs_[0]['BZ_path'])))

    dispersion_vals_over_BZ = np.zeros((phonon_selfenergy_vals.shape[0], len(objs_[0]['BZ_path'])))

    
    for VHidx in range((phonon_selfenergy_vals.shape[0])):
        for i in range(len(objs_[0]['BZ_path'])):
            momidx = objs_[0]['BZ_path'][i]
            phonon_se_val = phonon_selfenergy_vals[VHidx, momidx]
            phonon_selfenergy_vals_over_BZ[VHidx, i] = phonon_se_val
            dispersion_vals_over_BZ[VHidx, i] = np.sqrt(1.0 + 2.0*phonon_se_val/Omega0_vals[0])

    #zero_crossing = [Vh_vals[-1], None] if dispersion_vals[-1] == 0 else None
    zero_crossing = None
    if (dispersion_vals[-1] == 0 and len(dispersion_vals) > 3):
        #print("hohohoho")
        #print(len(Vh_vals[:-1]))
        zero_crossing = find_zero_crossing(Vh_vals[:-1], dispersion_vals[:-1], degree=min(3, len(Vh_vals)-1) )
        #print(zero_crossing)
        Vh_vals[-1] = zero_crossing[0]
        dispersion_vals[-1] = 0
    
    
    
        
    return Vh_vals, np.sqrt(dispersion_vals), dispersion_vals_over_BZ, phonon_selfenergy_vals_over_BZ, zero_crossing, U_vals



def get_renormalised_omegas2(objs):
    print("howdy")
    objs_ = sorted(objs, key=lambda o: o["Vh"] )
    pure_chi_d_electronic_vals = []
    chi_d_vals = []
    Vh_vals = []
    U_vals = []
    Omega0_vals = []
    phonon_selfenergy_vals = []
    for o in objs_:
        
        if o["diverged"]:
            #print("susc_d")
            continue
        if abs(o["U"] - o["Vh"]) < 0.1:
            continue
        if abs(o["Vh"]) < 0.01:
            phonon_selfenergy_vals.append(np.zeros(o["phonon_selfenergy"].shape))
        else:
            phonon_selfenergy_vals.append(o["phonon_selfenergy"][:])
        if "pure_susc_d" in o.keys():
            pure_chi_d_electronic_vals.append(o["pure_susc_d"][:])
        chi_d_vals.append(o["susc_d"][:])
        Vh_vals.append(o["Vh"])
        U_vals.append(o["U"])
        Omega0_vals.append(o["Omega0"])
        
        if (1.0 - 2.0*(phonon_selfenergy_vals[-1][136]/Omega0_vals[-1]))<0:
            break
    
    Vh_vals = np.array(Vh_vals)
    Omega0_vals = np.array(Omega0_vals)
    pure_chi_d_electronic_vals = np.array(pure_chi_d_electronic_vals)
    phonon_selfenergy_vals = np.array(phonon_selfenergy_vals)
    
    print(phonon_selfenergy_vals.shape)
    #sqrt_vals = -Vh_vals*chi_d_vals+ 1
    sqrt_vals = 1.0 + 2.0*(phonon_selfenergy_vals[:, 136]/Omega0_vals)
    dispersion_vals = np.array([(0 if e < 0 else e) for e in sqrt_vals])
    
    
    phonon_selfenergy_vals_over_BZ = np.zeros((phonon_selfenergy_vals.shape[0], len(objs_[0]['BZ_path'])))

    dispersion_vals_over_BZ = np.zeros((phonon_selfenergy_vals.shape[0], len(objs_[0]['BZ_path'])))

    
    for VHidx in range((phonon_selfenergy_vals.shape[0])):
        for i in range(len(objs_[0]['BZ_path'])):
            momidx = objs_[0]['BZ_path'][i]
            phonon_se_val = phonon_selfenergy_vals[VHidx, momidx]
            phonon_selfenergy_vals_over_BZ[VHidx, i] = phonon_se_val
            dispersion_vals_over_BZ[VHidx, i] = np.sqrt(1.0 + 2.0*phonon_se_val/Omega0_vals[0])

    #zero_crossing = [Vh_vals[-1], None] if dispersion_vals[-1] == 0 else None
    zero_crossing = None
    if (dispersion_vals[-1] == 0 and len(dispersion_vals) > 3):
        #print("hohohoho")
        #print(len(Vh_vals[:-1]))
        zero_crossing = find_zero_crossing(Vh_vals[:-1], dispersion_vals[:-1], degree=min(3, len(Vh_vals)-1) )
        #print(zero_crossing)
        Vh_vals[-1] = zero_crossing[0]
        dispersion_vals[-1] = 0
    
    
    
        
    return Vh_vals, np.sqrt(dispersion_vals), dispersion_vals_over_BZ, phonon_selfenergy_vals_over_BZ, zero_crossing, U_vals
    

#xs, ys =get_renormalised_omegas(objs_vanhove_halffilling_omega0p01_Us[1.0] )
#plt.plot(xs, ys/0.01, "*-")

plt.xlim(0, 3)
plt.ylim(-1.5, 1.5)

plt.ylabel(r"$\Omega_{eff}(\pi, \pi)/\Omega_0$")
plt.xlabel(r"$V_H$")

for Uval in [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 2.5, 2.75]:
    xs, ys, disps, ses,_ =get_renormalised_omegas(objs_beta20_halffilling_omega1_Us[Uval])
    plt.plot(xs, ys, "*-", label= r"$U$= {0}t, $\Omega_0 = 1t$".format(Uval))




plt.legend()

#xs, ys =get_renormalised_omegas(objs_vanhove_halffilling_omega0p01_Us[1.0] )
#plt.plot(xs, ys/0.01, "*-")

plt.xlim(0, 3)
plt.ylim(0, 2)

plt.ylabel(r"$\Omega_{eff}(\pi, \pi)/\Omega_0$")
plt.xlabel(r"$V_H$")


for Uval in [0.5, 1.0, 1.5, 2.0, 2.5]:
    xs, ys, disps, ses,_ =get_renormalised_omegas(objs_vanhove_halffilling_omega01_Us[Uval])
    plt.plot(xs, ys, "*-", label= r"$U$= {0}t, $\Omega_0 = 1t$".format(Uval))


plt.legend()
print(_)

def plot_leading_fluctuations_beta20(objs, Us_list_,  title_ = "", crystal_instability= None, ax_ = None, fig_ = None, no_xticks = False, no_yticks = False):
    #fig, ax = plt.subplots()
    if ax_ is None:
        ax_ = plt.gca()
    if fig_ is None:
        fig_ = plt.gcf()
    ax = ax_
    fig = fig_
    num = len(Us_list_)
    dU = Us_list[1]
    print(dU)
    
    ax.set_xlim(-0.5, num-0.5)
    ax.set_ylim(-0.5, num-0.5)
    ax.set_aspect('equal')

    
    XXsc = []
    YYsc = []
    ZZsc = []
    
    XXall = []
    YYall = []
    ZZall = []

    XXd = []
    YYd = []
    ZZd = []
    

    ZZdall = []

    XXm = []
    YYm = []
    ZZm = []

    ZZmall = [] 
        

    ZZwsc = []
    ZZwd = []
    ZZwm = []

    _pp_str = "_pp"
    
    vals_phasediagram = np.ones( (int(num), int(num))) * 130
    #dbg_phasediagram = np.zeroes
    ax.text(3.0/dU+0.05, 0/dU+0.05, r"$\square$", size=23, rotation=0,
             ha="center", va="center"
             )
    ax.text(3.0/dU+0.05, 0.5/dU+0.05, r"$\square$", size=23, rotation=0,
             ha="center", va="center"
             )
    for obj in objs: 
        if ((obj["U"]+0.001) % 0.5) > 1e-2 or ((obj["Vh"]+0.001) % 0.5) > 1e-2:
            continue
        color = np.array([0.0, 0.0, 0.0])
        leading = obj["leading_susc"]
        is_divergent =obj["diverged"] 
        if "sc" in leading:
            if np.abs(obj["max_sc"+_pp_str] - obj["max_d"+_pp_str])/(obj["max_d"+_pp_str] + obj["max_sc"+_pp_str]) < 0.05:
                leading += ", d"
            if np.abs(obj["max_sc"+_pp_str] - obj["max_m"+_pp_str])/(obj["max_m"+_pp_str] + obj["max_sc"+_pp_str]) < 0.05:
                leading += ", m"
        if "d" in leading:
            if np.abs(obj["max_sc"+_pp_str] - obj["max_d"+_pp_str])/(obj["max_d"+_pp_str] + obj["max_sc"+_pp_str]) < 0.05:
                leading += ", sc"
            if np.abs(obj["max_d"+_pp_str] - obj["max_m"+_pp_str])/(obj["max_m"+_pp_str] + obj["max_d"+_pp_str]) < 0.05:
                leading += ", m"
        if "m" in leading:
            if np.abs(obj["max_sc"+_pp_str] - obj["max_m"+_pp_str])/(obj["max_m"+_pp_str] + obj["max_sc"+_pp_str]) < 0.05:
                leading += ", sc"
            if np.abs(obj["max_d"+_pp_str] - obj["max_m"+_pp_str])/(obj["max_m"+_pp_str] + obj["max_d"+_pp_str]) < 0.05:
                leading += ", d"
        
        if False and is_divergent:
                ax.text(obj["Vh"]/dU, obj["U"]/dU, r"$\cdot$", size=12, rotation=0,
             ha="center", va="center"
             )
        if "sc" in leading:
            if np.abs(obj["max_sc"+_pp_str] - obj["max_d"+_pp_str])/(obj["max_d"+_pp_str] + obj["max_sc"+_pp_str]) < 0.1:
                leading += ", d"
            color += np.array([1.0, 0.0, 0.0])
            ax.text(obj["Vh"]/dU+0.05, obj["U"]/dU+0.05, r"$\bigcirc$", size=23, rotation=0,
             ha="center", va="center"
             )
            
        if "d" in leading:
            ax.text(obj["Vh"]/dU+0.05, obj["U"]/dU+0.05, r"$\square$", size=23, rotation=0,
             ha="center", va="center"
             )
            
            color += np.array([0.0, 1.0, 0.0])
        if "m" in leading:
            ax.text(obj["Vh"]/dU+0.05, obj["U"]/dU+0.05, r"$\triangle$", size=23, rotation=0,
         ha="center", va="center",
         )
            color += np.array([0.0, 0.0, 1.0])


        if (obj['max_sc'+_pp_str] > 0 and "sc" in obj["leading_susc"]):
            XXsc.append(obj["Vh"])
            YYsc.append(obj["U"])
            ZZsc.append(obj['max_sc'+_pp_str])
            ZZwsc.append(obj['max_w_sc'])

        if (obj['max_d'+_pp_str] > 0 and "d" in obj["leading_susc"]):
            XXd.append(obj["Vh"])
            YYd.append(obj["U"])
            ZZd.append(obj['max_d'+_pp_str])
            ZZwd.append(obj['max_w_d'])


        if (obj['max_m'+_pp_str] > 0 and "m" in obj["leading_susc"]):
            XXm.append(obj["Vh"])
            YYm.append(obj["U"])
            ZZm.append(obj['max_m'+_pp_str])
            ZZwm.append(obj['max_w_m'])
        
        max_coupling = max(obj['max_m'+_pp_str], max(obj['max_sc'+_pp_str], obj['max_d'+_pp_str]))
        
        if max_coupling == max_coupling:
            vals_phasediagram[int((obj["U"]/dU + 0.01)), int((obj["Vh"]/dU) + 0.01) ] = max(obj['max_m'+_pp_str], max(obj['max_sc'+_pp_str], obj['max_d'+_pp_str]))
        
    the_levels=  sorted(list(set(ZZsc))) #[0, 0.8, 1.0, 1.4, 1.9, 2.5, 2.7, 3.5, 4.4, 5.6, 7.5, 8.9, 14, 17, 20, 30.4, 55, 64 ,100, 210, 500]

    the_levelsd=  sorted(list(set(ZZd))) #[0, 0.8, 1.0, 1.4, 1.9, 2.5, 2.7, 3.5, 4.4, 5.6, 7.5, 8.9, 14, 17, 20, 30.4, 55, 64 ,100, 210, 500]
    the_levelsm=  sorted(list(set(ZZm))) #[0, 0.8, 1.0, 1.4, 1.9, 2.5, 2.7, 3.5, 4.4, 5.6, 7.5, 8.9, 14, 17, 20, 30.4, 55, 64 ,100, 210, 500]


    
    colorsr = [(.9 +(0.9-lam)/9.0, .9 + (0.9-lam)/9.0 - lam*lam, 0.9+(0.9-lam)/9.0 - lam*lam) for lam in np.linspace(0, 0.9, len(the_levels))]
    colorsg = [(.9+(0.9-lam)/9.0 - lam*lam, .9+(0.9-lam)/9.0, 0.9+(0.9-lam)/9.0 - lam*lam) for lam in np.linspace(0, 0.9, len(the_levelsd))]
    colorsb = [(.9+(0.9-lam)/9.0 - lam*lam, .9+(0.9-lam)/9.0 - lam*lam, .9+(0.9-lam)/9.0) for lam in np.linspace(0, 0.9, len(the_levelsm))]

    vals_phasediagram[vals_phasediagram > 130.0] = 130.0

    #for i in range(vals_phasediagram.shape[0]):
    #    dx = +1 if i+1 < vals_phasediagram.shape[0] else -1
    #    dy = +1 if i+1 < vals_phasediagram.shape[0] else -1
    #    vals_phasediagram[i, i] = dU*(vals_phasediagram[i, i+dy] + vals_phasediagram[i+dx, i])
    
    ax.set_title(title_, pad=10, fontsize=20, loc='left')
    
    if not no_xticks:
        ax.set_xlabel(r'$V_H/t$')
        ax_.set_xticks(np.linspace(0, num-1, 4), np.linspace(0, 3, 4))  # Relabel x-axis from 0 to 3
    else:
        ax_.set_xticks([])

        
    if not no_yticks:
        ax_.set_yticks(np.linspace(0, num-1, 4), np.linspace(0, 3, 4))  # Relabel y-axis from 0 to 3
        ax.set_ylabel(r'$U/t$')
    else:
        ax_.set_yticks([])
    #print(vals_phasediagram)
    
    clx = ax_.imshow(np.log(vals_phasediagram+1.0)/np.log(10), cmap='coolwarm', interpolation="nearest", vmax=np.log(130+1.0)/np.log(10))#interpolation='hamming') #hermit, hanning, hamming, byssel
    #fig_.colorbar(clx, ax = ax_, label=r'$log(1+max\chi)$')
    return clx
    #plt.show()

def plot_suscs_vs_Vh(chann, objs__, desired_U, title_ = "", label_ = "", max_Vh = 1e10, min_Vh = -1, linestyle_ = "-", ax = None, hide_x_labels = False, no_labels = False):
    objs_ = []
    for obj in objs__:
        if (abs(obj["U"] - desired_U))<1e-4:
            objs_.append(obj)

    objs = sorted(objs_, key=lambda o: o["Vh"] )
    #plt.grid(True)
    if ax == None:
        ax = plt.gca()
    #ax.set_title(r"$\omega_0 = "+str(int(objs[0]['Omega0']))+r"t$", y=1.0, fontsize=20)
    if not hide_x_labels:
        ax.set_xlabel(r"$V_H/t$")#(r"$2g_0^2/\Omega_0$")
    else:
        ax.set_xlabel(r"")
        ax.set_xticks([])  # Remove tick marks
        ax.set_xticklabels([]) 
        
    ax.set_ylabel(r"$\chi^{-1}/t$")
    #plt.xscale('log')
    chi_vals = []
    Vh_vals = []
    symbols = {"sc": "o", "d": "s"}
    colours = {"sc": "red", "d": "green"}
    for o in objs:
        if o["Vh"] > max_Vh:
            continue
        #if abs(o["Vh"]- o["U"]) < 0.001:
        #    continue
        #if o["max_"+chann] != o["max_"+chann] or math.isinf(o["max_"+chann]):
        #    continue
        if o["max_"+chann] == 0:
            continue
        else:
            #print(o[chann+"_pp"])
            chi_vals.append(1.0/(o[chann+"_pp"]))
        Vh_vals.append(o["Vh"])
    #print(chi_vals)
    #popt, pcov = curve_fit(linear, Vh_vals, chi_vals, p0=[0])
    #print(round(pcov[0][0], 4))
    
    #fitted_y_vals = [linear(x, popt[0]) for x in Vh_vals]
    #plt.gca().set_aspect('equal')
   
    ax.plot(Vh_vals, chi_vals,
             linestyle=linestyle_, marker=symbols[chann], markersize=8, color = colours[chann], linewidth=2, markeredgewidth=1, markeredgecolor='black', markerfacecolor=colours[chann],
             label=r"$\chi^{{"+chann.upper() + r"} }$, $U$ = "+str(objs[0]['U']) + label_ if not no_labels else None)


    #plt.plot(x, y, linestyle='--', marker='o', markersize=10, markerfacecolor='red', markeredgewidth=2, markeredgecolor='black', linewidth=2)
    #plt.plot(Vh_vals, fitted_y_vals, "-.")