#!/usr/bin/env python3
"""
Dilute cylinders on substrate.
Test to investigate intensity changes with different numbers of slices.
"""
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm
import matplotlib.pyplot as plt


def get_sample(n_slices=1):
    # Materials (with absorption for sharper tests)
    particle_color = (0.86, 0.24, 0.18)
    particle_mat = ba.RefractiveMaterial("Particle", particle_color, 6e-04, 2e-05)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.RefractiveMaterial("Substrate", substrate_color, 6e-06, 2e-08)
    vacuum = ba.Vacuum()

    # Particle
    ff = ba.Cylinder(5*nm, 5*nm)
    particle = ba.Particle(particle_mat, ff)

    # Layers
    layer_1 = ba.Layer(vacuum)
    if n_slices > 1:
        layer_1.setNumberOfSlices(n_slices)
    layer_1.deposit2D(ba.Dilute2D(0.001, particle))
    layer_2 = ba.Layer(substrate_mat)

    # Sample
    sample = ba.Sample()
    sample.addLayer(layer_1)
    sample.addLayer(layer_2)

    return sample


def simulate_horizontal_cut(n_slices=1):
    """Simulate with Nx1 detector for horizontal cut at alpha_f=0.92 deg."""
    sample = get_sample(n_slices)
    beam = ba.Beam(1e9, 0.1*nm, 0.2*deg)
    n = 200
    # Single bin centered at alpha_f=0.92 deg
    detector = ba.SphericalDetector(n, -2*deg, 2*deg, 1, 0.91*deg, 0.93*deg)
    simulation = ba.ScatteringSimulation(beam, sample, detector)
    return simulation.simulate()


def simulate_vertical_cut(n_slices=1):
    """Simulate with 1xN detector for vertical cut at phi_f=0 deg."""
    sample = get_sample(n_slices)
    beam = ba.Beam(1e9, 0.1*nm, 0.2*deg)
    n = 200
    # Single bin centered at phi_f=0 deg
    detector = ba.SphericalDetector(1, -0.01*deg, 0.01*deg, n, 0, 2*deg)
    simulation = ba.ScatteringSimulation(beam, sample, detector)
    return simulation.simulate()


# Color cycle for different N values
colors = ['C0', 'C1', 'C2', 'C3']


if __name__ == '__main__':
    # Slice counts to compare
    slice_counts = [1, 2, 5, 20]

    # Create figure with 1 row, 2 columns
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    # Horizontal cuts at alpha_f=0.92 deg
    ax = axes[0]
    for i, n in enumerate(slice_counts):
        result = simulate_horizontal_cut(n_slices=n)
        ax.semilogy(result.xCenters(), result.flatVector(),
                    label=f'N={n}', color=colors[i])
    ax.set_title('Horizontal cut at alpha_f=0.92 deg')
    ax.set_xlabel('phi_f (deg)')
    ax.set_ylabel('Intensity')
    ax.legend()

    # Vertical cuts at phi_f=0 deg
    ax = axes[1]
    for i, n in enumerate(slice_counts):
        result = simulate_vertical_cut(n_slices=n)
        ax.semilogy(result.yCenters(), result.flatVector(),
                    label=f'N={n}', color=colors[i])
    ax.set_title('Vertical cut at phi_f=0 deg')
    ax.set_xlabel('alpha_f (deg)')
    ax.set_ylabel('Intensity')
    ax.legend()

    plt.tight_layout()
    ba.showSample3D(get_sample(), sample_size=120*nm, seed=0)
    plt.show()
