from __future__ import annotations

import numpy as np
import unittest

from rfg_adc_plotter.processing.calibration import calibrate_freqs, recalculate_calibration_c
from rfg_adc_plotter.processing.fft import compute_distance_axis, compute_fft_mag_row, compute_fft_row
from rfg_adc_plotter.processing.normalization import build_calib_envelopes, normalize_by_calib
from rfg_adc_plotter.processing.peaks import find_peak_width_markers, find_top_peaks_over_ref, rolling_median_ref


class ProcessingTests(unittest.TestCase):
    def test_recalculate_calibration_preserves_requested_edges(self):
        coeffs = recalculate_calibration_c(np.asarray([0.0, 1.0, 0.025], dtype=np.float64), 3.3, 14.3)
        y0 = coeffs[0] + coeffs[1] * 3.3 + coeffs[2] * (3.3 ** 2)
        y1 = coeffs[0] + coeffs[1] * 14.3 + coeffs[2] * (14.3 ** 2)
        self.assertTrue(np.isclose(y0, 3.3))
        self.assertTrue(np.isclose(y1, 14.3))

    def test_calibrate_freqs_returns_monotonic_axis_and_same_shape(self):
        sweep = {"F": np.linspace(3.3, 14.3, 32), "I": np.linspace(-1.0, 1.0, 32)}
        calibrated = calibrate_freqs(sweep)
        self.assertEqual(calibrated["F"].shape, (32,))
        self.assertEqual(calibrated["I"].shape, (32,))
        self.assertTrue(np.all(np.diff(calibrated["F"]) >= 0.0))

    def test_normalizers_and_envelopes_return_finite_ranges(self):
        calib = (np.sin(np.linspace(0.0, 4.0 * np.pi, 64)) * 5.0).astype(np.float32)
        raw = calib * 0.75
        lower, upper = build_calib_envelopes(calib)
        self.assertEqual(lower.shape, calib.shape)
        self.assertEqual(upper.shape, calib.shape)
        self.assertTrue(np.all(lower <= upper))

        simple = normalize_by_calib(raw, calib + 10.0, norm_type="simple")
        projector = normalize_by_calib(raw, calib, norm_type="projector")
        self.assertEqual(simple.shape, raw.shape)
        self.assertEqual(projector.shape, raw.shape)
        self.assertTrue(np.any(np.isfinite(simple)))
        self.assertTrue(np.any(np.isfinite(projector)))

    def test_fft_helpers_return_expected_shapes(self):
        sweep = np.sin(np.linspace(0.0, 4.0 * np.pi, 128)).astype(np.float32)
        freqs = np.linspace(3.3, 14.3, 128, dtype=np.float64)
        mag = compute_fft_mag_row(sweep, freqs, 513)
        row = compute_fft_row(sweep, freqs, 513)
        axis = compute_distance_axis(freqs, 513)
        self.assertEqual(mag.shape, (513,))
        self.assertEqual(row.shape, (513,))
        self.assertEqual(axis.shape, (513,))
        self.assertTrue(np.all(np.diff(axis) >= 0.0))

    def test_peak_helpers_find_reference_and_peak_boxes(self):
        xs = np.linspace(0.0, 10.0, 200)
        ys = np.exp(-((xs - 5.0) ** 2) / 0.4) * 10.0 + 1.0
        ref = rolling_median_ref(xs, ys, 2.0)
        peaks = find_top_peaks_over_ref(xs, ys, ref, top_n=3)
        width = find_peak_width_markers(xs, ys)
        self.assertEqual(ref.shape, ys.shape)
        self.assertEqual(len(peaks), 1)
        self.assertGreater(peaks[0]["x"], 4.0)
        self.assertLess(peaks[0]["x"], 6.0)
        self.assertIsNotNone(width)
        self.assertGreater(width["width"], 0.0)


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