"""Isotropic Migdal-Eliashberg Tc on the Matsubara axis for spectra made of Einstein modes.

alpha^2F(w) = sum_i (lam_i * w_i / 2) * delta(w - w_i), so that
lambda(i nu) = sum_i lam_i * w_i^2 / (w_i^2 + nu^2).

Linearised gap equations in the Allen-Dynes symmetric form, for n, m >= 0:

    K_nm = lambda(n - m) + lambda(n + m + 1) - 2 mu*
           - delta_nm * [2n + 1 + lambda(0) + 2 sum_{l=1..n} lambda(l)]

with lambda(l) evaluated at the bosonic frequency 2 pi T l. Tc is the temperature at
which the largest eigenvalue of K crosses zero. Frequencies and temperatures share one
unit (hbar = kB = 1).
"""

import json
import math
from pathlib import Path

import numpy as np
from scipy.linalg import eigh
from scipy.optimize import brentq

HERE = Path(__file__).parent

# Below this Tc / omega_max the Matsubara matrix exceeds about 3000 frequencies at a
# cutoff of ten times omega_max; such weak-coupling cases are reported as zero.
T_FLOOR = 5e-4


def lam_matsubara(l, t, modes):
    nu = 2.0 * math.pi * t * np.asarray(l, dtype=float)
    out = np.zeros_like(nu)
    for lam_i, w_i in modes:
        out += lam_i * w_i**2 / (w_i**2 + nu**2)
    return out


def max_eig(t, modes, mu_star, w_cut):
    """Largest eigenvalue of the Allen-Dynes kernel at temperature t."""
    n_max = max(int(math.floor(w_cut / (2.0 * math.pi * t) - 0.5)) + 1, 4)
    n = np.arange(n_max)
    lam_l = lam_matsubara(np.arange(2 * n_max + 1), t, modes)
    k = lam_l[np.abs(n[:, None] - n[None, :])] + lam_l[n[:, None] + n[None, :] + 1] - 2.0 * mu_star
    cum = np.concatenate([[0.0], np.cumsum(lam_l[1:n_max])])
    k[n, n] -= 2 * n + 1 + lam_l[0] + 2.0 * cum
    return eigh(k, eigvals_only=True, subset_by_index=[n_max - 1, n_max - 1])[0]


def tc(modes, mu_star=0.0, cut_factor=10.0):
    """Tc for a list of (lambda_i, omega_i) modes. Cutoff is cut_factor * max omega."""
    w_max = max(w for _, w in modes)
    w_cut = cut_factor * w_max
    f = lambda t: max_eig(t, modes, mu_star, w_cut)
    hi = 3.0 * w_max * math.sqrt(max(sum(l for l, _ in modes), 1.0))
    lo = hi
    while f(lo) < 0.0:
        lo *= 0.7
        if lo < T_FLOOR * w_max:
            return 0.0
    hi = lo / 0.7
    while f(hi) > 0.0:
        hi *= 1.3
    return brentq(f, lo, hi, xtol=1e-7 * w_max, rtol=1e-6)


def allen_dynes(lam, w_log, w2, mu_star):
    """Allen-Dynes 1975 with the strong-coupling (f1) and shape (f2) factors."""
    if lam <= mu_star * (1 + 0.62 * lam):
        return 0.0
    l1 = 2.46 * (1 + 3.8 * mu_star)
    l2 = 1.82 * (1 + 6.3 * mu_star) * (w2 / w_log)
    f1 = (1 + (lam / l1) ** 1.5) ** (1.0 / 3.0)
    f2 = 1 + (w2 / w_log - 1) * lam**2 / (lam**2 + l2**2)
    return f1 * f2 * (w_log / 1.2) * math.exp(-1.04 * (1 + lam) / (lam - mu_star * (1 + 0.62 * lam)))


def mcmillan(lam, w_log, mu_star):
    if lam <= mu_star * (1 + 0.62 * lam):
        return 0.0
    return (w_log / 1.2) * math.exp(-1.04 * (1 + lam) / (lam - mu_star * (1 + 0.62 * lam)))


def mu_at_cutoff(mu_star, ratio):
    """Rescale mu* from a cutoff w0 to a cutoff ratio * w0."""
    return mu_star / (1.0 - mu_star * math.log(ratio))


def moments(modes):
    lam = sum(l for l, _ in modes)
    w_log = math.exp(sum(l * math.log(w) for l, w in modes) / lam)
    w2 = math.sqrt(sum(l * w**2 for l, w in modes) / lam)
    return lam, w_log, w2


def einstein_table():
    lams = [0.3, 0.5, 0.75, 1.0, 1.5, 2.0, 2.5, 3.0, 4.0, 5.0, 7.0, 10.0]
    rows = []
    for mu in (0.0, 0.10, 0.16):
        for lam in lams:
            t10 = tc([(lam, 1.0)], mu, 10.0)
            rows.append({
                "lambda": lam,
                "mu_star": mu,
                "tc_over_omega_e": t10,
                "allen_dynes_same_mu": allen_dynes(lam, 1.0, 1.0, mu),
                "mcmillan_same_mu": mcmillan(lam, 1.0, mu),
                "asymptote": 0.1827 * math.sqrt(lam),
            })
    return rows


