# -*- coding:utf-8; eval: (blacken-mode) -*-
import os
import pytest
import numpy as np
from numpy.testing import assert_allclose
from numpy import pi
from scipy.special import zeta
from scipy.linalg import expm
from usadelndsoc.solver import *
from usadelndsoc.bcs import BCS_Delta
from usadelndsoc.matsubara import get_matsubara_sum

try:
    import usadel1
except ImportError:
    usadel1 = None


def example_sns_1d(phi, L, nx=500, ny=1, Delta0=1.0, eta=0.0):
    solver = Solver(nx=nx, ny=ny)
    solver.mask[...] = MASK_NONE
    solver.mask[0, :] = MASK_TERMINAL
    solver.mask[-1, :] = MASK_TERMINAL
    solver.Lx = L
    solver.Ly = L
    solver.eta = eta

    solver.Omega[...] = 0

    solver.Delta[0, :] = Delta0 * np.exp(-1j * phi / 2) * np.eye(2)
    solver.Delta[-1, :] = Delta0 * np.exp(+1j * phi / 2) * np.eye(2)

    return solver


def test_example_sns_1d():
    # Test against analytic solution
    omega = 30
    Lx = 10
    phi = pi / 2
    sol = example_sns_1d(phi, Lx)
    res = sol.solve(omega=omega)

    x = np.linspace(0, Lx, sol.shape[0])

    theta0 = np.arctan(1 / omega)
    theta_a = 4 * np.arctan(np.exp(-np.sqrt(2 * omega) * x) * np.tan(theta0 / 4))
    theta_b = 4 * np.arctan(np.exp(-np.sqrt(2 * omega) * (Lx - x)) * np.tan(theta0 / 4))

    F = np.sin(theta_a) * np.exp(-1j * phi / 2) + np.sin(theta_b) * np.exp(1j * phi / 2)

    rF = res.F[:, :, 0, 0]
    assert_allclose(rF.squeeze(), F, rtol=0, atol=5e-3 * abs(F).max())


@pytest.mark.skipif(usadel1 is None, reason="usadel1 not installed")
def test_example_sns_1d_J_usadel1():
    Lx = 10
    phi = pi / 2
    sol = example_sns_1d(phi, Lx)
    sol.Ly = 1

    T = 0.2
    Delta, I0, _, success = sol.self_consistency(T=T, T_c0=0, tol=1e-10)

    tau3 = np.diag([1, 1, -1, -1])
    I0 = tr(I0[:-1, 0, 0] @ tau3)

    I0_mean = I0[1:-1].real.mean()

    # Current is real-valued and conserved
    assert_allclose(I0.imag, 0, atol=1e-8)
    assert_allclose(I0[1:-1].real, I0_mean, rtol=1e-5)

    # Compare to Usadel1
    g = usadel1.Geometry(1, 2)
    g.t_type = [usadel1.NODE_CLEAN_S_TERMINAL, usadel1.NODE_CLEAN_S_TERMINAL]
    g.w_type = usadel1.WIRE_TYPE_N
    g.w_ends[0, :] = [0, 1]
    g.w_length = Lx
    g.w_conductance = 1
    g.t_delta = [1, 1]
    g.t_phase = [-phi / 2, phi / 2]
    g.t_t = T
    g.t_mu = 0

    s = usadel1.CurrentSolver(g, maxE=4, ne=2000)
    s.solve()
    I1, _ = s.get_currents()
    I1_mean = I1.mean()

    assert_allclose(I0_mean, I1_mean, rtol=1e-3)


@pytest.mark.skipif(usadel1 is None, reason="usadel1 not installed")
def test_example_sss_1d_J_usadel1():
    Lx = 10
    phi = pi / 2
    sol = example_sns_1d(phi, 30 * Lx / (30 - 2), nx=30)
    sol.Ly = 1

    T_c0 = np.exp(np.euler_gamma) / np.pi
    T = 0.2
    Delta, I0, _, success = sol.self_consistency(T=T, T_c0=T_c0, workers=1)

    tau3 = np.diag([1, 1, -1, -1])
    I0 = tr(I0[:-1, 0, 0] @ tau3)

    I0_mean = I0.real.mean()

    # Current is real-valued and conserved
    assert_allclose(I0.imag, 0, atol=1e-8)
    assert_allclose(I0.real, I0_mean, rtol=5e-3)

    # Compare to Usadel1
    g = usadel1.Geometry(1, 2)
    g.t_type = [usadel1.NODE_CLEAN_S_TERMINAL, usadel1.NODE_CLEAN_S_TERMINAL]
    g.w_type = usadel1.WIRE_TYPE_S
    g.w_ends[0, :] = [0, 1]
    g.w_delta = 1.0
    g.w_phase = np.linspace(-phi / 2, phi / 2, len(g.x))
    g.w_length = Lx
    g.w_conductance = 1
    g.t_delta = [1, 1]
    g.t_phase = [-phi / 2, phi / 2]
    g.t_t = T
    g.t_mu = 0

    g.omega_D = 100
    g.coupling_lambda = 1 / np.arccosh(g.omega_D / 1.0)

    for k, d, v in usadel1.self_consistent_matsubara_iteration(g):
        if d.relative_residual_norm() < 1e-6:
            break

    s = usadel1.CurrentSolver(g, ne=2000)
    s.solve()
    I1, _ = s.get_currents()
    I1_mean = I1.mean()

    assert_allclose(I0_mean, I1_mean, rtol=5e-2)


