from __future__ import annotations

import numpy as np
import pytest
from numpy.testing import assert_allclose
from pytest import approx

import boost_histogram as bh


@pytest.fixture
def v():
    h = bh.Histogram(bh.axis.Integer(0, 4), storage=bh.storage.Weight())
    h.fill([1, 1, 1, 2, 2, 3])
    return h.view()


def test_basic_view(v):
    assert_allclose(v.value, [0, 3, 2, 1])
    assert_allclose(v.variance, [0, 3, 2, 1])


def test_view_mul(v):
    v2 = v * 2
    assert_allclose(v2.value, [0, 6, 4, 2])
    assert_allclose(v2.variance, [0, 12, 8, 4])

    v2 = 2 * v
    assert_allclose(v2.value, [0, 6, 4, 2])
    assert_allclose(v2.variance, [0, 12, 8, 4])

    v2 = v * (-2)
    assert_allclose(v2.value, [0, -6, -4, -2])
    assert_allclose(v2.variance, [0, 12, 8, 4])

    v *= 2
    assert_allclose(v.value, [0, 6, 4, 2])
    assert_allclose(v.variance, [0, 12, 8, 4])


def test_view_div(v):
    v2 = v / 2
    assert_allclose(v2.value, [0, 1.5, 1, 0.5])
    assert_allclose(v2.variance, [0, 0.75, 0.5, 0.25])

    v2 = v / (-0.5)
    assert_allclose(v2.value, [0, -6, -4, -2])
    assert_allclose(v2.variance, [0, 12, 8, 4])

    v2 = 1 / v[1:]
    assert_allclose(v2.value, [1 / 3, 1 / 2, 1])
    assert_allclose(v2.variance, [1 / 3, 1 / 2, 1])

    v /= 0.5
    assert_allclose(v.value, [0, 6, 4, 2])
    assert_allclose(v.variance, [0, 12, 8, 4])


def test_view_add(v):
    v2 = v + 1
    assert_allclose(v2.value, [1, 4, 3, 2])
    assert_allclose(v2.variance, [1, 4, 3, 2])

    v2 = v + 2
    assert_allclose(v2.value, [2, 5, 4, 3])
    assert_allclose(v2.variance, [4, 7, 6, 5])

    v2 = 2 + v
    assert_allclose(v2.value, [2, 5, 4, 3])
    assert_allclose(v2.variance, [4, 7, 6, 5])

    v2 = v.copy()
    v2 += 2
    assert_allclose(v2.value, [2, 5, 4, 3])
    assert_allclose(v2.variance, [4, 7, 6, 5])

    v2 = v + v
    assert_allclose(v2.value, v.value * 2)
    assert_allclose(v2.variance, v.variance * 2)


def test_view_sub(v):
    v2 = v - 1
    assert_allclose(v2.value, [-1, 2, 1, 0])
    assert_allclose(v2.variance, [1, 4, 3, 2])

    v2 = v - 2
    assert_allclose(v2.value, [-2, 1, 0, -1])
    assert_allclose(v2.variance, [4, 7, 6, 5])

    v2 = 1 - v
    assert_allclose(v2.value, [1, -2, -1, 0])
    assert_allclose(v2.variance, [1, 4, 3, 2])

    v2 = v.copy()
    v2 -= 2
    assert_allclose(v2.value, [-2, 1, 0, -1])
    assert_allclose(v2.variance, [4, 7, 6, 5])

    v2 = v - v
    assert_allclose(v2.value, [0, 0, 0, 0])
    assert_allclose(v2.variance, v.variance * 2)


def test_view_unary(v):
    v2 = +v
    assert_allclose(v.value, v2.value)
    assert_allclose(v.variance, v2.variance)

    v2 = -v
    assert_allclose(-v.value, v2.value)
    assert_allclose(v.variance, v2.variance)


def test_view_add_same(v):
    v2 = v + v

    assert_allclose(v.value * 2, v2.value)
    assert_allclose(v.variance * 2, v2.variance)

    v2 = v + v[1]
    assert_allclose(v.value + 3, v2.value)
    assert_allclose(v.variance + 3, v2.variance)

    v2 = v + bh.accumulators.WeightedSum(5, 6)
    assert_allclose(v.value + 5, v2.value)
    assert_allclose(v.variance + 6, v2.variance)

    with pytest.raises(TypeError):
        v2 = v + bh.accumulators.WeightedMean(1, 2, 5, 6)


