#!/usr/bin/env python

#  Copyright (c) 2006, The Regents of the University of California, through
#  Lawrence Berkeley National Laboratory (subject to receipt of any required
#  approvals from the U.S. Dept. of Energy).  All rights reserved.

#  This software is distributed under the new BSD Open Source License.
#  <http://www.opensource.org/licenses/bsd-license.html>
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#  (1) Redistributions of source code must retain the above copyright notice,
#  this list of conditions and the following disclaimer.
#
#  (2) Redistributions in binary form must reproduce the above copyright
#  notice, this list of conditions and the following disclaimer in the
#  documentation and or other materials provided with the distribution.
#
#  (3) Neither the name of the University of California, Lawrence Berkeley
#  National Laboratory, U.S. Dept. of Energy nor the names of its contributors
#  may be used to endorse or promote products derived from this software
#  without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.

import unittest
from math import log, sqrt
from typing import Tuple

import pytest
from numpy import all, array, float64, ones, zeros
from scipy.stats import entropy

from weblogo import (
    GhostscriptAPI,
    LogoData,
    LogoFormat,
    LogoOptions,
    equiprobable_distribution,
    parse_prior,
)
from weblogo.color import Color
from weblogo.colorscheme import ColorScheme, IndexColor, RefSeqColor, SymbolColor
from weblogo.logomath import Dirichlet, Gamma
from weblogo.seq import (
    Alphabet,
    unambiguous_dna_alphabet,
    unambiguous_protein_alphabet,
    unambiguous_rna_alphabet,
)
from weblogo.utils import ArgumentError


class test_logoformat(unittest.TestCase):
    def test_options(self) -> None:
        LogoOptions()


def test_logoformat_errors() -> None:
    logodata = LogoData()
    logodata.length = 100

    # Negactive logo_margin
    with pytest.raises(ArgumentError):
        logooptions = LogoOptions()
        logooptions.logo_margin = -1
        LogoFormat(logodata, logooptions)

    # logo_start before start of sequecne
    with pytest.raises(ArgumentError):
        logooptions = LogoOptions()
        logooptions.first_index = 10
        logooptions.logo_start = -10
        LogoFormat(logodata, logooptions)

    # logo_end before logo_staRT
    with pytest.raises(ArgumentError):
        logooptions = LogoOptions()
        logooptions.first_index = 1
        logooptions.logo_start = 10
        logooptions.logo_end = -10
        LogoFormat(logodata, logooptions)

    # logo_end past lenght of sequence
    with pytest.raises(ArgumentError):
        logooptions = LogoOptions()
        logooptions.first_index = 1
        logooptions.logo_start = 10
        logooptions.logo_end = 200
        LogoFormat(logodata, logooptions)

    # No alphabet
    with pytest.raises(ArgumentError):
        logooptions = LogoOptions()
        logooptions.first_index = 1
        logooptions.logo_start = 10
        logooptions.logo_end = 20
        LogoFormat(logodata, logooptions)

    with pytest.raises(ArgumentError):
        logooptions = LogoOptions()
        logooptions.yaxis_scale = -1
        LogoFormat(logodata, logooptions)


def test_logoformats() -> None:
    # Make sure all different logo option code gets run
    logodata = LogoData()
    logodata.alphabet = unambiguous_rna_alphabet
    logodata.length = 100

    logooptions = LogoOptions()
    logooptions.fineprint = ""
    logooptions.xaxis_label = ""
    logooptions.yaxis_label = "Label"
    LogoFormat(logodata, logooptions)

    logooptions.yaxis_label = ""
    logooptions.unit_name = "probability"
    LogoFormat(logodata, logooptions)

    logooptions.show_yaxis = False
    LogoFormat(logodata, logooptions)

    logooptions.yaxis_label = "Label"
    logooptions.show_ends = True
    logooptions.show_xaxis = True
    LogoFormat(logodata, logooptions)

    logooptions.rotate_numbers = True
    LogoFormat(logodata, logooptions)

    logooptions.show_xaxis = False
    LogoFormat(logodata, logooptions)

    logodata.alphabet = Alphabet("ABCD")
    LogoFormat(logodata, logooptions)


class test_ghostscript(unittest.TestCase):
    def test_version(self) -> None:
        GhostscriptAPI().version()


