#!/usr/bin/env python3
"""
Cylindrical mesocrystal with hexagonal lattice and two-material basis.
Test to investigate intensity changes with different numbers of slices.
"""
import math
import numpy as np
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm, nm2, R3
import matplotlib.pyplot as plt


def get_sample(n_slices=1):
    # Materials (zero absorption)
    particle1_color = (0.86, 0.24, 0.18)
    particle1_mat = ba.RefractiveMaterial("Particle1", particle1_color, 4e-05, 0)
    particle2_color = (0.25, 0.65, 0.35)
    particle2_mat = ba.RefractiveMaterial("Particle2", particle2_color, 8e-06, 0)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.RefractiveMaterial("Substrate", substrate_color, 6e-06, 0)
    vacuum = ba.Vacuum()

    # Basis particles
    inner_ff = ba.FuzzySphere(4.7*nm, .4*nm)

    inner_particle1 = ba.Particle(particle1_mat, inner_ff)
    inner_particle2 = ba.Particle(particle2_mat, inner_ff)

    # 3D lattice
    lattice_a = 6.21
    lattice_c = 6.57
    sigma_a = 1.16
    position_variance = sigma_a/math.sqrt(3)

    lattice = ba.HexagonalLattice(lattice_a*2, lattice_c*2*2.3)

    bas_a = lattice.basisVectorA()
    bas_b = lattice.basisVectorB()
    bas_c = lattice.basisVectorC()

    # Basis composition
    position_0 = R3(0, 0, 0)
    position_1 = 1.0/3*(2*bas_a + bas_b + bas_c)
    basis = ba.Compound()
    basis.addComponent(inner_particle1, position_0)
    basis.addComponent(inner_particle2, position_1)

    # Crystal
    crystal = ba.Crystal(basis, lattice, position_variance)

    # Mesocrystal: shaped crystal
    outer_ff = ba.Cylinder(20*nm, 50*nm)
    outer_particle = ba.Mesocrystal(crystal, outer_ff)

    # Layers
    layer_1 = ba.Layer(vacuum)
    layer_1.deposit2D(ba.Dilute2D(0.0001/nm2, outer_particle))
    if n_slices > 1:
        layer_1.setNumberOfSlices(n_slices)
    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, reciprocal=False):
    """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)
    if reciprocal:
        simulation.options().setMesoReciprocalSum(True, 2.5)
    return simulation.simulate()


def simulate_vertical_cut(n_slices=1, reciprocal=False):
    """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)
    if reciprocal:
        simulation.options().setMesoReciprocalSum(True, 2.5)
    return simulation.simulate()


def simulate_depthprobe(n_slices=1, reciprocal=False):
    """Simulate intensity vs depth."""
    sample = get_sample(n_slices)
    n = 100
    alpha_f = 0.001*deg if reciprocal else 0.0064*deg
    scan = ba.AlphaScan([alpha_f])
    scan.setWavelength(0.1*nm)
    z_axis = ba.EquiDivision("z (nm)", n, -20*nm, 120*nm)
    simulation = ba.DepthprobeSimulation(scan, sample, z_axis)
    if reciprocal:
        simulation.options().setMesoReciprocalSum(True, 2.5)
    return simulation.simulate(), alpha_f/deg


def get_sld_profile(n_slices=1):
    """Get SLD profile from resampled sample."""
    sample = get_sample(n_slices)
    n = 200
    z_min, z_max = -20*nm, 60*nm
    z = np.linspace(z_min, z_max, n)
    sld = ba.materialProfileSLD(sample, n, z_min, z_max)
    sld_real = np.array([s.real for s in sld])
    return z/nm, sld_real


# 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 2 rows, 4 columns
    fig, axes = plt.subplots(2, 4, figsize=(18, 10))

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

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

    ax = axes[0, 2]
    for i, n in enumerate(slice_counts):
        result, alpha_f = simulate_depthprobe(n_slices=n, reciprocal=False)
        ax.semilogy(result.yCenters(), result.flatVector(),
                    label=f'N={n}', color=colors[i])
    ax.set_title(f'Real-space: Intensity vs depth at alpha_f={alpha_f:.4f} deg')
    ax.set_xlabel('z (nm)')
    ax.set_ylabel('Intensity')
    ax.legend()

    ax = axes[0, 3]
    for i, n in enumerate(slice_counts):
        z, sld = get_sld_profile(n_slices=n)
        ax.plot(z, sld, label=f'N={n}', color=colors[i])
    ax.set_title('Real-space: SLD profile')
    ax.set_xlabel('z (nm)')
    ax.set_ylabel('SLD (real part)')
    ax.legend()

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

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

    ax = axes[1, 2]
    for i, n in enumerate(slice_counts):
        result, alpha_f = simulate_depthprobe(n_slices=n, reciprocal=True)
        ax.semilogy(result.yCenters(), result.flatVector(),
                    label=f'N={n}', color=colors[i])
    ax.set_title(f'Fourier: Intensity vs depth at alpha_f={alpha_f:.3f} deg')
    ax.set_xlabel('z (nm)')
    ax.set_ylabel('Intensity')
    ax.legend()

    ax = axes[1, 3]
    for i, n in enumerate(slice_counts):
        z, sld = get_sld_profile(n_slices=n)
        ax.plot(z, sld, label=f'N={n}', color=colors[i])
    ax.set_title('Fourier: SLD profile')
    ax.set_xlabel('z (nm)')
    ax.set_ylabel('SLD (real part)')
    ax.legend()

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