import numpy as np


def calculate_anisotropy_tensor(mag_ent):
    """_summary_

    Args:
        mag_ent (_type_): _description_

    Returns:
        _type_: _description_
    """

    energies = mag_ent["energies"]
    Kxx = energies[1, 1] - energies[1, 0]
    Kyy = energies[0, 1] - energies[0, 0]
    Kzz = 0

    calculated_diff = Kyy - Kxx
    expected_diff = energies[2, 0] - energies[2, 1]
    consistency_check = abs(calculated_diff - expected_diff)

    return Kxx, Kyy, Kzz, consistency_check


def calculate_exchange_tensor(pair):
    """_summary_

    Args:
        pair (_type_): _description_

    Returns:
        _type_: _description_
    """
    energies = pair["energies"]
    # Initialize output arrays
    J = np.zeros((3, 3))
    D = np.zeros(3)

    # J matrix calculations
    # J(1,1) = mean([DEij(2,2,2), DEij(2,2,3)])
    J[0, 0] = np.mean([energies[1, 3], energies[2, 3]])

    # J(1,2) = -mean([DEij(1,2,3), DEij(2,1,3)])
    J[0, 1] = -np.mean([energies[2, 1], energies[2, 2]])
    J[1, 0] = J[0, 1]

    # J(1,3) = -mean([DEij(1,2,2), DEij(2,1,2)])
    J[0, 2] = -np.mean([energies[1, 1], energies[1, 2]])
    J[2, 0] = J[0, 2]

    # J(2,2) = mean([DEij(2,2,1), DEij(1,1,3)])
    J[1, 1] = np.mean([energies[0, 3], energies[2, 0]])

    # J(2,3) = -mean([DEij(1,2,1), DEij(2,1,1)])
    J[1, 2] = -np.mean([energies[0, 1], energies[0, 2]])
    J[2, 1] = J[1, 2]

    # J(3,3) = mean([DEij(1,1,1), DEij(1,1,2)])
    J[2, 2] = np.mean([energies[0, 0], energies[1, 0]])

    # D vector calculations
    # D(1) = mean([DEij(1,2,1), -DEij(2,1,1)])
    D[0] = np.mean([energies[0, 1], -energies[0, 2]])

    # D(2) = mean([DEij(2,1,2), -DEij(1,2,2)])
    D[1] = np.mean([energies[1, 2], -energies[1, 1]])

    # D(3) = mean([DEij(1,2,3), -DEij(2,1,3)])
    D[2] = np.mean([energies[2, 1], -energies[2, 2]])

    J_iso = np.trace(J) / 3
    J_S = (J - J_iso * np.eye(3)).flatten()

    return J_iso, J_S, D, J


def int_de_ke(traced, we):
    """_summary_

    Args:
        traced (_type_): _description_
        we (_type_): _description_

    Returns:
        _type_: _description_
    """
    return np.trapz(-1 / np.pi * np.imag(traced * we))