def cutoff_check():
    out = []
    for lam in (0.5, 1.0, 2.0, 5.0, 10.0):
        a, b, c = (tc([(lam, 1.0)], 0.0, f) for f in (10.0, 20.0, 40.0))
        out.append({"lambda": lam, "cut10": a, "cut20": b, "cut40": c, "rel_change_10_to_20": abs(b - a) / b})
    return out


def dense_curve(mu, cut=10.0):
    lams = np.concatenate([np.arange(0.4, 2.0, 0.05), np.arange(2.0, 6.01, 0.1)])
    return [{"lambda": float(l), "tc_over_omega_e": tc([(float(l), 1.0)], mu, cut)} for l in lams]


def two_mode(total_lambda, w_low, w_high, mu, fractions):
    """Fixed total lambda split between a soft mode and a stiff mode."""
    out = []
    for x in fractions:
        modes = []
        if x < 1.0:
            modes.append(((1 - x) * total_lambda, w_low))
        if x > 0.0:
            modes.append((x * total_lambda, w_high))
        lam, w_log, w2 = moments(modes)
        # one absolute cutoff for every split, so mu* means the same thing throughout
        cut = 10.0 * w_high / max(w for _, w in modes)
        out.append({
            "fraction_on_stiff_mode": float(x),
            "tc": tc(modes, mu, cut),
            "omega_log": w_log,
            "omega_2": w2,
            "allen_dynes": allen_dynes(lam, w_log, w2, mu),
        })
    return out


def functional_derivative(base_modes, mu, probes, eps=2e-3):
    """dTc / d(alpha^2F(Omega)): response of Tc to a small delta-function of weight eps at Omega."""
    t0 = tc(base_modes, mu, 10.0)
    w_cut_factor = 10.0
    w_max = max(w for _, w in base_modes)
    out = []
    for om in probes:
        d_lam = 2.0 * eps / om
        modes = base_modes + [(d_lam, om)]
        # hold the absolute cutoff fixed so the probe does not move it
        f = lambda t: max_eig(t, modes, mu, w_cut_factor * w_max)
        t1 = brentq(f, 0.5 * t0, 2.0 * t0, xtol=1e-10, rtol=1e-10)
        out.append({"omega_over_tc": om / t0, "dtc_dalpha2f": (t1 - t0) / eps})
    return t0, out


if __name__ == "__main__":
    res = {}
    res["einstein_table"] = einstein_table()
    res["cutoff_check_mu0"] = cutoff_check()
    res["einstein_curves"] = {f"{mu:.2f}": dense_curve(mu) for mu in (0.0, 0.10, 0.16)}

    # Hydride-like two-mode spectrum: 15 meV host mode and 150 meV hydrogen mode.
    fr = np.linspace(0.0, 1.0, 21)
    res["two_mode"] = {
        f"{L:.1f}": two_mode(L, 15.0, 150.0, 0.13, fr) for L in (1.0, 2.0, 3.0)
    }

    t0, fd = functional_derivative([(2.0, 1.0)], 0.13, np.geomspace(0.05, 12.0, 60))
    res["functional_derivative"] = {"base": "Einstein, lambda=2, mu*=0.13 at 10 omega_E", "tc": t0, "curve": fd}

    (HERE / "results.json").write_text(json.dumps(res, indent=1))

    print("Einstein spectrum, Tc / omega_E (cutoff 10 omega_E)")
    print(f"{'lambda':>7} {'mu*':>5} {'Eliashberg':>11} {'AllenDynes':>11} {'AD/El':>7} {'asym':>7}")
    for r in res["einstein_table"]:
        ratio = r["allen_dynes_same_mu"] / r["tc_over_omega_e"] if r["tc_over_omega_e"] else float("nan")
        print(f"{r['lambda']:7.2f} {r['mu_star']:5.2f} {r['tc_over_omega_e']:11.5f} {r['allen_dynes_same_mu']:11.5f} {ratio:7.3f} {r['asymptote']:7.4f}")
    print("\ncutoff check, mu*=0")
    for r in res["cutoff_check_mu0"]:
        print(r)
    pk = max(fd, key=lambda d: d["dtc_dalpha2f"])
    print(f"\nfunctional derivative peaks at omega = {pk['omega_over_tc']:.2f} kB Tc (base Tc = {t0:.4f} omega_E)")
    for L, rows in res["two_mode"].items():
        print(f"two-mode lambda={L}: Tc(all soft)={rows[0]['tc']:.2f} meV, Tc(half)={rows[10]['tc']:.2f}, Tc(all stiff)={rows[-1]['tc']:.2f} meV")