import torch


def interface_r_vec(polarization, n_i, n_f, th_i, th_f):
    """
    reflection amplitude (from Fresnel equations)
    polarization is either "s" or "p" for polarization
    n_i, n_f are (complex) refractive index for incident and final
    th_i, th_f are (complex) propegation angle for incident and final
    (in radians, where 0=normal). "th" stands for "theta".
    """
    if polarization == 's':
        ni_thi = torch.einsum('sij,skij->skji', n_i, torch.cos(th_i))
        nf_thf = torch.einsum('sij,skij->skji', n_f, torch.cos(th_f))
        return (ni_thi - nf_thf) / (ni_thi + nf_thf)
    elif polarization == 'p':
        nf_thi = torch.einsum('sij,skij->skji', n_f, torch.cos(th_i))
        ni_thf = torch.einsum('sij,skij->skji', n_i, torch.cos(th_f))
        return (nf_thi - ni_thf) / (nf_thi + ni_thf)
    else:
        raise ValueError("Polarization must be 's' or 'p'")