#!/usr/bin/env python3
"""
GISAS with rectangular and spherical detector ROI and masks.

Same CylindersAndPrisms sample with two detector setups:
1. Spherical detector (40x30), rectangular mask, ROI #1 (RectDetWithRoi)
2. Spherical detector (40x30), rectangular mask, ROI #2 (SphericalDetWithRoi)
Both use the standard beam (intensity=1, wavelength=0.1 nm, alpha=0.2 deg).
"""
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm


def get_sample():
    vacuum = ba.Vacuum()
    particle_color = (0.86, 0.24, 0.18)
    particle_mat = ba.RefractiveMaterial("Particle", particle_color, 6e-4, 2e-8)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.RefractiveMaterial("Substrate", substrate_color, 6e-6, 2e-8)

    cylinder = ba.Particle(particle_mat, ba.Cylinder(5*nm, 5*nm))
    prism3 = ba.Particle(particle_mat, ba.Prism3(10*nm, 5*nm))

    vacuum_layer = ba.Layer(vacuum)
    vacuum_layer.deposit2D(ba.Dilute2D(0.0004, cylinder))
    vacuum_layer.deposit2D(ba.Dilute2D(0.0004, prism3))

    sample = ba.Sample()
    sample.addLayer(vacuum_layer)
    sample.addLayer(ba.Layer(substrate_mat))
    return sample


def simulate_rect_roi(sample):
    beam = ba.Beam(1, 0.1*nm, 0.2*deg)
    detector = ba.SphericalDetector(40, -2*deg, 2*deg, 30, -1*deg, 2*deg)
    sim = ba.ScatteringSimulation(beam, sample, detector)
    sim.detector().addMask(
        ba.Rectangle(0.35*deg, -1.15*deg, 0.75*deg, 1.95*deg))
    sim.detector().setRegionOfInterest(
        -2.55*deg, 0.25*deg, 1.75*deg, 1.85*deg)
    result = sim.simulate()
    result.setTitle("RectDetWithRoi")
    return result


def simulate_spherical_roi(sample):
    beam = ba.Beam(1, 0.1*nm, 0.2*deg)
    detector = ba.SphericalDetector(40, -2*deg, 2*deg, 30, -1*deg, 2*deg)
    sim = ba.ScatteringSimulation(beam, sample, detector)
    sim.detector().addMask(
        ba.Rectangle(-0.499*deg, 0.601*deg, -0.201*deg, 0.999*deg))
    sim.detector().setRegionOfInterest(
        -1.499*deg, 0.251*deg, 1.499*deg, 1.749*deg)
    result = sim.simulate()
    result.setTitle("SphericalDetWithRoi")
    return result


if __name__ == '__main__':
    sample = get_sample()
    results = [
        simulate_rect_roi(sample),
        simulate_spherical_roi(sample),
    ]
    plotargs = bp.parse_commandline()
    bp.plot2d_to_row(results, **plotargs)
    bp.export(**plotargs)
