# -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import random
import pytest

from asv.step_detect import (solve_potts, solve_potts_autogamma, solve_potts_approx,
                             detect_regressions, golden_search, median, rolling_median_dev,
                             L1Dist, detect_steps)
from asv import step_detect


try:
    import numpy as np
    np.random.seed
    HAVE_NUMPY = True
except (ImportError, NameError):
    HAVE_NUMPY = False

try:
    from asv import _rangemedian
    HAVE_RANGEMEDIAN = True
except ImportError:
    HAVE_RANGEMEDIAN = False


@pytest.fixture(params=[
    "python",
    pytest.param("rangemedian", marks=pytest.mark.skipif(not HAVE_RANGEMEDIAN, reason="compiled asv._rangemedian required"))
])
def use_rangemedian(request):
    if request.param == "rangemedian":
        assert isinstance(step_detect.get_mu_dist([0], [1]), _rangemedian.RangeMedian)
        return True
    else:
        step_detect._rangemedian = None

        def restore():
            if HAVE_RANGEMEDIAN:
                step_detect._rangemedian = _rangemedian
        request.addfinalizer(restore)

        assert isinstance(step_detect.get_mu_dist([0], [1]), L1Dist)
        return False


@pytest.mark.skipif(not HAVE_NUMPY, reason="test needs numpy")
def test_solve_potts(use_rangemedian):
    np.random.seed(1234)

    # Easy case, exact solver
    y = [1, 1, 1, 2, 2, 2, 3, 3, 3]
    right, values, dists = solve_potts(y, w=[1]*len(y), gamma=0.1)
    assert right == [3, 6, 9]
    assert np.allclose(values, [1, 2, 3], atol=0)

    right, values, dists = solve_potts(y, w=[1]*len(y), gamma=8.0)
    assert right == [9]
    assert np.allclose(values, [2], atol=0)

    # l1 norm
    right, values, dists = solve_potts(y, w=[1]*len(y), gamma=0.1)
    assert right == [3, 6, 9]
    assert np.allclose(values, [1, 2, 3], atol=0)

    # Bigger case, exact solver
    t = np.arange(100)
    y0 = (+ 0.4 * (t >= 5)
          + 0.9 * (t >= 10)
          - 0.2 * (t >= 20)
          + 0.2 * (t >= 50)
          + 1.1 * (t >= 70))
    y = y0 + 0.05 * np.random.rand(y0.size)
    right, values, dists = solve_potts(y.tolist(), w=[1]*len(y), gamma=0.1)
    assert right == [5, 10, 20, 50, 70, 100]

    # Bigger case, approximative solver
    right, values, dists = solve_potts_approx(y.tolist(), w=[1]*len(y), gamma=0.1)
    assert right == [5, 10, 20, 50, 70, 100]

    # Larger case
    t = np.arange(3000)
    y0 = (+ 0.4 * (t >= 10)
          + 0.9 * (t >= 30)
          - 0.2 * (t >= 200)
          + 0.2 * (t >= 600)
          + 1.1 * (t >= 2500)
          - 0.5 * (t >= 2990))

    # Small amount of noise shouldn't disturb step finding
    y = y0 + 0.05 * np.random.randn(y0.size)
    right, values, dists, gamma = solve_potts_autogamma(y.tolist(), w=[1]*len(y))
    assert right == [10, 30, 200, 600, 2500, 2990, 3000]

    # Large noise should prevent finding any steps
    y = y0 + 5.0 * np.random.randn(y0.size)
    right, values, dists, gamma = solve_potts_autogamma(y.tolist(), w=[1]*len(y))
    assert right == [3000]

    # The routine shouldn't choke on datasets with 10k points.
    # Appending noisy data to weakly noisy data should retain the
    # steps in the former
    y = y0 + 0.025 * np.random.rand(y.size)
    ypad = 0.025 * np.random.randn(10000 - 3000)
    right, values, dists, gamma = solve_potts_autogamma(y.tolist() + ypad.tolist(),
                                                        w=[1]*(len(y)+len(ypad)))
    assert right == [10, 30, 200, 600, 2500, 2990, 3000, 10000]


@pytest.mark.skipif(not HAVE_NUMPY, reason="test needs numpy")
def test_autocorrelated():
    # Check that a low-amplitude cosine signal is not interpreted as
    # multiple steps
    j = np.arange(1000)
    y = 0.2 * np.cos(j/100.0) + 1.0 * (j >= 500)
    right, values, dists, gamma = solve_potts_autogamma(y.tolist(), w=[1]*len(y))
    assert right == [500, 1000]


def test_zero_variance():
    # Should not choke on this data
    y = [1.0]*1000
    right, values, dists, gamma = solve_potts_autogamma(y, w=[1]*len(y))
    assert right == [1000]


