# -*- 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 usadelndsoc.solver import *
from usadelndsoc.bcs import BCS_Delta

try:
    import usadel1
except ImportError:
    usadel1 = None


def example_sns_1d(phi, L, nx=500, ny=1, Delta0=1.0, alpha=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.alpha = alpha

    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)

    tau3 = np.diag([1, 1, -1, -1])
    I0 = (I0[:-1, 0, 0] @ tau3).trace(axis1=-2, axis2=-1)

    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_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=5e-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 = (I0[:-1, 0, 0] @ tau3).trace(axis1=-2, axis2=-1)

    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, L, 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 = (Js @ tau3).trace(axis1=-2, axis2=-1)
    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.04 * Imax
    assert curve_separation(phis, Is, phis_k, Is_k) < 0.075 * 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()