class test_parse_prior(unittest.TestCase):
    def test_parse_prior_none(self) -> None:
        self.assertEqual(None, parse_prior(None, unambiguous_protein_alphabet))
        self.assertEqual(None, parse_prior("none", unambiguous_protein_alphabet))

    def test_parse_prior_equiprobable(self) -> None:
        self.assertTrue(
            all(
                20.0 * equiprobable_distribution(20)
                == parse_prior(
                    "equiprobable", unambiguous_protein_alphabet, weight=20.0
                )
            )
        )

        self.assertTrue(
            all(
                1.2 * equiprobable_distribution(3)
                == parse_prior(" equiprobablE  ", Alphabet("123"), 1.2)
            )
        )

    def test_parse_prior_percentage(self) -> None:
        # print(parse_prior('50%', unambiguous_dna_alphabet, 1.))
        self.assertTrue(
            all(
                equiprobable_distribution(4)
                == parse_prior("50%", unambiguous_dna_alphabet, 1.0)
            )
        )

        self.assertTrue(
            all(
                equiprobable_distribution(4)
                == parse_prior(" 50.0 % ", unambiguous_dna_alphabet, 1.0)
            )
        )

        self.assertTrue(
            all(
                array((0.3, 0.2, 0.2, 0.3), float64)
                == parse_prior(" 40.0 % ", unambiguous_dna_alphabet, 1.0)
            )
        )

    def test_parse_prior_float(self) -> None:
        self.assertTrue(
            all(
                equiprobable_distribution(4)
                == parse_prior("0.5", unambiguous_dna_alphabet, 1.0)
            )
        )

        self.assertTrue(
            all(
                equiprobable_distribution(4)
                == parse_prior(" 0.500 ", unambiguous_dna_alphabet, 1.0)
            )
        )

        self.assertTrue(
            all(
                array((0.3, 0.2, 0.2, 0.3), float64)
                == parse_prior(" 0.40 ", unambiguous_dna_alphabet, 1.0)
            )
        )

    def test_auto(self) -> None:
        self.assertTrue(
            all(
                2.0 * equiprobable_distribution(4)
                == parse_prior("auto", unambiguous_dna_alphabet)
            )
        )
        self.assertTrue(
            all(
                2.0 * equiprobable_distribution(4)
                == parse_prior("automatic", unambiguous_dna_alphabet)
            )
        )

        parse_prior("automatic", unambiguous_protein_alphabet)
        parse_prior("E. coli", unambiguous_dna_alphabet)

    def test_weight(self) -> None:
        self.assertTrue(
            all(
                2.0 * equiprobable_distribution(4)
                == parse_prior("automatic", unambiguous_dna_alphabet)
            )
        )
        self.assertTrue(
            all(
                123.123 * equiprobable_distribution(4)
                == parse_prior("auto", unambiguous_dna_alphabet, 123.123)
            )
        )

    def test_explicit(self) -> None:
        s = "{'A':10, 'C':40, 'G':40, 'T':10}"
        p = array((10, 40, 40, 10), float64) * 2.0 / 100.0
        self.assertTrue(all(p == parse_prior(s, unambiguous_dna_alphabet)))


def test_parse_prior_error() -> None:
    with pytest.raises(ValueError):
        parse_prior("0.5", unambiguous_protein_alphabet, weight=-10000.0)

    with pytest.raises(ValueError):
        s = "{'A':10, 'C':40, 'G':40, 'T':10}"
        parse_prior(s, unambiguous_protein_alphabet)

    with pytest.raises(ValueError):
        s = "{'A':'ljkasd', 'C':40, 'G':40, 'T':10}"
        parse_prior(s, unambiguous_dna_alphabet)

    with pytest.raises(ValueError):
        s = "asjnd"
        parse_prior(s, unambiguous_dna_alphabet)


class test_logooptions(unittest.TestCase):
    def test_create(self) -> None:
        opt = LogoOptions()
        opt.small_fontsize = 10
        repr(opt)

        opt = LogoOptions(logo_title="sometitle")  # type: ignore
        self.assertEqual(opt.logo_title, "sometitle")