def test_view_assign(v):
    v[...] = [[4, 1], [5, 2], [6, 1], [7, 2]]

    assert_allclose(v.value, [4, 5, 6, 7])
    assert_allclose(v.variance, [1, 2, 1, 2])


def test_view_assign_mean():
    h = bh.Histogram(bh.axis.Integer(0, 4), storage=bh.storage.Mean())
    m = h.copy().view()

    h[...] = [[10, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]
    assert_allclose(h.view().count, [10, 4, 7, 10])
    assert_allclose(h.view().value, [2, 5, 8, 11])
    assert_allclose(h.view().variance, [3, 6, 9, 12])

    # Make sure this really was a copy
    assert m.count[0] != 10

    # Assign directly on view
    m[...] = [[10, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]

    assert_allclose(m.count, [10, 4, 7, 10])
    assert_allclose(m.value, [2, 5, 8, 11])
    assert_allclose(m.variance, [3, 6, 9, 12])
    # Note: if counts <= 1, variance is undefined


def test_view_assign_wmean():
    h = bh.Histogram(bh.axis.Integer(0, 4), storage=bh.storage.WeightedMean())

    w = h.copy().view()

    h[...] = [[10, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]

    assert_allclose(h.view().sum_of_weights, [10, 5, 9, 13])
    assert_allclose(h.view().sum_of_weights_squared, [2, 6, 10, 14])
    assert_allclose(h.view().value, [3, 7, 11, 15])
    assert_allclose(h.view().variance, [4, 8, 12, 16])

    # Make sure this really was a copy
    assert w.sum_of_weights[0] != 10

    # Assign directly on view
    w[...] = [[10, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]

    assert_allclose(w.sum_of_weights, [10, 5, 9, 13])
    assert_allclose(w.sum_of_weights_squared, [2, 6, 10, 14])
    assert_allclose(w.value, [3, 7, 11, 15])
    assert_allclose(w.variance, [4, 8, 12, 16])
    # Note: if sum_of_weights <= 1, variance is undefined

    w[0] = [9, 1, 2, 3]
    assert w.sum_of_weights[0] == 9
    assert w[0].sum_of_weights_squared == 1
    assert w.value[0] == 2
    assert w[0].variance == 3


# Issue #826 - accessing fields of a 0-d MeanView raised "iteration over a 0-d array"
def test_0d_mean_view():
    h = bh.Histogram(bh.axis.Integer(0, 4), storage=bh.storage.Mean())
    h.fill([0, 1, 2, 3, 0, 1], sample=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
    profile = h.project()
    assert profile.shape == ()

    # view() should return a 0-d MeanView without error
    v = profile.view()
    assert v.shape == ()

    # Field access on a 0-d MeanView should work
    assert v.count == approx(6)
    assert v.value == approx((1 + 2 + 3 + 4 + 5 + 6) / 6)
    assert v["value"] == approx((1 + 2 + 3 + 4 + 5 + 6) / 6)

    # values(), variances(), and counts() on the histogram should work
    assert profile.values() == approx((1 + 2 + 3 + 4 + 5 + 6) / 6)
    assert profile.counts() == approx(6)
    assert profile.variances() == approx(v.variance / 6)


def test_0d_weighted_mean_view():
    h = bh.Histogram(bh.axis.Integer(0, 4), storage=bh.storage.WeightedMean())
    h.fill([0, 1, 2, 3], sample=[1.0, 2.0, 3.0, 4.0], weight=[1.0, 1.0, 1.0, 1.0])
    profile = h.project()
    assert profile.shape == ()

    # view() should return a 0-d WeightedMeanView without error
    v = profile.view()
    assert v.shape == ()

    # Field access on a 0-d WeightedMeanView should work
    assert v.value == approx(2.5)
    assert v.sum_of_weights == approx(4.0)
    assert v["value"] == approx(2.5)

    # values() and counts() on the histogram should work
    assert profile.values() == approx(2.5)
    assert profile.counts() == approx(4.0)


# Issue #696
def test_view_cumsum():
    h = bh.Histogram(
        bh.axis.Integer(1, 10, underflow=True, overflow=False),
        storage=bh.storage.Weight(),
    )
    h.fill([2, 3], weight=[1.5, 2.5])

    view = h.view()
    c_view = np.cumsum(view)
    assert c_view.value == approx(np.cumsum(view.value))
    assert c_view.variance == approx(np.cumsum(view.variance))
