#!/usr/bin/env python3
"""
Cylinders of two different sizes in Decoupling Approximation,
Local Monodisperse Approximation and Size-Spacing Coupling Approximation
"""
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm


def get_sample(approximation):
    """
    A sample with cylinders of two different sizes on a substrate
    in radial paracrystal ordering.
    """

    # Materials
    particle_color = (0.86, 0.24, 0.18)
    particle_mat = ba.RefractiveMaterial("Particle", particle_color, 0.0006, 2e-08)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.RefractiveMaterial("Substrate", substrate_color, 6e-06, 2e-08)
    vacuum = ba.Vacuum()

    # Particles
    ff_1 = ba.Cylinder(5*nm, 5*nm)
    ff_2 = ba.Cylinder(8*nm, 8*nm)
    particle_1 = ba.Particle(particle_mat, ff_1)
    particle_2 = ba.Particle(particle_mat, ff_2)

    # Layers
    layer_1 = ba.Layer(vacuum)
    layer_2 = ba.Layer(substrate_mat)

    # Distribution function
    profile = ba.Profile1DGauss(3*nm)

    if (not approximation=='LMA' and
        not approximation=='DA' and
        not approximation=='SSCA'):
        raise Exception('Unknown approximation')

    if approximation == 'LMA':
        """
        Local Monodisperse Approximation: two independent layouts
        with their own distinct interference functions
        """
        # Interference functions
        layout_1 = ba.RadialParacrystal(particle_1, 16.8*nm, 1000*nm)
        layout_2 = ba.RadialParacrystal(particle_2, 22.8*nm, 1000*nm)
        layout_1.setProbabilityDistribution(profile)
        layout_2.setProbabilityDistribution(profile)

        # Populate layer with two independent layouts
        layer_1.addDeposit2D(0.5, layout_1)
        layer_1.addDeposit2D(0.5, layout_2)

    else: # 'DA/SSCA'
        # Particle mixture
        mix = ba.Mixture()
        mix.addParticle(particle_1, 4)
        mix.addParticle(particle_2, 1)

        layout = ba.RadialParacrystal(mix, 18*nm, 1000*nm)
        profile = ba.Profile1DGauss(3*nm)
        layout.setProbabilityDistribution(profile)

        """
        The only difference between Decoupling Approximation and
        Size-Spacing Coupling Approximation is that in SSCA there is
        a position phase offset between fractions defined by 'kappa' parameter:
        DA: kappa = 0 (default value)
        SSCA: kappa = 1
        """
        if approximation == 'SSCA':
            layout.setKappa(1)

        # Populate layer with the layout
        layer_1.deposit2D(layout)

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

    return sample


def get_simulation(sample):
    beam = ba.Beam(1e9, 0.1*nm, 0.2*deg)
    n = 200
    detector = ba.SphericalDetector(n, 0., 2*deg, n, 0., 2*deg)
    simulation = ba.ScatteringSimulation(beam, sample, detector)
    simulation.options().setUseAvgMaterials(False)
    return simulation


def simulate(approximation):
    sample = get_sample(approximation)
    simulation = get_simulation(sample)
    result = simulation.simulate()
    result.setTitle(approximation)
    return result


if __name__ == '__main__':
    results = [
        simulate('LMA'),
        simulate('DA'),
        simulate('SSCA')
    ]
    plotargs = bp.parse_commandline()
    bp.plot2d_to_row(results, unit_aspect=1, **plotargs)
    bp.export(**plotargs)