class test_colorscheme(unittest.TestCase):
    def test_symbol_color(self) -> None:
        sc = SymbolColor("ABC", "black", "Because")
        self.assertEqual(sc.description, "Because")
        self.assertEqual(sc.symbol_color(0, "A", 0), Color.by_name("black"))
        self.assertEqual(sc.symbol_color(1, "D", 0), None)

    def test_index_color(self) -> None:
        ic = IndexColor([1, 3], "black", "Because")
        self.assertEqual(ic.description, "Because")
        self.assertEqual(ic.symbol_color(0, "A", 0), None)
        self.assertEqual(ic.symbol_color(1, "A", 0), Color.by_name("black"))

    def test_ref_seq_color(self) -> None:
        rc = RefSeqColor("ABC", "black", "Because")
        self.assertEqual(rc.description, "Because")

        self.assertEqual(rc.symbol_color(0, "A", 0), Color.by_name("black"))
        self.assertEqual(rc.symbol_color(1, "A", 0), None)
        self.assertEqual(rc.symbol_color(2, "A", 0), None)

        self.assertEqual(rc.symbol_color(0, "B", 0), None)
        self.assertEqual(rc.symbol_color(1, "B", 0), Color.by_name("black"))
        self.assertEqual(rc.symbol_color(2, "B", 0), None)

        self.assertEqual(rc.symbol_color(0, "C", 0), None)
        self.assertEqual(rc.symbol_color(1, "C", 0), None)
        self.assertEqual(rc.symbol_color(2, "C", 0), Color.by_name("black"))

    def test_colorscheme(self) -> None:
        cs = ColorScheme(
            [
                SymbolColor("G", "orange"),
                SymbolColor("TU", "red"),
                SymbolColor("C", "blue"),
                SymbolColor("A", "green"),
            ],
            title="title",
            description="description",
        )

        self.assertEqual(cs.symbol_color(1, "G", 1), Color.by_name("orange"))
        self.assertEqual(cs.symbol_color(1, "T", 1), Color.by_name("red"))
        self.assertEqual(cs.symbol_color(1, "C", 1), Color.by_name("blue"))
        self.assertEqual(cs.symbol_color(1, "A", 1), Color.by_name("green"))
        self.assertEqual(cs.symbol_color(1, "X", 1), cs.default_color)

        cs = ColorScheme(
            [
                SymbolColor("G", "orange"),
                SymbolColor("TU", "red"),
                SymbolColor("C", "blue"),
                SymbolColor("A", "green"),
            ],
            title="title",
            description="description",
            alphabet=Alphabet("GTUCA"),
        )
        self.assertRaises(KeyError, cs.symbol_color, 1, "X", 1)