OMEGAS = [30, 5, (0.3 + 0.001j) / 1j]


@pytest.mark.skipif(usadel1 is None, reason="usadel1 not installed")
@pytest.mark.parametrize("omega", OMEGAS)
def test_example_sns_usadel1(omega):
    # Test against another numerical solution
    Lx = 10
    phi = pi / 2
    sol = example_sns_1d(phi, Lx)
    res = sol.solve(omega=omega)

    g = usadel1.Geometry(1, 2)
    g.t_type = [usadel1.NODE_CLEAN_S_TERMINAL, usadel1.NODE_CLEAN_S_TERMINAL]
    g.w_type = usadel1.WIRE_TYPE_N
    g.w_ends[0, :] = [0, 1]
    g.w_length = Lx
    g.w_conductance = 1
    g.t_delta = [1, 1]
    g.t_phase = [-phi / 2, phi / 2]
    g.t_t = 1
    g.t_mu = 0

    s = usadel1.Solver()
    s.set_geometry(g)
    x_scl = np.linspace(0, 1, sol.shape[0])
    res2 = s.sp_solve(np.array([1j * omega]), x_scl)
    F2 = 2j * res2.a / (1 - res2.a * res2.b)
    F2c = 2j * res2.b / (1 - res2.a * res2.b)

    assert_allclose(
        res.F[..., 0, 0].squeeze(), F2.squeeze(), rtol=1e-2, atol=1e-3 * abs(F2).max()
    )
    assert_allclose(
        res.Fc[..., 0, 0].squeeze(),
        F2c.squeeze(),
        rtol=1e-2,
        atol=1e-3 * abs(F2c).max(),
    )


@pytest.mark.slow
@pytest.mark.parametrize("T", [0.28, 0.53, 0.8])
def test_selfcons_kupriyanov(request, T):
    fns_k = {
        0.28: "kupriyanov1981-cpr-a1.dat",
        0.53: "kupriyanov1981-cpr-a2.dat",
        0.8: "kupriyanov1981-cpr-a3.dat",
    }
    fns_u = {
        0.28: "L3.989_T0.28_TcR1_TcW1_gamma1e-06_r0_ns7.dat",
        0.53: "L3.989_T0.53_TcR1_TcW1_gamma1e-06_r0_ns7.dat",
        0.8: "L3.989_T0.8_TcR1_TcW1_gamma1e-06_r0_ns7.dat",
    }

    fn_k = os.path.join(os.path.dirname(__file__), "data", fns_k[T])
    fn_u = os.path.join(os.path.dirname(__file__), "data", fns_u[T])

    L = 3.989
    Tc0 = 1.0

    Delta0 = BCS_Delta(T, Tc0)
    Delta00 = BCS_Delta(0.05, Tc0)

    # L is the distance between terminals, but Lx the system size where
    # the two outermost layers are terminals
    nx = 40
    Lp = (L / (nx - 2)) * nx
    sol = example_sns_1d(0, Lp, nx=nx, ny=1, Delta0=Delta0)
    sol.Ly = 1
    phase_mask = (sol.mask == MASK_TERMINAL) & (sol.x[:, None] > 0)

    dn = request.config.cache.makedir("test_selfcons_kupriyanov")
    cache_fn = os.path.join(dn, f"T-{T}.npz")
    res = cpr(
        sol,
        T=T,
        T_c0=Tc0,
        phase_mask=phase_mask,
        filename=cache_fn,
        workers=1,
        ds=0.4,
    )

    phis, Deltas, Js = res

    tau3 = np.diag([1, 1, -1, -1])
    Js = tr(Js @ tau3)
    Is = Js[:, :-1, 0, 0].squeeze().mean(axis=1)

    m = ~np.isnan(Is)
    Is = Is[m]
    phis = phis[m]

    phis_k, Is_k = np.loadtxt(fn_k).T
    phis_u, Is_u = np.loadtxt(fn_u).T

    Is_k *= Delta00 / L

    Imax = abs(Is_u).max()

    assert curve_separation(phis, Is, phis_u, Is_u) < 0.075 * Imax
    assert curve_separation(phis, Is, phis_k, Is_k) < 0.11 * Imax


