import scipy as _sp
import sympy as syp
import numpy as np
import pandas as pd
import logging
from transforms3d import _gohlketransforms as tr
logger = logging.getLogger(__name__)


def _get_key_props(phase=None, diameter='throat.diameter',
                   surface_tension='pore.surface_tension',
                   contact_angle='pore.contact_angle'):
    r"""
    Many of the methods are generic to pores and throats. Some information may
    be stored on either the pore or throat and needs to be interpolated.
    This is a helper method to return the properties in the correct format.
    To do:
        Check for method to convert throat to pore data
    """
    element = diameter.split('.')[0]
    if element == 'pore':
        if 'throat' in surface_tension:
            sigma = phase.interpolate_data(propname=surface_tension)
        else:
            sigma = phase[surface_tension]
        if 'throat' in contact_angle:
            theta = phase.interpolate_data(propname=contact_angle)
        else:
            theta = phase[contact_angle]
    if element == 'throat':
        if 'pore' in surface_tension:
            sigma = phase.interpolate_data(propname=surface_tension)
        else:
            sigma = phase[surface_tension]
        if 'pore' in contact_angle:
            theta = phase.interpolate_data(propname=contact_angle)
        else:
            theta = phase[contact_angle]

    return element, sigma, theta


def washburn(target, surface_tension='pore.surface_tension',
             contact_angle='pore.contact_angle',
             diameter='throat.diameter'):
    r"""
    Computes the capillary entry pressure assuming the throat in a cylindrical
    tube.

    Parameters
    ----------
    target : OpenPNM Object
        The object for which these values are being calculated.  This
        controls the length of the calculated array, and also provides
        access to other necessary thermofluid properties.

    sigma : string
        The dictionary key containing the surface tension values to be used. If
        a pore property is given, it is interpolated to a throat list.

    theta : string
        The dictionary key containing the contact angle values to be used. If
        a pore property is given, it is interpolated to a throat list.

    diameter : string
        The dictionary key containing the throat diameter values to be used.

    Notes
    -----
    The Washburn equation is:

    .. math::
        P_c = -\frac{2\sigma(cos(\theta))}{r}

    This is the most basic approach to calculating entry pressure and is
    suitable for highly non-wetting invading phases in most materials.

    """
    network = target.project.network
    phase = target.project.find_phase(target)
    element, sigma, theta = _get_key_props(phase=phase,
                                           diameter=diameter,
                                           surface_tension=surface_tension,
                                           contact_angle=contact_angle)
    r = network[diameter]/2
    value = -2*sigma*_sp.cos(_sp.radians(theta))/r
    if diameter.split('.')[0] == 'throat':
        value = value[phase.throats(target.name)]
    else:
        value = value[phase.pores(target.name)]
    value[_sp.absolute(value) == _sp.inf] = 0
    return value


def purcell(target, r_toroid, surface_tension='pore.surface_tension',
            contact_angle='pore.contact_angle',
            diameter='throat.diameter'):
    r"""
    Computes the throat capillary entry pressure assuming the throat is a
    toroid.

    Parameters
    ----------
    target : OpenPNM Object
        The object for which these values are being calculated.  This
        controls the length of the calculated array, and also provides
        access to other necessary thermofluid properties.

    r_toroid : float or array_like
        The radius of the toroid surrounding the pore

    sigma : dict key (string)
        The dictionary key containing the surface tension values to be used.
        If a pore property is given, it is interpolated to a throat list.

    theta : dict key (string)
        The dictionary key containing the contact angle values to be used.
        If a pore property is given, it is interpolated to a throat list.

    diameter : dict key (string)
        The dictionary key containing the throat diameter values to be used.

    Notes
    -----
    This approach accounts for the converging-diverging nature of many throat
    types.  Advancing the meniscus beyond the apex of the toroid requires an
    increase in capillary pressure beyond that for a cylindical tube of the
    same radius. The details of this equation are described by Mason and
    Morrow [1]_, and explored by Gostick [2]_ in the context of a pore network
    model.

    References
    ----------

    .. [1] G. Mason, N. R. Morrow, Effect of contact angle on capillary
           displacement curvatures in pore throats formed by spheres. J.
           Colloid Interface Sci. 168, 130 (1994).
    .. [2] J. Gostick, Random pore network modeling of fibrous PEMFC gas
           diffusion media using Voronoi and Delaunay tessellations. J.
           Electrochem. Soc. 160, F731 (2013).

    """
    network = target.project.network
    phase = target.project.find_phase(target)
    element, sigma, theta = _get_key_props(phase=phase,
                                           diameter=diameter,
                                           surface_tension=surface_tension,
                                           contact_angle=contact_angle)
    r = network[diameter]/2
    R = r_toroid
    alpha = theta - 180 + _sp.arcsin(_sp.sin(_sp.radians(theta)/(1+r/R)))
    value = (-2*sigma/r) * \
        (_sp.cos(_sp.radians(theta - alpha)) /
            (1 + R/r*(1 - _sp.cos(_sp.radians(alpha)))))
    if diameter.split('.')[0] == 'throat':
        value = value[phase.throats(target.name)]
    else:
        value = value[phase.pores(target.name)]
    return value