class test_color(unittest.TestCase):
    def test_color_names(self) -> None:
        names = Color.names()
        self.assertEqual(len(names), 147)
        for n in names:
            c = Color.by_name(n)
            self.assertTrue(c is not None)

    def test_color_components(self) -> None:
        white = Color.by_name("white")
        self.assertEqual(1.0, white.red)
        self.assertEqual(1.0, white.green)
        self.assertEqual(1.0, white.blue)

        c = Color(0.3, 0.4, 0.2)
        self.assertEqual(0.3, c.red)
        self.assertEqual(0.4, c.green)
        self.assertEqual(0.2, c.blue)

        c = Color(0, 128, 0)
        self.assertEqual(0.0, c.red)
        self.assertEqual(128.0 / 255.0, c.green)
        self.assertEqual(0.0, c.blue)

    def test_color_from_rgb(self) -> None:
        white = Color.by_name("white")
        self.assertEqual(white, Color(1.0, 1.0, 1.0))
        self.assertEqual(white, Color(255, 255, 255))
        self.assertEqual(white, Color.from_rgb(1.0, 1.0, 1.0))
        self.assertEqual(white, Color.from_rgb(255, 255, 255))

    def test_color_from_hsl(self) -> None:
        red = Color.by_name("red")
        lime = Color.by_name("lime")
        saddlebrown = Color.by_name("saddlebrown")
        darkgreen = Color.by_name("darkgreen")
        blue = Color.by_name("blue")
        Color.by_name("green")
        self.assertEqual(red, Color.from_hsl(0, 1.0, 0.5))
        self.assertEqual(lime, Color.from_hsl(120, 1.0, 0.5))
        self.assertEqual(blue, Color.from_hsl(240, 1.0, 0.5))
        self.assertEqual(Color.by_name("gray"), Color.from_hsl(0, 0, 0.5))
        self.assertEqual(saddlebrown, Color.from_hsl(25, 0.76, 0.31))
        self.assertEqual(darkgreen, Color.from_hsl(120, 1.0, 0.197))

    def test_color_by_name(self) -> None:
        white = Color.by_name("white")
        self.assertEqual(white, Color.by_name("white"))
        self.assertEqual(white, Color.by_name("WHITE"))
        self.assertEqual(white, Color.by_name(" wHiTe \t\n\t"))
        self.assertEqual(Color(255, 255, 240), Color.by_name("ivory"))
        self.assertEqual(Color(70, 130, 180), Color.by_name("steelblue"))
        self.assertEqual(Color(0, 128, 0), Color.by_name("green"))

    def test_color_from_invalid_name(self) -> None:
        self.assertRaises(ValueError, Color.by_name, "not_a_color")

    def test_color_clipping(self) -> None:
        red = Color.by_name("red")
        self.assertEqual(red, Color(255, 0, 0))
        self.assertEqual(red, Color(260, -10, 0))
        self.assertEqual(red, Color(1.1, -0.0, -1.0))
        self.assertEqual(Color(1.0001, 213.0, 1.2).red, 1.0)
        self.assertEqual(Color(-0.001, -2183.0, -1.0).red, 0.0)
        self.assertEqual(Color(1.0001, 213.0, 1.2).green, 1.0)
        self.assertEqual(Color(-0.001, -2183.0, -1.0).green, 0.0)
        self.assertEqual(Color(1.0001, 213.0, 1.2).blue, 1.0)
        self.assertEqual(Color(-0.001, -2183.0, -1.0).blue, 0.0)

    def test_color_fail_on_mixed_type(self) -> None:
        self.assertRaises(TypeError, Color.from_rgb, 1, 1, 1.0)
        self.assertRaises(TypeError, Color.from_rgb, 1.0, 1, 1.0)

    def test_color_red(self) -> None:
        # Check Usage comment in Color
        red = Color.by_name("red")
        self.assertEqual(red, Color(255, 0, 0))
        self.assertEqual(red, Color(1.0, 0.0, 0.0))
        self.assertEqual(red, Color.from_rgb(1.0, 0.0, 0.0))
        self.assertEqual(red, Color.from_rgb(255, 0, 0))
        self.assertEqual(red, Color.from_hsl(0.0, 1.0, 0.5))
        self.assertEqual(red, Color.from_string("red"))
        self.assertEqual(red, Color.from_string("RED"))
        self.assertEqual(red, Color.from_string("#F00"))
        self.assertEqual(red, Color.from_string("#FF0000"))
        self.assertEqual(red, Color.from_string("rgb(255, 0, 0)"))
        self.assertEqual(red, Color.from_string("rgb(100%, 0%, 0%)"))
        self.assertEqual(red, Color.from_string("hsl(0, 100%, 50%)"))

    def test_color_from_string(self) -> None:
        Color(128, 0, 128)  # purple
        red = Color(255, 0, 0)
        skyblue = Color(135, 206, 235)

        red_strings = (
            "red",
            "ReD",
            "RED",
            "   Red \t",
            "#F00",
            "#FF0000",
            "rgb(255, 0, 0)",
            "rgb(100%, 0%, 0%)",
            "hsl(0, 100%, 50%)",
        )
        for s in red_strings:
            self.assertEqual(red, Color.from_string(s))

        skyblue_strings = (
            "skyblue",
            "SKYBLUE",
            "  \t\n SkyBlue  \t",
            "#87ceeb",
            "rgb(135,206,235)",
        )
        for s in skyblue_strings:
            self.assertEqual(skyblue, Color.from_string(s))

        self.assertRaises(ValueError, Color.from_string, "#not_a_color")
        self.assertRaises(ValueError, Color.from_string, "rgb(not_a_color)")
        self.assertRaises(ValueError, Color.from_string, "hsl(not_a_color)")
        self.assertRaises(ValueError, Color.from_string, "not_a_color")

    def test_color_equality(self) -> None:
        c1 = Color(123, 99, 12)
        c2 = Color(123, 99, 12)
        self.assertEqual(c1, c2)
        self.assertNotEqual(c1, "not_a_color")


