"""An example of use of KuboBastin kernels for surface and sea contributions.
See. Bonbien and Manchon, Phys. Rev. B 102, 085113 (2020).
Slightly different set of parameters are used.
"""

import matplotlib.pyplot as plt
import numpy as np
from sympy import Matrix, cos, sin, symbols

from py4mulas.mu_kernels import KuboBastinKernel
from py4mulas.models import Kmodel
from py4mulas.responses import Kubo

sigma_0 = Matrix([[1, 0], [0, 1]])
sigma_x = Matrix([[0, 1], [1, 0]])
sigma_y = Matrix([[0, -1j], [1j, 0]])
sigma_z = Matrix([[1, 0], [0, -1]])


def rashba_model():
    # Rashba model as a sympy expression
    k_x, k_y = symbols("k_x k_y")
    t, tR, delta = symbols("t tR delta")
    model = 2 * t * sigma_0 - 2 * t * (cos(k_x) + cos(k_y)) * sigma_0

    # add ferromagnetism
    model += delta * sigma_z

    # add spin-orbit coupling
    model += tR * (sin(k_y) * sigma_x - sin(k_x) * sigma_y)
    return model


def plot_conductivity(nk=200, beta="y"):
    toy_model = rashba_model()
    kxs = np.linspace(-np.pi, np.pi, nk)
    kys = kxs
    params = dict(t=1, tR=1.4, delta=0.2)

    # define a py4mulas model
    model = Kmodel(toy_model, k_1=kxs, k_2=kys, params=params, real_space=True)

    alpha = "x"

    options = dict(precomp=True, chunk_size=10000)
    G = Kubo(
        model,
        alpha=alpha,
        beta=beta,
        kspace_options=options,
        mu_kernel=KuboBastinKernel(None),
    )

    energies = np.linspace(-3, -1, 50)

    # for display and matplotlib plotting
    styles = {None: "s", "BMI": "ob", "BMII": "o", "StredaI": "-r", "StredaII": "-b"}
    contributions = ["BMI", "StredaI"]
    labels = {"BMI": r"$\tilde\sigma^{I}_{xx}$", "StredaI": r"$\sigma^I_{xx}$"}
    plt.xticks([-3, -2.5, -2, -1.5, -1], fontsize=12)
    plt.yticks(np.arange(0, 30, 5), fontsize=12)
    if beta == "y":
        contributions = [None, "BMII", "StredaI", "StredaII"]
        labels = {
            None: r"$\sigma_{xy}$",
            "BMI": r"$\tilde\sigma^{I}_{xy}$",
            "BMII": r"$\tilde\sigma^{II}_{xy}$",
            "StredaI": r"$\sigma^I_{xy}$",
            "StredaII": r"$\sigma^{II}_{xy}$",
        }
        plt.yticks([-0.4, -0.2, 0, 0.2], fontsize=12)
    # end of display related

    def compute_and_plot(kernel_name):
        if kernel_name is not None:
            # we update the kernel name from the same object
            # if the object is recreated the operator kernel will have to be re-evaluated
            # which does not make sense as all these kernels have the same operator kernel
            G.mu_kernel = KuboBastinKernel(kernel_name)
        results = []
        for energy in energies:
            kwargs = dict(mu=energy, temperature=0.1, eta=0.1)
            results.append(G(**kwargs))

        plt.plot(energies, results, styles[kernel_name], label=labels[kernel_name])

    for term in contributions:
        compute_and_plot(term)

    plt.legend(fontsize=14)
    # plt.savefig(alpha+beta+'.png')
    plt.show()


def main():
    plot_conductivity(nk=400, beta="y")


if __name__ == "__main__":
    main()