//! CMB polarization power spectra
//!
//! This example demonstrates:
//! - E-mode power spectrum (scalar perturbations)
//! - B-mode power spectrum (tensor perturbations)
//! - Temperature-E cross-correlation (TE)
//! - Tensor-to-scalar ratio constraint

use andam::cmb::polarization::*;
use andam::dynamics::Universe;
use plotters::prelude::*;
use std::f64::consts::PI;
use std::error::Error;

fn main() -> Result<(), Box<dyn Error>> {
    println!("=== CMB Polarization Power Spectra ===\n");

    let universe = Universe::benchmark();
    let l_max = 2000;

    println!("Computing polarization spectra...");
    println!("  l_max = {}", l_max);

    let pol_spectrum = PolarizationSpectrum::from_boltzmann(&universe, l_max);

    let r = pol_spectrum.tensor_to_scalar_ratio();
    println!("  Tensor-to-scalar ratio r = {:.6}\n", r);

    let filename = "cmb_polarization.png";
    let root = BitMapBackend::new(filename, (2400, 1800)).into_drawing_area();
    root.fill(&WHITE)?;

    // Split into 3 panels vertically
    let areas = root.split_evenly((3, 1));

    // Panel 1: EE spectrum
    {
        let mut chart = ChartBuilder::on(&areas[0])
            .caption("CMB E-mode Power Spectrum", ("sans-serif", 40))
            .margin(15)
            .x_label_area_size(60)
            .y_label_area_size(80)
            .build_cartesian_2d(2..l_max, 0.1..100.0)?;

        chart.configure_mesh()
            .x_desc("Multipole l")
            .y_desc("l(l+1)C_l^EE/2π [μK²]")
            .draw()?;

        let data_ee: Vec<_> = (2..=l_max)
            .map(|l| {
                let cl = pol_spectrum.c_l_ee[l];
                let dl = (l * (l + 1)) as f64 * cl / (2.0 * PI);
                (l, dl)
            })
            .collect();

        chart.draw_series(LineSeries::new(
            data_ee,
            BLUE.stroke_width(3),
        ))?
        .label("E-mode (scalar perturbations)")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], BLUE.stroke_width(3)));

        chart.configure_series_labels()
            .background_style(&WHITE.mix(0.8))
            .border_style(&BLACK)
            .draw()?;

        // Add annotation
        chart.draw_series(std::iter::once(Text::new(
            "Generated by Thomson scattering",
            (100, 80.0),
            ("sans-serif", 28).into_font(),
        )))?;
    }

    // Panel 2: BB spectrum
    {
        let mut chart = ChartBuilder::on(&areas[1])
            .caption("CMB B-mode Power Spectrum", ("sans-serif", 40))
            .margin(15)
            .x_label_area_size(60)
            .y_label_area_size(80)
            .build_cartesian_2d(
                2..l_max,
                (0.0001..0.1).log_scale()
            )?;

        chart.configure_mesh()
            .x_desc("Multipole l")
            .y_desc("l(l+1)C_l^BB/2π [μK²]")
            .draw()?;

        let data_bb: Vec<_> = (2..=l_max)
            .filter_map(|l| {
                let cl = pol_spectrum.c_l_bb[l];
                if cl > 0.0 {
                    let dl = (l * (l + 1)) as f64 * cl / (2.0 * PI);
                    Some((l, dl))
                } else {
                    None
                }
            })
            .collect();

        chart.draw_series(LineSeries::new(
            data_bb,
            RED.stroke_width(3),
        ))?
        .label("B-mode (tensor perturbations)")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RED.stroke_width(3)));

        chart.configure_series_labels()
            .background_style(&WHITE.mix(0.8))
            .border_style(&BLACK)
            .draw()?;

        // Add annotations
        chart.draw_series(std::iter::once(Text::new(
            "Primordial gravitational waves",
            (10, 0.01),
            ("sans-serif", 28).into_font(),
        )))?;

        chart.draw_series(std::iter::once(Text::new(
            format!("r = {:.6} (tensor-to-scalar ratio)", r),
            (10, 0.005),
            ("sans-serif", 28).into_font(),
        )))?;
    }

    // Panel 3: TE cross-correlation
    {
        let mut chart = ChartBuilder::on(&areas[2])
            .caption("Temperature-E Cross-Correlation", ("sans-serif", 40))
            .margin(15)
            .x_label_area_size(60)
            .y_label_area_size(80)
            .build_cartesian_2d(2..l_max, -15.0..10.0)?;

        chart.configure_mesh()
            .x_desc("Multipole l")
            .y_desc("l(l+1)C_l^TE/2π [μK²]")
            .draw()?;

        let data_te: Vec<_> = (2..=l_max)
            .map(|l| {
                let cl = pol_spectrum.c_l_te[l];
                let dl = (l * (l + 1)) as f64 * cl / (2.0 * PI);
                (l, dl)
            })
            .collect();

        chart.draw_series(LineSeries::new(
            data_te,
            GREEN.stroke_width(3),
        ))?
        .label("TE cross-correlation")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], GREEN.stroke_width(3)));

        // Zero line
        chart.draw_series(LineSeries::new(
            vec![(2, 0.0), (l_max, 0.0)],
            BLACK.stroke_width(1),
        ))?;

        chart.configure_series_labels()
            .background_style(&WHITE.mix(0.8))
            .border_style(&BLACK)
            .draw()?;

        // Add annotation
        chart.draw_series(std::iter::once(Text::new(
            "Correlation between T and E modes",
            (100, 8.0),
            ("sans-serif", 28).into_font(),
        )))?;
    }

    root.present()?;
    println!("Created: {}", filename);

    // Print some statistics
    println!("\n=== Polarization Spectrum Statistics ===");
    println!("E-mode power at l=200: {:.2} μK²",
             200.0 * 201.0 * pol_spectrum.c_l_ee[200] / (2.0 * PI));
    println!("B-mode power at l=80:  {:.6} μK²",
             80.0 * 81.0 * pol_spectrum.c_l_bb[80] / (2.0 * PI));
    println!("TE power at l=200:     {:.2} μK²",
             200.0 * 201.0 * pol_spectrum.c_l_te[200] / (2.0 * PI));
    println!("\nNote: B-modes are tiny because r << 1 (no primordial GW detection yet)");

    Ok(())
}