class test_gamma(unittest.TestCase):
    def test_create(self) -> None:
        a = 1.213
        b = 3.210
        g = Gamma(a, b)
        self.assertEqual(g.alpha, a)
        self.assertEqual(g.beta, b)

    def test_mean_variance(self) -> None:
        g = Gamma.from_mean_variance(2.0, 3.0)
        self.assertEqual(g.mean(), 2.0)
        self.assertEqual(g.variance(), 3.0)

        g = Gamma.from_mean_variance(2.0123, 3.01283)
        self.assertEqual(g.mean(), 2.0123)
        self.assertEqual(g.variance(), 3.01283)

    def test_from_shape_scale(self) -> None:
        g = Gamma.from_shape_scale(1.0, 8.0)
        self.assertEqual(g.alpha, 1.0)
        self.assertEqual(g.beta, 1.0 / 8.0)

    def test_invalid_args(self) -> None:
        self.assertRaises(ValueError, Gamma, 1.0, -1.0)
        self.assertRaises(ValueError, Gamma, 0.0, 1.0)
        self.assertRaises(ValueError, Gamma, 1.0, 0.0)

    def test_sample(self) -> None:
        m = 10.0
        v = 2.0
        g = Gamma.from_mean_variance(m, v)
        # print(g.alpha, g.beta)
        S = 1000
        total = 0.0
        for s in range(S):
            total += g.sample()
        mean = total / S

        # The estimated mean will differ from true mean by a small amount

        error = 4.0 * sqrt(g.variance() / S)
        # print(mean, m, error)
        self.assertTrue(abs(mean - m) < error)

    def test_pdf(self) -> None:
        m = 3.0
        v = 2.0
        g = Gamma.from_mean_variance(m, v)
        upper = 30.0

        norm = integrate(g.pdf, 0, upper)
        self.assertAlmostEqual(norm, 1.0)

        def fx(x: float) -> float:
            return x * g.pdf(x)

        mean = integrate(fx, 0, upper)
        self.assertAlmostEqual(mean, m)

        def fx2(x: float) -> float:
            return x * x * g.pdf(x)

        x2 = integrate(fx2, 0, upper)
        var = x2 - mean**2
        self.assertAlmostEqual(var, v)

    def test_cdf(self) -> None:
        m = 3.0
        v = 2.0
        g = Gamma.from_mean_variance(m, v)
        # Numerical integration
        S = 1000
        M = 10.0
        total_p = 0.0
        epsilon = 1e-4
        last = 0.0
        for s in range(S):
            x = s * M / S
            p = g.pdf(x) * M / S
            total_p += (last - p) / 2.0
            last = p
            # print(x, total_p, g.cdf(x))

            self.assertTrue((total_p - g.cdf(x)) < epsilon)

    def test_inverse_cdf(self) -> None:
        g = Gamma.from_mean_variance(2.34, 4)
        self.assertAlmostEqual(3.9, g.inverse_cdf(g.cdf(3.9)))
        self.assertAlmostEqual(1.92, g.inverse_cdf(g.cdf(1.92)))

        g = Gamma.from_mean_variance(10.34, 2)
        self.assertAlmostEqual(3.9, g.inverse_cdf(g.cdf(3.9)))
        self.assertAlmostEqual(10.92, g.inverse_cdf(g.cdf(10.92)))

        g = Gamma.from_mean_variance(10.34, 2)
        self.assertAlmostEqual(0.975, g.cdf(g.inverse_cdf(0.975)))
        self.assertAlmostEqual(0.025, g.cdf(g.inverse_cdf(0.025)))

        g = Gamma.from_mean_variance(1.34, 4)
        self.assertAlmostEqual(0.975, g.cdf(g.inverse_cdf(0.975)))
        self.assertAlmostEqual(0.025, g.cdf(g.inverse_cdf(0.025)))


