"""Tests behavior of structured arrays"""

import matplotlib.pyplot as plt
import numpy as np

from gumbi import *


def test_1d_array_plotting():
    stdzr = Standardizer(
        c={"μ": -5.30, "σ": 0.582},
        d={"μ": -0.307, "σ": 0.158},
        z={"μ": 2, "σ": 2},
        log_vars=["d", "c"],
        logit_vars=["z"],
    )
    m_pa = parray(m=np.arange(1, 5) / 10, stdzr=stdzr)
    r_upa = uparray("d", np.arange(1, 5) / 10 + 0.5, np.arange(1, 5) / 100 * 2, stdzr=stdzr)
    ParrayPlotter(m_pa, r_upa.t).plot()
    ParrayPlotter(m_pa, r_upa.z).plot()
    ParrayPlotter(m_pa.t, r_upa.t).plot()
    ParrayPlotter(m_pa.z, r_upa.z).plot()
    ParrayPlotter(m_pa, r_upa, y_scale="standardized").plot()
    ParrayPlotter(m_pa, r_upa, y_scale="transformed").plot()
    ParrayPlotter(m_pa, r_upa, y_scale="transformed", x_scale="standardized").plot()
    ParrayPlotter(m_pa, r_upa, y_scale="standardized", x_scale="transformed").plot()
    ParrayPlotter(m_pa, r_upa, y_scale="transformed", x_scale="standardized").plot()
    ParrayPlotter(m_pa, r_upa.z, y_tick_scale="standardized").plot()
    ParrayPlotter(m_pa, r_upa.t, y_tick_scale="transformed").plot()
    ParrayPlotter(m_pa.z, r_upa, x_tick_scale="standardized").plot()
    ParrayPlotter(m_pa.t, r_upa, x_tick_scale="transformed").plot()


def test_2d_array_plotting():
    stdzr = Standardizer(c={"μ": -5.30, "σ": 0.582}, d={"μ": -0.307, "σ": 0.158}, log_vars=["d", "c"])
    c = np.arange(1, 10, 0.25)
    d = np.arange(1, 10, 0.25)
    c, d = np.meshgrid(c, d)
    Z = np.sin(np.sqrt((c - 5) ** 2 + (d - 5) ** 2))
    mrz = ParameterArray(c=c, d=d, z=Z, stdzr=stdzr)
    pp = ParrayPlotter(mrz["c"], mrz["d"], mrz["z"])
    cs = pp(plt.contourf)
    pp = ParrayPlotter(mrz["c"].t, mrz["d"].t, mrz["z"].t)
    cs = pp(plt.contourf)
    pp = ParrayPlotter(mrz["c"].z, mrz["d"].z, mrz["z"].z)
    cs = pp(plt.contourf)
    pp = ParrayPlotter(
        mrz["c"],
        mrz["d"],
        mrz["z"],
        x_scale="standardized",
        y_scale="standardized",
        z_scale="standardized",
    )
    cs = pp(plt.contourf)
    pp = ParrayPlotter(
        mrz["c"],
        mrz["d"],
        mrz["z"],
        x_scale="transformed",
        y_scale="transformed",
        z_scale="transformed",
    )
    cs = pp(plt.contourf)
    pp = ParrayPlotter(
        mrz["c"].z,
        mrz["d"].t,
        mrz["z"],
        x_tick_scale="standardized",
        y_tick_scale="transformed",
    )
    cs = pp(plt.contourf)
    pp = ParrayPlotter(
        mrz["c"].t,
        mrz["d"].z,
        mrz["z"],
        x_tick_scale="transformed",
        y_tick_scale="standardized",
    )
    cs = pp(plt.contourf)

    cs


def test_colorbar():
    c = np.arange(1, 10, 0.25)
    d = np.arange(1, 10, 0.25)
    c, d = np.meshgrid(c, d)
    Z = np.sin(np.sqrt((c - 5) ** 2 + (d - 5) ** 2))
    stdzr = Standardizer(
        c={"μ": -5, "σ": 0.5},
        d={"μ": -0.3, "σ": 0.15},
        z={"μ": 2, "σ": 2},
        log_vars=["d", "c"],
        logit_vars=["z"],
    )
    mrz = ParameterArray(c=c, d=d, z=Z**2, stdzr=stdzr)

    pp = ParrayPlotter(mrz["c"], mrz["d"], mrz["z"])
    cs = pp(plt.contourf)
    cbar = pp.colorbar(cs)

    pp = ParrayPlotter(mrz["c"], mrz["d"], mrz["z"].t)
    cs = pp(plt.contourf)
    cbar = pp.colorbar(cs)

    pp = ParrayPlotter(mrz["c"], mrz["d"], mrz["z"].z)
    cs = pp(plt.contourf)
    cbar = pp.colorbar(cs)

    pp = ParrayPlotter(mrz["c"], mrz["d"], mrz["z"].t)
    cs = pp(plt.contourf)
    cbar = pp.colorbar(cs)

    pp = ParrayPlotter(mrz["c"], mrz["d"], mrz["z"], z_scale="transformed")
    cs = pp(plt.contourf)
    cbar = pp.colorbar(cs)

    pp = ParrayPlotter(mrz["c"], mrz["d"], mrz["z"].z)
    cs = pp(plt.contourf)
    cbar = pp.colorbar(cs)

    pp = ParrayPlotter(mrz["c"], mrz["d"], mrz["z"], z_scale="standardized")
    cs = pp(plt.contourf)
    cbar = pp.colorbar(cs)

    cbar
