import copy


def uvp_zscore(uvp, error_field='bs_std', inplace=False):
    """
    Calculate a zscore of a UVPSpec object using
    entry 'error_field' in its stats_array. This
    assumes that the UVPSpec object has been already
    mean subtracted using
    hera_pspec.uvpspec_utils.subtract_uvp().

    The resultant zscore is stored in the stats_array
    as error_field + "_zscore".

    Parameters
    ----------
    uvp : UVPSpec object

    error_field : str, optional
        Key of stats_array to use as z-score normalization.

    inplace : bool, optional
        If True, add zscores into input uvp, else
        make a copy of uvp and return with zscores.

    Returns
    -------
    if inplace:
        uvp : UVPSpec object
    """
    if not inplace:
        uvp = copy.deepcopy(uvp)

    # check error_field
    assert error_field in list(uvp.stats_array.keys()), "{} not found in stats_array" \
           .format(error_field)
    new_field = "{}_zscore".format(error_field)

    # iterate over spectral windows
    for i, spw in enumerate(uvp.spw_array):
        # iterate over polarizations
        for j, polpair in enumerate(uvp.polpair_array):
            # iterate over blpairs
            for k, blp in enumerate(uvp.blpair_array):
                key = (spw, blp, polpair)

                # calculate z-score: real and imag separately
                d = uvp.get_data(key)
                e = uvp.get_stats(error_field, key)
                zsc = d.real / e.real + 1j * d.imag / e.imag

                # set into uvp
                uvp.set_stats(new_field, key, zsc)

    if not inplace:
        return uvp