class test_Dirichlet(unittest.TestCase):
    def test_init(self) -> None:
        Dirichlet(
            (
                1,
                1,
                1,
                1,
            )
        )

    def test_random(self) -> None:
        def do_test(alpha: Tuple[float, ...], samples: int = 1000) -> None:
            ent = zeros((samples,), float64)
            # alpha = ones( ( K,), Float64 ) * A/K

            # pt = zeros( (len(alpha) ,), Float64)
            d = Dirichlet(alpha)
            for s in range(samples):
                p = d.sample()
                # print(p)
                # pt +=p
                ent[s] = entropy(p)

            # print(pt/samples)

            m = mean(ent)
            v = var(ent)

            dm = d.mean_entropy()
            dv = d.variance_entropy()

            # print(alpha, ':', m, v, dm, dv)
            error = 4.0 * sqrt(v / samples)
            self.assertTrue(abs(m - dm) < error)
            self.assertTrue(abs(v - dv) < error)  # dodgy error estimate

        do_test((1.0, 1.0))
        do_test((2.0, 1.0))
        do_test((3.0, 1.0))
        do_test((4.0, 1.0))
        do_test((5.0, 1.0))
        do_test((6.0, 1.0))

        do_test((1.0, 1.0))
        do_test((20.0, 20.0))
        do_test((1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0))
        do_test((0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1))
        do_test((0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01))
        do_test((2.0, 6.0, 1.0, 1.0))

    def test_mean(self) -> None:
        alpha = ones((10,), float64) * 23.0
        d = Dirichlet(alpha)
        m = d.mean()
        self.assertAlmostEqual(m[2], 1.0 / 10)
        self.assertAlmostEqual(sum(m), 1.0)

    def test_covariance(self) -> None:
        alpha = ones((4,), float64)
        d = Dirichlet(alpha)
        cv = d.covariance()
        self.assertEqual(cv.shape, (4, 4))
        self.assertAlmostEqual(cv[0, 0], 1.0 * (1.0 - 1.0 / 4.0) / (4.0 * 5.0))
        self.assertAlmostEqual(cv[0, 1], -1 / (4.0 * 4.0 * 5.0))

    def test_mean_x(self) -> None:
        alpha = (1.0, 2.0, 3.0, 4.0)
        xx = (2.0, 2.0, 2.0, 2.0)
        m = Dirichlet(alpha).mean_x(xx)
        self.assertEqual(m, 2.0)

        xx2 = (2.0, 2.0, 2.0, 2.0, 2.0)
        self.assertRaises(ValueError, Dirichlet(alpha).mean_x, xx2)

        alpha = (1.0, 1.0, 1.0, 1.0)
        xx = (2.0, 3.0, 4.0, 3.0)
        m = Dirichlet(alpha).mean_x(xx)
        self.assertEqual(m, 3.0)

    def test_variance_x(self) -> None:
        alpha = (1.0, 1.0, 1.0, 1.0)
        xx = (2.0, 2.0, 2.0, 2.0)
        v = Dirichlet(alpha).variance_x(xx)
        self.assertAlmostEqual(v, 0.0)

        alpha = (1.0, 2.0, 3.0, 4.0)
        xx = (2.0, 0.0, 1.0, 10.0)
        v = Dirichlet(alpha).variance_x(xx)
        # print(v)
        # TODO: Don't actually know if this is correct

        xx2 = (2.0, 2.0, 2.0, 2.0, 2.0)
        self.assertRaises(ValueError, Dirichlet(alpha).variance_x, xx2)

    def test_relative_entropy(self) -> None:
        alpha = (2.0, 10.0, 1.0, 1.0)
        d = Dirichlet(alpha)
        pvec = (0.1, 0.2, 0.3, 0.4)

        rent = d.mean_relative_entropy(pvec)
        vrent = d.variance_relative_entropy(pvec)
        low, high = d.interval_relative_entropy(pvec, 0.95)

        # print()
        # print('> ', rent, vrent, low, high)

        # This test can fail randomly, but the precision from a few
        # thousand samples is low. Increasing samples, 1000->2000
        samples = 2000
        sent = zeros((samples,), float64)

        for s in range(samples):
            post = d.sample()
            e = -entropy(post)
            for k in range(4):
                e += -post[k] * log(pvec[k])
            sent[s] = e
        sent.sort()
        self.assertTrue(abs(sent.mean() - rent) < 4.0 * sqrt(vrent))
        self.assertAlmostEqual(sent.std(), sqrt(vrent), 1)
        self.assertTrue(abs(low - sent[int(samples * 0.025)]) < 0.2)
        self.assertTrue(abs(high - sent[int(samples * 0.975)]) < 0.2)


class _from_URL_fileopen_Tests(unittest.TestCase):
    def test_URLscheme(self) -> None:
        """test for http, https, or ftp scheme"""
        from weblogo.logo import _from_URL_fileopen

        broken_url = "file://foo.txt"
        self.assertRaises(ValueError, _from_URL_fileopen, (broken_url))


def mean(a) -> float:  # type: ignore
    return sum(a) / len(a)


def var(a) -> float:  # type: ignore
    return (sum(a * a) / len(a)) - mean(a) ** 2


def integrate(f, a, b, n=1000):  # type: ignore
    """
    Numerically integrate the function 'f' from 'a' to 'b' using a discretization with 'n' points.

    Args:
    - f -- A function that eats a float and returns a float.
    - a -- Lower integration bound (float)
    - b -- Upper integration bound (float)
    - n -- number of sample points (int)

    Status :
        Alpha (very primitive.)
    """
    h = (b - a) / (n - 1.0)
    total = 0.0
    for i in range(n):
        total += f(a + (i) * h)
    result = h * (total - 0.5 * f(a) - 0.5 * f(b))
    return result


if __name__ == "__main__":
    unittest.main()
