"""Elastic free-energy derivatives for Flory, phantom, and Miao
network models."""


def _elastic_prefactor(partial_molar_vol_pol, elasticity_c, vol_frac_pol):
    return (elasticity_c * vol_frac_pol) / partial_molar_vol_pol


def flory_rehner_dad_v(
    partial_molar_vol_pol, elasticity_c, vol_frac_pol_0, vol_frac_pol
):
    """Simple Flory-Rehner elasticity contribution."""
    return _elastic_prefactor(
        partial_molar_vol_pol, elasticity_c, vol_frac_pol
    ) * ((vol_frac_pol_0 / vol_frac_pol) ** (2 / 3) - 0.5)


def phantom_dad_v(
    partial_molar_vol_pol, elasticity_c, vol_frac_pol_0, vol_frac_pol
):
    """Simple phantom network elasticity contribution."""
    return _elastic_prefactor(
        partial_molar_vol_pol, elasticity_c, vol_frac_pol
    ) * (vol_frac_pol_0 / vol_frac_pol) ** (2 / 3)


def miao_dad_v(
    partial_molar_vol_pol,
    elasticity_c,
    vol_frac_pol_0,
    vol_frac_pol,
    vol_frac_pol_max=0.01,
):
    """Miao finite extensibility elasticity contribution."""
    stretch = (vol_frac_pol_0 / vol_frac_pol) ** (2 / 3)
    denom = 1 - (vol_frac_pol_max / vol_frac_pol) ** (2 / 3)
    return _elastic_prefactor(
        partial_molar_vol_pol, elasticity_c, vol_frac_pol
    ) * (
        (stretch / denom) + ((stretch - 1) / (denom**2)) * (1 - denom) - 0.5
    )