@pytest.mark.skipif(not HAVE_NUMPY, reason="test needs numpy")
def test_weighted():
    np.random.seed(1234)

    t = np.arange(100)
    y0 = (+ 0.4 * (t >= 5)
          + 0.9 * (t >= 10)
          - 0.2 * (t >= 20)
          + 0.2 * (t >= 50)
          + 1.1 * (t >= 70))
    y = y0 + 0.05 * np.random.rand(y0.size)

    y = y.tolist()
    w = [1]*len(y)

    y[15] = 2
    right, values, dists = solve_potts(y, w=w, gamma=0.1)
    assert right == [5, 10, 15, 16, 20, 50, 70, 100]

    steps = detect_steps(y, w=w)
    steps_pos = [s[0] for s in steps]
    assert steps_pos == [0, 5, 10, 15, 16, 20, 50, 70]

    w[15] = 0.1
    right, values, dists = solve_potts(y, w=w, gamma=0.1)
    assert right == [5, 10, 20, 50, 70, 100]

    steps = detect_steps(y, w=w)
    steps_pos = [s[0] for s in steps]
    assert steps_pos == [0, 5, 10, 20, 50, 70]

    # Missing weights and data should be handled properly
    y[35] = None
    w[23] = None

    steps = detect_steps(y, w=w)
    steps_pos = [s[0] for s in steps]
    assert steps_pos == [0, 5, 10, 20, 50, 70]


@pytest.mark.skipif(not HAVE_NUMPY, reason="test needs numpy")
def test_detect_regressions(use_rangemedian):
    np.random.seed(1234)

    for seed in [1234, 5678, 8901, 2345]:
        np.random.seed(seed)
        t = np.arange(4000)
        y = 0.7 * np.random.rand(t.size)

        y -= 0.3 * (t >= 1000)
        y += 2.0 * (t >= 3234 + (seed % 123))
        y += 2.0 * ((t >= 3264 + (seed % 53)) & (t < 3700))
        y -= 2.7 * ((t >= 3350) & (t < 3500 + (seed % 71)))

        y = y.tolist()
        y[123] = None
        y[1234] = np.nan
        steps = detect_steps(y)

        steps_lr = [(l, r) for l, r, _, _, _ in steps]
        k = steps[0][1]
        assert 990 <= k <= 1010
        assert steps_lr == [(0, k),
                            (k, 3234 + (seed % 123)),
                            (3234 + (seed % 123), 3264 + (seed % 53)),
                            (3264 + (seed % 53), 3350),
                            (3350, 3500 + (seed % 71)),
                            (3500 + (seed % 71), 3700),
                            (3700, 4000)]
        steps_v = [x[2] for x in steps]
        assert np.allclose(steps_v, [0.35, 0.05, 2.05, 4.05, 1.15, 4.05, 2.05], rtol=0.3)

        # The expected mean error is 0.7 <|U(0,1) - 1/2|> = 0.7/4
        steps_err = [x[4] for x in steps]
        assert np.allclose(steps_err, [0.7/4]*7, rtol=0.3)

        # Check detect_regressions
        new_value, best_value, regression_pos = detect_regressions(steps)
        assert regression_pos == [(3233 + (seed % 123), (3233 + (seed % 123) + 1),
                                   steps_v[1], min(steps_v[2:])),
                                  (3499 + (seed % 71), 3499 + (seed % 71) + 1,
                                   steps_v[4], min(steps_v[5:]))]
        assert np.allclose(best_value, 0.7/2 - 0.3, rtol=0.3, atol=0)
        assert np.allclose(new_value, 0.7/2 - 0.3 + 2, rtol=0.3, atol=0)

def test_golden_search():
    def f(x):
        return 1 + x**3 + x**4
    x = golden_search(f, -1, -0.25, xatol=1e-5, ftol=0)
    assert abs(x - (-3/4)) < 1e-4
    x = golden_search(f, -0.25, 0.25, xatol=1e-5, ftol=0)
    assert abs(x - (-0.25)) < 1e-4


def rolling_median_dev_naive(items):
    for j in range(1, len(items)):
        m = median(items[:j])
        d = sum(abs(x - m) for x in items[:j])
        yield m, d


def test_rolling_median():
    random.seed(1)

    datasets = [
        [1, 1, 10, 3, 5, 1, -16, -3, 4, 9],
        [random.gauss(0, 1) for j in range(500)]
    ]

    for x in datasets:
        got = list(rolling_median_dev(x))
        expected = rolling_median_dev_naive(x)
        for j, b in enumerate(expected):
            a = got[j]
            assert abs(a[0] - b[0]) < 1e-10, (a, b)
            assert abs(a[1] - b[1]) < 1e-10, (a, b)


