from dataclasses import dataclass
import numpy as np
import time
import ctypes


clib = ctypes.cdll.LoadLibrary('./overlap_algorithm.so')
clib.contact_function.argtypes = [
    ctypes.c_double,
    ctypes.c_double,
    ctypes.c_double,
    ctypes.c_double,
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=1, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=1, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=1, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=1, flags='C_CONTIGUOUS'),
    ctypes.c_int,
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=1, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.intc, ndim=1, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.intc, ndim=1, flags='C_CONTIGUOUS'),
]
clib.contact_function.restype = ctypes.c_double

clib.contact_function_multi.argtypes = [
    ctypes.c_int,
    ctypes.c_double,
    ctypes.c_double,
    ctypes.c_double,
    ctypes.c_double,
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=2, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=2, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=2, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=2, flags='C_CONTIGUOUS'),
    ctypes.c_int,
    np.ctypeslib.ndpointer(dtype=np.intc, ndim=1, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=1, flags='C_CONTIGUOUS'),
]
clib.contact_function_multi.restype = ctypes.c_void_p


def minimizer_map(minimizer: str) -> int:
    """Maps minimizer name to an integer for use in C functions from "overlap_algorithm.c"."""
    if minimizer == 'brent':
        return 0
    elif minimizer == 'brent_early':
        return 1
    elif minimizer == 'gss':
        return 2
    elif minimizer == 'gss_early':
        return 3
    else:
        raise ValueError('Unknown minimizer: "' + minimizer + '"')


@dataclass
class Contact:
    """Stores contact data."""
    contact_f: float
    min_eig_T: float
    med_eig_T: float
    min_vec_scalar: float  # result of scalar product test fot the eigenvector corresponding to minimal eigenvalue
    med_vec_scalar: float  # result of scalar product test fot the eigenvector corresponding to median eigenvalue
    feval_min: int  # number of minimal eigenvalue evaluations
    feval_med: int  # number of median eigenvalue evaluations
    branch: int  # contact function solution branch (0: min eig, 1: med eig, 2: Omega)


def contact_function(a0: float, a1: float, b0: float, b1: float,
                    coord0: np.ndarray, orient0: np.ndarray, coord1: np.ndarray, orient1: np.ndarray,
                    minimizer: str = 'brent') -> Contact:
    """Calculates contact data for a pair of spherical ellipses."""

    other_results = np.zeros(4, dtype=np.float64)
    fevals = np.zeros(2, dtype=np.intc)
    branch = np.zeros(1, dtype=np.intc)

    contact_f = clib.contact_function(a0, a1, b0, b1, coord0, orient0, coord1, orient1,
                                      minimizer_map(minimizer), other_results, fevals, branch)

    return Contact(contact_f, other_results[0], other_results[1], other_results[2], other_results[3],
                   fevals[0], fevals[1], branch[0])


@dataclass
class ContactMulti:
    """
    Stores the contact function for many pairs of ellipses, along with average numbers of eigenvalue evaluations
    and total calculation time.
    """
    contact_f: np.ndarray
    avg_evals_mineig: float
    avg_evals_medeig: float
    time: float


def contact_function_multi(a0: float, a1: float, b0: float, b1: float,
                           coord0: np.ndarray, orient0: np.ndarray, coord1: np.ndarray, orient1: np.ndarray,
                           minimizer: str = 'brent') -> ContactMulti:
    """Calculates contact function for multiple configurations of spherical ellipses."""

    n = len(coord0)
    results = np.zeros(n, dtype=np.float64)
    all_evals = np.zeros(2, dtype=np.intc)

    t0 = time.perf_counter()
    clib.contact_function_multi(n, a0, a1, b0, b1, coord0, orient0, coord1, orient1,
                                minimizer_map(minimizer), all_evals, results)
    t1 = time.perf_counter()

    return ContactMulti(results, all_evals[0] / n, all_evals[1] / n, t1 - t0)


class EllipsePair:

    def __init__(self, a0: float, a1: float, epsilon: float):
        """
        Class to calculate the contact function between  two spherical ellipses.
        :param a0: semi-major axis for the first elliptical cylinder
        :param a1: semi-major axis for the second elliptical cylinder
        :param epsilon: aspect ratio
        """
        self.a0 = max(a0, a1)  # defined so that always a0 > a1
        self.a1 = min(a0, a1)
        self.b0 = self.a0 / epsilon
        self.b1 = self.a1 / epsilon
        self.size_ratio = self.a0 / self.a1
        self.epsilon = epsilon

    def contact(self, coord0, orient0, coord1=np.array([0., 0., 1.]), orient1=np.array([1., 0., 0.]),
                minimizer='brent') -> Contact:
        """Evaluates the contact function along with oth contact data."""
        return contact_function(self.a0, self.a1, self.b0, self.b1, coord0, orient0, coord1, orient1, minimizer)

    def contact_multi_dist(self, distances, num_ang, axis=np.array([1, 0, 0]), minimizer='brent')\
            -> (np.ndarray, np.ndarray, np.ndarray):
        """
        Evaluates the contact function at different distances between two ellipses as well as multiple
        mutual orientations. Returns mean calculation time, mean number of minimum eigenvalue evaluations and
        mean number of median eigenvalue evaluations at each distance.
        """

        timings = np.zeros(len(distances))
        avg_evals_mineig = np.zeros(len(distances))
        avg_evals_medeig = np.zeros(len(distances))
        for i, dist in enumerate(distances):
            coords0, coords1, orients0, orients1 = generate_confgs(dist, num_ang, axis=axis)
            contact_multi = contact_function_multi(self.a0, self.a1, self.b0, self.b1,
                                                   coords0, orients0, coords1, orients1,
                                                   minimizer=minimizer)
            timings[i] = contact_multi.time
            avg_evals_mineig[i] = contact_multi.avg_evals_mineig
            avg_evals_medeig[i] = contact_multi.avg_evals_medeig

        print(f'Total time: {np.sum(timings)}')

        return timings, avg_evals_mineig, avg_evals_medeig


def generate_confgs(dist: np.ndarray, n: int, axis: np.ndarray = np.array([1, 0, 0])) \
        -> (np.ndarray, np.ndarray, np.ndarray, np.ndarray):
    """
    At a given distance, generate all possible orientational configurations of two spherical ellipses,
    with n steps in a half-circle rotation. This results in n x n configurations.
    The first ellipse will always be centered at the north pole while the position of the second is determined by
    parameter axis (and obviously distance).
    """

    angles = np.linspace(1e-7, np.pi - 1e-7, n)

    coords0 = np.zeros((n, 3))
    coords0[:, 2] = 1
    orients0 = np.array([1., 0., 0.])[None, :] * np.cos(angles)[:, None] + \
               np.array([0., 1., 0.])[None, :] * np.sin(angles)[:, None]

    coords1 = coords0 * np.cos(dist) + \
              np.cross(axis, coords0) * np.sin(dist) + \
              axis[None, :] * np.einsum('m, nm -> n', axis, coords0)[:, None] * (1 - np.cos(dist))
    orients1 = orients0 * np.cos(dist) + \
               np.cross(axis, orients0) * np.sin(dist) + \
               axis[None, :] * np.einsum('m, nm -> n', axis, orients0)[:, None] * (1 - np.cos(dist))

    indices = np.indices((n, n))
    coords0 = coords0[indices[0].flatten()]
    orients0 = orients0[indices[0].flatten()]
    coords1 = coords1[indices[1].flatten()]
    orients1 = orients1[indices[1].flatten()]

    return coords0, coords1, orients0, orients1