#!/usr/bin/env python3
"""
SAS form-factor examples.

Demonstrates the scattered intensity for a dilute layer of particles
with a given form factor.  The particle geometry and the relative
absorption are selected via command-line arguments:

    particle_geometry=<name>          (required)
    relative_absorption=<float>       (default 0; beta = value * delta)
"""

import sys
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm, std_simulations


# ---------------------------------------------------------------------------
# Per-geometry factory functions
# Each returns a list of (particle, title) pairs.
# ---------------------------------------------------------------------------

def particles_BarLorentz(material):
    ff = ba.BarLorentz(20*nm, 6.3*nm, 4.7*nm)
    result = []
    for omega in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_Bipyramid4(material):
    ff = ba.Bipyramid4(8*nm, 4.3*nm, 0.5, 70*deg)
    result = []
    for omega in [0, 30, 45]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_Box(material):
    ff = ba.Box(12*nm, 4.1*nm, 6*nm)
    result = []
    for omega in [0, 30, 45]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_CantellatedCube(material):
    ff = ba.CantellatedCube(8*nm, 2.5*nm)
    result = []
    for omega in [0, 45]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_Cone(material):
    ff = ba.Cone(4.25*nm, 12*nm, 75*deg)
    result = []
    for theta in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationY(theta*deg))
        result.append((p, r'$\theta=%d^\circ$' % theta))
    return result

def particles_CosineRippleBox(material):
    ff = ba.CosineRippleBox(20*nm, 6.3*nm, 4.7*nm)
    result = []
    for omega in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_Cylinder(material):
    ff = ba.Cylinder(3.6*nm, 7.2*nm)
    result = []
    for theta in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationY(theta*deg))
        result.append((p, r'$\theta=%d^\circ$' % theta))
    return result

def particles_Dodecahedron(material):
    ff = ba.Dodecahedron(3.38*nm)
    titles = ['face normal', 'vertex normal', 'edge normal']
    angles = [26.5651, -52.6226, 58.2825]
    result = []
    for angle, title in zip(angles, titles):
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationY(angle*deg))
        result.append((p, title))
    return result

def particles_Ellipsoid(material):
    ff = ba.Ellipsoid(4.4*nm, 8*nm, 4*nm)
    result = []
    for omega in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_EllipsoidalCylinder(material):
    ff = ba.EllipsoidalCylinder(7.8*nm, 3*nm, 4*nm)
    result = []
    for omega in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_EllipsoidalSegment(material):
    ff = ba.EllipsoidalSegment(4.4*nm, 8*nm, 4*nm, 0, 4*nm)
    result = []
    for omega in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_HorizontalCylinder(material):
    ff = ba.HorizontalCylinder(4*nm, 9*nm, 1.7*nm, 2*nm)
    result = []
    for omega in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_Icosahedron(material):
    ff = ba.Icosahedron(5.13*nm)
    titles = ['face normal', 'vertex normal', 'edge normal']
    angles = [48.1897, -52.6226, 69.0948]
    result = []
    for angle, title in zip(angles, titles):
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationY(angle*deg))
        result.append((p, title))
    return result

def particles_PlatonicOctahedron(material):
    ff = ba.PlatonicOctahedron(8.55*nm)
    titles = ['face normal', 'vertex normal', 'edge normal']
    angles = [30, 90, 0]
    result = []
    for angle, title in zip(angles, titles):
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationY(angle*deg))
        result.append((p, title))
    return result

def particles_PlatonicTetrahedron(material):
    ff = ba.PlatonicTetrahedron(13.6*nm)
    titles = ['vertex normal', 'edge normal']
    angles = [-90, 38.68218]
    result = []
    for angle, title in zip(angles, titles):
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationY(angle*deg))
        result.append((p, title))
    return result

def particles_Prism3(material):
    ff = ba.Prism3(9.2*nm, 8*nm)
    result = []
    for omega in [0, 15, 30]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_Prism6(material):
    ff = ba.Prism6(3.77*nm, 8*nm)
    result = []
    for omega in [0, 15, 30]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_Pyramid2(material):
    ff = ba.Pyramid2(6.4*nm, 13*nm, 7*nm, 70*deg)
    result = []
    for omega in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_Pyramid3(material):
    ff = ba.Pyramid3(14*nm, 6.8*nm, 70*deg)
    result = []
    for omega in [0, 15, 30]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_Pyramid4(material):
    ff = ba.Pyramid4(10*nm, 4*nm, 70*deg)
    result = []
    for omega in [0, 30]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_Pyramid6(material):
    ff = ba.Pyramid6(5.43*nm, 6.8*nm, 70*deg)
    result = []
    for omega in [0, 15, 30]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_SawtoothRippleBox(material):
    ff = ba.SawtoothRippleBox(20*nm, 6.3*nm, 4.7*nm, 3*nm)
    result = []
    for omega in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result

def particles_Sphere(material):
    ff = ba.Sphere(4.13*nm)
    return [(ba.Particle(material, ff), None)]

def particles_SphericalSegment(material):
    ff = ba.SphericalSegment(7*nm, 1*nm, 9.75*nm)
    result = []
    for theta in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationY(theta*deg))
        result.append((p, r'$\theta=%d^\circ$' % theta))
    return result

def particles_Spheroid(material):
    ff = ba.Spheroid(6.85*nm, 1.5*nm)
    result = []
    for theta in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationY(theta*deg))
        result.append((p, r'$\theta=%d^\circ$' % theta))
    return result

def particles_SpheroidalSegment(material):
    ff = ba.SpheroidalSegment(7*nm, 13.3*nm, 2.3*nm, 20.6*nm)
    result = []
    for theta in [0, 45, 90]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationY(theta*deg))
        result.append((p, r'$\theta=%d^\circ$' % theta))
    return result

def particles_TruncatedCube(material):
    ff = ba.TruncatedCube(6.74*nm, 2*nm)
    result = []
    for omega in [0, 15, 45]:
        p = ba.Particle(material, ff)
        p.rotate(ba.RotationZ(omega*deg))
        result.append((p, r'$\omega=%d^\circ$' % omega))
    return result


# ---------------------------------------------------------------------------
# Simulation
# ---------------------------------------------------------------------------

def get_simulation(particle):
    n = 201
    return std_simulations.particle(particle, n)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

if __name__ == '__main__':
    # Extract script-specific args and clean sys.argv so that
    # ba_check.persistence_test only sees datfile/reference/tolerance.
    particle_geometry = None
    relative_absorption = 0.0
    clean_argv = [sys.argv[0]]
    for arg in sys.argv[1:]:
        s = arg.split('=', 1)
        if s[0] == 'particle_geometry':
            particle_geometry = s[1]
        elif s[0] == 'relative_absorption':
            relative_absorption = float(s[1])
        else:
            clean_argv.append(arg)
    sys.argv = clean_argv

    if particle_geometry is None:
        raise ValueError("Required argument 'particle_geometry' is missing")

    delta = 1e-6
    particle_color = (0.86, 0.24, 0.18)
    material = ba.RefractiveMaterial(
        "Particle", particle_color, delta, relative_absorption*delta)

    factory = globals()[f'particles_{particle_geometry}']
    pairs = factory(material)

    results = []
    for particle, title in pairs:
        simulation = get_simulation(particle)
        result = simulation.simulate()
        if title is not None:
            result.setTitle(title)
        results.append(result)

    plotargs = bp.parse_commandline()
    bp.plot_ff_to_row(results, **plotargs)
    bp.export(**plotargs)