def test_l1dist(use_rangemedian):
    random.seed(1)

    datasets = [
        ([1, 1, 10, 3, 5, 1, -16, -3, 4, 9], [1]*10),
        ([random.gauss(0, 1) for j in range(50)], [1]*50),
        ([1, 2, 3, 4], [1, 2, 1, 2])
    ]

    def median_iter(y, w):
        if all(ww == w[0] for ww in w):
            for m, d in rolling_median_dev(y):
                yield m, d
        else:
            for j in range(1, len(y) + 1):
                m = step_detect.weighted_median(y[:j], w[:j])
                d = sum(ww*abs(yy - m) for yy, ww in zip(y[:j], w[:j]))
                yield m, d

    for y, w in datasets:
        dist = step_detect.get_mu_dist(y, w)

        for i in range(len(y)):
            for p, (m2, d2) in enumerate(median_iter(y[i:], w[i:])):
                j = i + p

                m = dist.mu(i, j)
                d = dist.dist(i, j)

                assert m == m2, (i, j)
                assert abs(d - d2) < 1e-10, (i, j)


def test_regression_threshold():
    steps = [(0, 1,   1.0, 1.0, 0.0),
             (1, 2,   1.1, 1.1, 0.0),
             (2, 3,   2.0, 2.0, 0.0)]

    latest, best, pos = detect_regressions(steps, threshold=0.05, min_size=1)
    assert latest == 2
    assert best == 1
    assert pos == [(0, 1, 1.0, 1.1), (1, 2, 1.1, 2.0)]

    latest, best, pos = detect_regressions(steps, threshold=0.2, min_size=1)
    assert latest == 2
    assert best == 1
    assert pos == [(1, 2, 1.1, 2.0)]

    latest, best, pos = detect_regressions(steps, threshold=0.8, min_size=1)
    assert latest == 2
    assert best == 1
    assert pos == [(1, 2, 1.1, 2.0)]

    latest, best, pos = detect_regressions(steps, threshold=1.1, min_size=1)
    assert latest == None
    assert best == None
    assert pos == None

    steps = [(0, 1,   1.0, 1.0, 0.0),
             (1, 2,   1.3, 1.3, 0.0),
             (2, 3,   1.1, 1.1, 0.0)]

    latest, best, pos = detect_regressions(steps, threshold=0.2, min_size=1)
    assert latest == None
    assert best == None
    assert pos == None

    # Gradual change should result to a regression detected somewhere,
    # even if the individual steps are smaller than the threshold
    steps = [(0, 1,   1.0, 1.0, 0.0),
             (1, 2,   1.04, 1.04, 0.0),
             (2, 3,   1.08, 1.08, 0.0),]

    latest, best, pos = detect_regressions(steps, threshold=0.05)
    assert pos == [(0, 1, 1.0, 1.04)]


def test_zero_weight():
    t = list(range(50))
    y = [1.0 * (tt > 25) for tt in t]
    w = [0 if tt >= 20 and tt < 30 else 1 for tt in t]

    # Due to zero weight, part of the data is not assigned to belong to a step
    steps = detect_steps(y, w=w)
    steps_pos = [s[:2] for s in steps]
    assert steps_pos == [(0, 20), (30, 50)]

    # A small weight allows to find the actual step
    w = [1e-6 if ww == 0 else 1 for ww in w]
    steps = detect_steps(y, w=w)
    steps_pos = [s[:2] for s in steps]
    assert steps_pos == [(0, 26), (26, 50)]


def test_regression_min_size():
    steps = [(0, 2,   1.0, 1.0, 0.0),
             (2, 3,   0.0, 0.0, 0.0),
             (3, 5,   1.0, 1.0, 0.0)]

    latest, best, pos = detect_regressions(steps)
    assert latest is None and best is None and pos is None

    latest, best, pos = detect_regressions(steps, min_size=1)
    assert latest == 1.0
    assert best == 0.0
    assert pos == [(2, 3, 0.0, 1.0)]

    steps = [(0, 2,   1.0, 1.0, 0.0),
             (2, 3,   0.0, 0.0, 0.0),
             (3, 5,   2.0, 1.0, 0.0)]

    latest, best, pos = detect_regressions(steps)
    assert latest == 2.0
    assert best == 0.0
    assert pos == [(2, 3, 0.0, 2.0)]


def test_solve_potts_approx_bug():
    y = [2.9, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1,
         3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1,
         3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.2, 3.1, 3.1, 3.1, 3.1,
         3.1, 3.1, 3.1, 3.1, 2.9, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1]
    w = [86.1, 1.0, 1.0, 1.0, 0.8, 1.0, 0.8, 1.0, 1.0, 0.9, 0.9, 1.0,
         0.8, 1.0, 1.0, 1.0, 1.0, 0.9, 0.6, 0.9, 0.5, 1.0, 1.0, 1.0, 1.0,
         1.0, 1.0, 1.0, 1.0, 1.0, 0.9, 1.0, 1.0, 0.7, 1.0, 1.0, 1.0, 1.0,
         1.0, 1.0, 0.8, 0.8, 50.0, 0.8, 1.0, 1.0, 0.6, 0.8, 1.0, 1.0]
    gamma = 0.3

    r0, v0, d0 = solve_potts(y, w, gamma=gamma)
    r, v, d = solve_potts_approx(y, w, gamma=gamma)
    assert r == r0
