#!/usr/bin/env python3
"""
Simple test of basic CMB alignment computation.

Verifies that the AMD alignment method produces valid (non-NaN) angles
in the [0, 90] degree range for random realisations.
"""

import numpy as np
from src.cmb_generator import CMBGenerator, CMBParameters


def test_basic_alignment():
    """Test basic alignment calculation with multiple seeds."""
    print("Testing basic alignment calculation (AMD method)...")

    cmb_params = CMBParameters(
        nside=64,
        lmax=10,
        seed=42,
        quadrupole_power=1.0,
        octopole_power=1.0,
    )

    all_ok = True

    for i in range(10):
        generator = CMBGenerator(CMBParameters(
            nside=64, lmax=10, seed=i,
            quadrupole_power=1.0, octopole_power=1.0,
        ))

        modes = generator.generate_low_l_modes()
        alignment = generator.compute_multipole_alignment(modes['alm'])
        angle = alignment['alignment_angle']

        valid = (not np.isnan(angle)) and (0 <= angle <= 90)
        status = "OK" if valid else "FAIL"

        print(f"  Seed {i}: alignment = {angle:.2f} deg [{status}]")

        if not valid:
            print(f"    quadrupole_vector: {alignment['quadrupole_vector']}")
            print(f"    octopole_vector:   {alignment['octopole_vector']}")
            print(f"    abs_dot_product:   {alignment['abs_dot_product']}")
            all_ok = False

    if all_ok:
        print("All alignment calculations valid.")
    else:
        print("Some calculations produced invalid results!")

    return all_ok


def test_basic_correction():
    """Test that a small perturbation produces valid results."""
    print("\nTesting basic correction...")

    rng = np.random.default_rng(42)

    generator = CMBGenerator(CMBParameters(
        nside=64, lmax=10, seed=42,
    ))

    modes = generator.generate_low_l_modes()
    initial_alm = modes['alm']

    for i in range(5):
        corrected_alm = initial_alm + 0.001 * rng.standard_normal(
            initial_alm.shape
        )

        initial_align = generator.compute_multipole_alignment(initial_alm)
        corrected_align = generator.compute_multipole_alignment(corrected_alm)

        i_angle = initial_align['alignment_angle']
        c_angle = corrected_align['alignment_angle']

        valid = not (np.isnan(i_angle) or np.isnan(c_angle))
        status = "OK" if valid else "FAIL"

        print(f"  Perturbation {i+1}: "
              f"initial={i_angle:.2f} deg, "
              f"corrected={c_angle:.2f} deg [{status}]")

    print("Basic correction test complete.")
    return True


if __name__ == "__main__":
    test_basic_alignment()
    test_basic_correction()