#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
from sklearn.metrics import roc_curve as skroc_curve


# In[2]:


def roc_curve(y_true, y_score):
    desc_score_indices = np.argsort(y_score, kind="mergesort")[::-1]
    y_true = y_true[desc_score_indices]
    y_score = y_score[desc_score_indices]
    distinct_value_indices = np.where(np.diff(y_score))[0]
    threshold_idxs = np.r_[distinct_value_indices, y_true.size - 1]
    tps = np.cumsum(y_true)[threshold_idxs]
    fps = 1 + threshold_idxs - tps
    thresholds = y_score[threshold_idxs]
    tps = np.r_[0, tps]
    fps = np.r_[0, fps]
    thresholds = np.r_[thresholds[0] + 1, thresholds]
    tpr = tps / tps[-1]
    fpr = fps / fps[-1]
    return fpr, tpr, thresholds


# In[3]:


for i in range(10):
    rng = np.random.RandomState(i)
    y_true = rng.randint(2, size=10)
    y_score = rng.randint(11, size=10) / 10
    ans1 = roc_curve(y_true, y_score)
    ans2 = skroc_curve(y_true, y_score, drop_intermediate=False)
    assert np.allclose(ans1[0], ans2[0])
    assert np.allclose(ans1[1], ans2[1])
    assert np.allclose(ans1[2], ans2[2])