def ransohoff_snap_off(target,
                       shape_factor=2.0,
                       require_pair=False,
                       contact_angle='pore.contact_angle',
                       surface_tension='pore.surface_tension',
                       diameter='throat.diameter',
                       wavelength=5e-6,
                       vertices='throat.offset_vertices',
                       **kwargs):
    r"""
    Computes the capillary snap-off pressure assuming the throat is cylindrical
    with converging-diverging change in diamater - like the Purcell model.
    The wavelength of the change in diamater is the fiber radius.
    Ref: Ransohoff, T.C., Gauglitz, P.A. and Radke, C.J., 1987. Snap‐off of gas
    bubbles in smoothly constricted noncircular capillaries. AIChE Journal,
    33(5), pp.753-765.

    Parameters
    ----------
    target : OpenPNM Object
        The object for which these values are being calculated.  This
        controls the length of the calculated array, and also provides
        access to other necessary thermofluid properties.
    shape_factor :
        constant dependent on the shape of throat cross-section 1.75 - 2.0, see
        Ref
    sigma : dict key (string)
        The dictionary key containing the surface tension values to be used. If
        a pore property is given, it is interpolated to a throat list.
    throat_diameter : dict key (string)
        The dictionary key containing the throat diameter values to be used.
    wavelength :
        The dictionary key containing the radius of the transverse interfacial
        radius of curvature at the neck (fiber radius in fibrous media)

    Notes
    -----
    This equation should be used to calculate the snap off capillary pressure
    in fribrous media

    """
    phase = target.project.find_phase(target)
    geometry = target.project.find_geometry(target)
    element, sigma, theta = _get_key_props(phase=phase,
                                           diameter=diameter,
                                           surface_tension=surface_tension,
                                           contact_angle=contact_angle)
    try:
        all_verts = geometry[vertices]
        # Work out whether throat geometry can support at least one pair of
        # adjacent arc menisci that can grow and merge to form snap-off
        # Only works if throat vertices are in convex hull order
        angles_ok = np.zeros(geometry.Nt, dtype=bool)
        for T in range(geometry.Nt):
            verts = all_verts[T]
            x = verts[:, 0]
            y = verts[:, 1]
            z = verts[:, 2]
            # PLus
            p = 1
            # Minus
            m = -1
            verts_p = np.vstack((np.roll(x, p),
                                 np.roll(y, p),
                                 np.roll(z, p))).T
            verts_m = np.vstack((np.roll(x, m),
                                 np.roll(y, m),
                                 np.roll(z, m))).T
            v1 = verts_p - verts
            v2 = verts_m - verts
            corner_angles = np.rad2deg(tr.angle_between_vectors(v1,
                                                                v2,
                                                                axis=1))
            # Logical test for existence of arc menisci
            am = theta[T] <= 90 - corner_angles/2
            if require_pair:
                # Logical test for two adjacent arc menisci
                pair_p = np.logical_and(am, np.roll(am, + p))
                pair_m = np.logical_and(am, np.roll(am, + m))
                am_pair = np.any(np.logical_or(pair_p, pair_m))
                angles_ok[T] = am_pair
            else:
                # Logical test for any arc menisci
                angles_ok[T] = np.any(am)
    except:
        logger.warning("Model is designed to work with property: " +
                       vertices)
        angles_ok = np.ones(geometry.Nt, dtype=bool)

    # Condition for arc menisci to form in corners
    rad_Ts = geometry[diameter]/2
    # Ransohoff and Radke eq. 4
    C = 1/rad_Ts - 1/wavelength
    value = sigma[phase.throats(target.name)]*C
    # Only throats that can support arc menisci can snap-off
    value[~angles_ok] = np.nan
    logger.info("Snap off pressures calculated for " +
                str(np.around(100*np.sum(angles_ok)/np.size(angles_ok), 0)) +
                "% of throats")
    return value
