import unittest

import matplotlib.pyplot as plt
import numpy as np

import wfdb
from wfdb.plot import plot


class TestPlotWfdb(unittest.TestCase):
    """
    Tests for the wfdb.plot_wfdb function
    """

    def assertAxesMatchSignal(self, axes, signal, t_divisor=1):
        """
        Check that axis limits are reasonable for plotting a signal array.

        Parameters
        ----------
        axes : matplotlib.axes.Axes
            An Axes object.
        signal : numpy.ndarray
            A one-dimensional array of sample values.
        t_divisor : float, optional
            The intended plotting resolution (number of samples of `signal`
            per unit of the X axis.)

        """
        xmin, xmax = axes.get_xlim()
        tmin = 0
        tmax = (len(signal) - 1) / t_divisor
        # The range from tmin to tmax should fit within the plot.
        self.assertLessEqual(
            xmin,
            tmin,
            msg=f"X range is [{xmin}, {xmax}]; expected [{tmin}, {tmax}]",
        )
        self.assertGreaterEqual(
            xmax,
            tmax,
            msg=f"X range is [{xmin}, {xmax}]; expected [{tmin}, {tmax}]",
        )
        # The padding on left and right sides should be approximately equal.
        self.assertAlmostEqual(
            xmin - tmin,
            tmax - xmax,
            delta=(tmax - tmin) / 10 + 1 / t_divisor,
            msg=f"X range is [{xmin}, {xmax}]; expected [{tmin}, {tmax}]",
        )

        ymin, ymax = axes.get_ylim()
        vmin = np.nanmin(signal)
        vmax = np.nanmax(signal)
        # The range from vmin to vmax should fit within the plot.
        self.assertLessEqual(
            ymin,
            vmin,
            msg=f"Y range is [{ymin}, {ymax}]; expected [{vmin}, {vmax}]",
        )
        self.assertGreaterEqual(
            ymax,
            vmax,
            msg=f"Y range is [{ymin}, {ymax}]; expected [{vmin}, {vmax}]",
        )
        # The padding on top and bottom should be approximately equal.
        self.assertAlmostEqual(
            ymin - vmin,
            vmax - ymax,
            delta=(vmax - vmin) / 10,
            msg=f"Y range is [{ymin}, {ymax}]; expected [{vmin}, {vmax}]",
        )

    def test_physical_smooth(self):
        """
        Plot a record with physical, single-frequency data
        """
        record = wfdb.rdrecord(
            "sample-data/100",
            sampto=1000,
            physical=True,
            smooth_frames=True,
        )
        self.assertIsNotNone(record.p_signal)

        annotation = wfdb.rdann("sample-data/100", "atr", sampto=1000)

        fig = wfdb.plot_wfdb(
            record,
            annotation,
            time_units="samples",
            ecg_grids="all",
            return_fig=True,
        )
        plt.close(fig)

        self.assertEqual(len(fig.axes), record.n_sig)
        for ch in range(record.n_sig):
            self.assertAxesMatchSignal(fig.axes[ch], record.p_signal[:, ch])

    def test_digital_smooth(self):
        """
        Plot a record with digital, single-frequency data
        """
        record = wfdb.rdrecord(
            "sample-data/drive02",
            sampto=1000,
            physical=False,
            smooth_frames=True,
        )
        self.assertIsNotNone(record.d_signal)

        fig = wfdb.plot_wfdb(record, time_units="seconds", return_fig=True)
        plt.close(fig)

        self.assertEqual(len(fig.axes), record.n_sig)
        for ch in range(record.n_sig):
            self.assertAxesMatchSignal(
                fig.axes[ch], record.d_signal[:, ch], record.fs
            )

    def test_physical_multifrequency(self):
        """
        Plot a record with physical, multi-frequency data
        """
        record = wfdb.rdrecord(
            "sample-data/wave_4",
            sampto=10,
            physical=True,
            smooth_frames=False,
        )
        self.assertIsNotNone(record.e_p_signal)

        fig = wfdb.plot_wfdb(record, time_units="seconds", return_fig=True)
        plt.close(fig)

        self.assertEqual(len(fig.axes), record.n_sig)
        for ch in range(record.n_sig):
            self.assertAxesMatchSignal(
                fig.axes[ch],
                record.e_p_signal[ch],
                record.fs * record.samps_per_frame[ch],
            )

    def test_digital_multifrequency(self):
        """
        Plot a record with digital, multi-frequency data
        """
        record = wfdb.rdrecord(
            "sample-data/multi-segment/041s/041s",
            sampto=1000,
            physical=False,
            smooth_frames=False,
        )
        self.assertIsNotNone(record.e_d_signal)

        fig = wfdb.plot_wfdb(record, time_units="seconds", return_fig=True)
        plt.close(fig)

        self.assertEqual(len(fig.axes), record.n_sig)
        for ch in range(record.n_sig):
            self.assertAxesMatchSignal(
                fig.axes[ch],
                record.e_d_signal[ch],
                record.fs * record.samps_per_frame[ch],
            )


class TestPlotInternal(unittest.TestCase):
    """
    Unit tests for internal wfdb.plot.plot functions
    """

    def test_get_plot_dims(self):
        sampfrom = 0
        sampto = 3000
        record = wfdb.rdrecord(
            "sample-data/100", physical=True, sampfrom=sampfrom, sampto=sampto
        )
        ann = wfdb.rdann(
            "sample-data/100", "atr", sampfrom=sampfrom, sampto=sampto
        )
        sig_len, n_sig, n_annot, n_subplots = plot._get_plot_dims(
            signal=record.p_signal, ann_samp=[ann.sample]
        )

        assert sig_len == sampto - sampfrom
        assert n_sig == record.n_sig
        assert n_annot == 1
        assert n_subplots == record.n_sig

    def test_create_figure_single_subplots(self):
        n_subplots = 1
        fig, axes = plot._create_figure(
            n_subplots, sharex=True, sharey=True, figsize=None
        )
        assert fig is not None
        assert axes is not None
        assert len(axes) == n_subplots

    def test_create_figure_multiple_subplots(self):
        n_subplots = 5
        fig, axes = plot._create_figure(
            n_subplots, sharex=True, sharey=True, figsize=None
        )
        assert fig is not None
        assert axes is not None
        assert len(axes) == n_subplots


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