#!/usr/bin/env python
#
# Copyright (c) 2018 10X Genomics, Inc. All rights reserved.
#

"""
Simple Good-Turing estimator.
Based on S implementation in
  William A. Gale & Geoffrey Sampson (1995) Good-turing frequency estimation without tears,
  Journal of Quantitative Linguistics, 2:3, 217-237, DOI: 10.1080/09296179508590051
"""

import numpy as np
import scipy.stats as sp_stats


class SimpleGoodTuringError(Exception):
    pass


def _averaging_transform(r, nr):
    d = np.concatenate((np.ones(1, dtype=int), np.diff(r)))
    dr = np.concatenate((
        0.5 * (d[1:] + d[0:-1]),
        np.array((d[-1],), dtype=float),
    ))
    return nr.astype(float)/dr


def _rstest(r, coef):
    return r * np.power(1 + 1/r, 1 + coef)


def simple_good_turing(xr, xnr):
    """Make a Simple Good-Turing estimate of the frequencies.
    Args:
      xr (np.array(int)): Non-zero item frequencies
      xnr (np.array(int)): Non-zero frequencies of frequencies
    Returns:
      (rstar (np.array(float)), p0 (float)):
        rstar: The adjusted non-zero frequencies
        p0: The total probability of unobserved items
    """

    xr = xr.astype(float)
    xnr = xnr.astype(float)

    xN = np.sum(xr*xnr)

    # Get Linear Good-Turing estimate
    xnrz = _averaging_transform(xr, xnr)
    slope, _intercept, _, _, _ = sp_stats.linregress(np.log(xr), np.log(xnrz))

    if slope > -1:
        raise SimpleGoodTuringError(
            "The log-log slope is > -1 (%d); the SGT estimator is not applicable to these data." % slope)

    xrst = _rstest(xr, slope)
    xrstrel = xrst/xr

    # Get traditional Good-Turing estimate
    xrtry = xr == np.concatenate((xr[1:]-1, np.zeros(1)))
    xrstarel = np.zeros(len(xr))
    xrstarel[xrtry] = (xr[xrtry]+1) / xr[xrtry] * np.concatenate((xnr[1:], np.zeros(1)))[xrtry] / xnr[xrtry]

    # Determine when to switch from GT to LGT estimates
    tursd = np.ones(len(xr))
    for i in range(len(xr)):
        if xrtry[i]:
            tursd[i] = float(i+2) / xnr[i] * np.sqrt(xnr[i+1] * (1 + xnr[i+1]/xnr[i]))

    xrstcmbrel = np.zeros(len(xr))
    useturing = True
    for r in range(len(xr)):
        if not useturing:
            xrstcmbrel[r] = xrstrel[r]
        else:
            if np.abs(xrstrel[r]-xrstarel[r]) * (1+r)/tursd[r] > 1.65:
                xrstcmbrel[r] = xrstarel[r]
            else:
                useturing = False
                xrstcmbrel[r] = xrstrel[r]

    # Renormalize the probabilities for observed objects
    sumpraw = np.sum(xrstcmbrel * xr * xnr / xN)

    xrstcmbrel = xrstcmbrel * (1 - xnr[0] / xN) / sumpraw
    p0 = xnr[0]/xN

    return (xr * xrstcmbrel, p0)


def sgt_proportions(frequencies):
    """Use Simple Good-Turing estimate to adjust for unobserved items
    Args:
      frequencies (np.array(int)): Nonzero frequencies of items
    Returns:
      (pstar (np.array(float)), p0 (float)):
        pstar: The adjusted non-zero proportions
        p0: The total probability of unobserved items
    """
    if len(frequencies) == 0:
        raise ValueError("Input frequency vector is empty")
    if np.count_nonzero(frequencies) != len(frequencies):
        raise ValueError("Frequencies must be greater than zero")

    freqfreqs = np.bincount(frequencies)
    assert freqfreqs[0] == 0
    use_freqs = np.flatnonzero(freqfreqs)

    if len(use_freqs) < 10:
        raise SimpleGoodTuringError("Too few non-zero frequency items (%d). Aborting SGT." % len(use_freqs))

    rstar, p0 = simple_good_turing(use_freqs, freqfreqs[use_freqs])

    # rstar contains the smoothed frequencies.
    # Map each original frequency r to its smoothed rstar.
    rstar_dict = dict(zip(use_freqs, rstar))

    rstar_sum = np.sum(freqfreqs[use_freqs] * rstar)
    rstar_i = np.fromiter((rstar_dict[f] for f in frequencies),
                          dtype=float, count=len(frequencies))
    pstar = (1 - p0) * (rstar_i / rstar_sum)

    assert np.isclose(p0 + np.sum(pstar), 1)
    return (pstar, p0)


def test_prosody():
    data = (
        (1, 120),
        (2, 40),
        (3, 24),
        (4, 13),
        (5, 15),
        (6, 5),
        (7, 11),
        (8, 2),
        (9, 2),
        (10, 1),
        (12, 3),
        (14, 2),
        (15, 1),
        (16, 1),
        (17, 3),
        (19, 1),
        (20, 3),
        (21, 2),
        (23, 3),
        (24, 3),
        (25, 3),
        (26, 2),
        (27, 2),
        (28, 1),
        (31, 2),
        (32, 2),
        (33, 1),
        (34, 2),
        (36, 2),
        (41, 3),
        (43, 1),
        (45, 3),
        (46, 1),
        (47, 1),
        (50, 1),
        (71, 1),
        (84, 1),
        (101, 1),
        (105, 1),
        (121, 1),
        (124, 1),
        (146, 1),
        (162, 1),
        (193, 1),
        (199, 1),
        (224, 1),
        (226, 1),
        (254, 1),
        (257, 1),
        (339, 1),
        (421, 1),
        (456, 1),
        (481, 1),
        (483, 1),
        (1140, 1),
        (1256, 1),
        (1322, 1),
        (1530, 1),
        (2131, 1),
        (2395, 1),
        (6925, 1),
        (7846, 1),
    )

    # Computed using R 3.5.1 w/ the Gale S code
    expect_p0 = 0.003883244
    expect_rstar = np.array((
        0.7628079,
        1.706448,
        2.679796,
        3.663988,
        4.653366,
        5.645628,
        6.63966,
        7.634856,
        8.63086,
        9.627446,
        11.62182,
        13.61725,
        14.61524,
        15.61336,
        16.6116,
        18.60836,
        19.60685,
        20.6054,
        22.60264,
        23.60133,
        24.60005,
        25.5988,
        26.59759,
        27.59639,
        30.59294,
        31.59183,
        32.59073,
        33.58964,
        35.58751,
        40.58235,
        42.58035,
        44.57836,
        45.57738,
        46.57641,
        49.57351,
        70.55399,
        83.54229,
        100.5272,
        104.5237,
        120.5097,
        123.507,
        145.4879,
        161.474,
        192.4472,
        198.4421,
        223.4205,
        225.4188,
        253.3947,
        256.3922,
        338.3218,
        420.2514,
        455.2215,
        480.2,
        482.1983,
        1138.636,
        1254.537,
        1320.48,
        1528.302,
        2128.788,
        2392.562,
        6918.687,
        7838.899,
    ))

    xr = np.array([d[0] for d in data], dtype=int)
    xnr = np.array([d[1] for d in data], dtype=int)

    rstar, p0 = simple_good_turing(xr, xnr)

    assert np.abs(p0 - expect_p0) < 1e-9
    assert np.all(np.abs(rstar - expect_rstar) < 1e-3)
    assert np.all((np.abs(rstar - expect_rstar))/expect_rstar < 1e-4)