def curve_separation(x1, y1, x2, y2):
    # Resample
    k1 = np.arange(len(x1))
    k1p = np.linspace(k1[0], k1[-1], 2000)
    k2 = np.arange(len(x2))
    k2p = np.linspace(k2[0], k2[-1], 2000)

    x1 = np.interp(k1p, k1, x1)
    y1 = np.interp(k1p, k1, y1)

    x2 = np.interp(k2p, k2, x2)
    y2 = np.interp(k2p, k2, y2)

    # Min-max separation
    s = np.hypot(abs(x1[:, None] - x2[None, :]), abs(y1[:, None] - y2[None, :]))
    return s.min(axis=1).max()


@pytest.mark.parametrize(
    "T,n,alpha_soc", [(0.1, 10, 0.01), (0.4, 10, 0.01), (0.4, 5, 0.0123)]
)
def test_soc_analytic(T, n, alpha_soc):
    sx = np.array([[0, 1], [1, 0]])
    sy = np.array([[0, -1j], [1j, 0]])
    sz = np.array([[1, 0], [0, -1]])
    s0 = np.array([[1, 0], [0, 1]])

    def get_sol(soc_alpha, eta, L=10, h=0.5, D=1.0):
        sol = Solver(nx=n, ny=n)
        sol.mask[...] = MASK_NONE
        sol.Lx = L
        sol.Ly = L

        sol.Omega[...] = 0
        sol.Delta[:, :] = np.eye(2)

        sol.D = D
        sol.eta = eta

        nx, ny = sol.shape
        dx = sol.Lx / nx
        dy = sol.Ly / ny

        Ax = 2 * soc_alpha * np.kron(s0, sy)
        Ay = -2 * soc_alpha * np.kron(s0, sx)

        sol.Ux[...] = expm(1j * dx * Ax)
        sol.Uy[...] = expm(1j * dy * Ay)

        sol.Omega[...] += h * np.kron(sz, sx)

        return sol

    # Solver without SOC
    h = 0.5
    sol0 = get_sol(soc_alpha=0, eta=0, h=h)

    # Solver for evaluating current perturbatively in SOC
    eta = 1.0
    sol = get_sol(soc_alpha=alpha_soc, eta=eta, h=h, D=0)

    E_typical = 10 + 10 * T
    w, a = get_matsubara_sum(T, E_typical)

    t3 = np.kron(sz, s0)

    Phis = []
    Fs = []
    Gs = []
    Js = []
    Jans = []

    dd = (sol.Lx / sol.shape[0], sol.Ly / sol.shape[1])

    Jtot = 0
    Jantot = 0
    for ww, aa in zip(w, a):
        res0 = sol0.solve(omega=ww)
        sol.omega = ww
        J = sol._core.grad_A(res0.Phi.copy()).transpose(0, 1, 2, 4, 3) * (1 / 16)
        Jtot += -1j * pi * aa * J

        gp = (ww - 1j * h) / np.sqrt((ww - 1j * h) ** 2 + 1)
        gm = (ww + 1j * h) / np.sqrt((ww + 1j * h) ** 2 + 1)
        fp = 1 / np.sqrt((ww - 1j * h) ** 2 + 1)
        fm = 1 / np.sqrt((ww + 1j * h) ** 2 + 1)
        # XXX: check the 8*eta
        Jy_an = (gp - gm) * (1 + gp * gm + fp * fm) * (-1j * alpha_soc**3) * (8 * eta)
        Jy_an *= dd[1]  # multiply current density by the width of a single cell
        Jantot += pi * aa * Jy_an

        Jy = -1j * tr(J @ t3)[n // 2, n // 2, 1]

        assert_allclose(Jy, Jy_an, rtol=1e-2)

        Fs.append(res0.F.copy())
        Gs.append(res0.G.copy())
        Phis.append(res0.Phi.copy())
        Js.append(J)
        Jans.append(Jy_an)


def test_selfcons_h():
    T = 0.4
    h = 0.5

    solver = Solver(nx=5, ny=2)
    solver.mask[...] = MASK_NONE
    solver.Lx = 5
    solver.Ly = 5
    solver.eta = 0
    solver.Delta[...] = 1.0
    solver.Omega[...] += h * np.kron(S_z, S_y)

    T_c0 = np.exp(np.euler_gamma) / np.pi
    Delta, I0, _, success = solver.self_consistency(T=T, T_c0=T_c0, workers=1)

    Delta0 = BCS_Delta(T, T_c0, h=h)
    Deltam = np.broadcast_to((Delta0 * S_0)[None, None, :, :], Delta.shape)

    assert_allclose(Delta.copy(), Deltam, rtol=1e-3)