import jax.numpy as jnp
import jax
from curvature_assembly import oriented_particle

Array = jnp.ndarray


def quadrupolar_eigenvalues(q0: Array, theta: Array) -> Array:
    return q0 * jnp.array([(jnp.cos(theta) + 3) / 4, (jnp.cos(theta) - 3) / 4, -jnp.cos(theta) / 2])


def quadrupolar_interaction(dr: Array, eigsys1: Array, eigsys2: Array, eigvals: Array) -> Array:
    """General quadrupolar interaction. However, it is really slow to evaluate in gradient-based simulations."""

    distance2 = jnp.sum(dr ** 2)
    distance4 = distance2 ** 2
    distance = jnp.sqrt(distance2)

    qf1 = oriented_particle.qf_from_rotation(eigsys1, oriented_particle.make_diagonal(eigvals))
    qf2 = oriented_particle.qf_from_rotation(eigsys2, oriented_particle.make_diagonal(eigvals))

    dr2 = jax.lax.dot_general(dr, dr, dimension_numbers=(((), ()), ((), ())))
    dr4 = jax.lax.dot_general(dr2, dr2, dimension_numbers=(((), ()), ((), ())))

    term1 = jnp.einsum('ijkl, ij, kl', dr4, qf1, qf2)
    term2 = jnp.einsum('jk, ij, ik', dr2, qf1, qf2)
    term3 = jnp.einsum('ij, ij', qf1, qf2)

    return 1 / (3 * distance ** 5) * (35 * term1 / distance4 - 20 * term2 / distance2 + 2 * term3)


def lin_quad_energy(dr: Array, eigsys1: Array, eigsys2: Array, eigvals: Array):
    """Interaction between two linear quadrupoles with eigenvalues [1. -1, 0] in this exact order."""

    q0 = eigvals[0]

    mi = eigsys1[:, 0]
    ni = eigsys1[:, 1]

    mj = eigsys2[:, 0]
    nj = eigsys2[:, 1]

    dist = jnp.sqrt(jnp.sum(dr * dr))
    rij_hat = dr / dist

    mi_rij = jnp.sum(mi * rij_hat)
    mj_rij = jnp.sum(mj * rij_hat)
    ni_rij = jnp.sum(ni * rij_hat)
    nj_rij = jnp.sum(nj * rij_hat)
    mi_mj = jnp.sum(mi * mj)
    ni_nj = jnp.sum(ni * nj)
    mi_nj = jnp.sum(mi * nj)
    ni_mj = jnp.sum(ni * mj)

    Aij = mi_rij ** 2 * mj_rij ** 2 - mi_rij ** 2 * nj_rij ** 2 - ni_rij ** 2 * mj_rij ** 2 + ni_rij ** 2 * nj_rij ** 2
    Bij = mi_mj * mi_rij * mj_rij - mi_nj * mi_rij * nj_rij - ni_mj * ni_rij * mj_rij + ni_nj * ni_rij * nj_rij
    Cij = mi_mj ** 2 - mi_nj ** 2 - ni_mj ** 2 + ni_nj ** 2

    return q0 ** 2 / (3 * dist ** 5) * (35 * Aij - 20 * Bij + 2 * Cij)


def ferro_orientational_energy(dr: Array, eigsys1: Array, eigsys2: Array, softness: float = 1.5):
    """
    Ferromagnetic-like interaction between a pair of particles. Must be combined with some distance based term.
    Softness parameter is a factor that scales the second term of the expansion and relates to the energy sensitivity
    on deviations from the parallel configuration for side by side particles. Lower values mean more stiff potential.
    Increasing it too much can lead to the preference for dipolar-like ordering (at softness = 3, effects notable at
    softness >= 2).
    """
    pi = eigsys1[:, 2]
    pj = eigsys2[:, 2]

    dist = jnp.sqrt(jnp.sum(dr * dr))
    rij_hat = dr / dist

    # positive values for attraction as added distance based term should make the entire energy negative
    return jnp.sum(pi * pj) - softness * jnp.sum(pi * rij_hat) * jnp.sum(pj * rij_hat)