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:
    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:
    contact_f: float
    min_eig_T: float
    med_eig_T: float
    min_vec_scalar: float
    med_vec_scalar: float
    feval_min: int
    feval_med: int
    branch: int


def contact_function(a0, a1, b0, b1, coord0, orient0, coord1, orient1, minimizer='brent') -> Contact:
    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:
    contact_f: np.ndarray
    avg_evals_mineig: float
    avg_evals_medeig: float
    time: float


def contact_function_multi(a0, a1, b0, b1, coord0, orient0, coord1, orient1, minimizer='brent') -> ContactMulti:
    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, a1, epsilon):
        """
        Class in which optimization of ellipse configurations on a sphere is performed.
        :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:
        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):

        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, n, axis=np.array([1, 0, 0])):
    """
    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