#!/usr/bin/env python3
"""
Fitting experimental data: spherical nanoparticles with size distribution
in 3 layers system (experiment at GALAXI).
"""
import os
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm, ba_fitmonitor, R3
from matplotlib import pyplot as plt
import lmfit

wavelength = 1.34*ba.angstrom
alpha_i = 0.463*ba.deg

datadir = os.getenv('BA_DATA_DIR', '')
if not datadir:
    raise Exception("Environment variable BA_DATA_DIR not set")

# sample model

radius = 5.75*ba.nm
sigma = 0.4
distance = 53.6*ba.nm
disorder = 10.5*ba.nm
kappa = 17.5
ptfe_thickness = 22.1*ba.nm
hmdso_thickness = 18.5*ba.nm

def get_sample(P):
    radius = P["radius"]
    sigma = P["sigma"]
    distance = P["distance"]

    # defining materials
    vacuum = ba.Vacuum()
    si_color = (0.30, 0.62, 0.86)
    si_mat = ba.RefractiveMaterial("Si", si_color, 5.7816e-6, 1.0229e-7)
    ag_color = (0.86, 0.24, 0.18)
    ag_mat = ba.RefractiveMaterial("Ag", ag_color, 2.2475e-5, 1.6152e-6)
    ptfe_color = (0.93, 0.72, 0.25)
    ptfe_mat = ba.RefractiveMaterial("PTFE", ptfe_color, 5.20509e-6, 1.9694e-8)
    hmdso_color = (0.25, 0.65, 0.35)
    hmdso_mat = ba.RefractiveMaterial("HMDSO", hmdso_color, 2.0888e-6, 1.3261e-8)

    # collection of particles with size distribution
    nparticles = 20
    nfwhm = 2.0
    sphere_ff = ba.Sphere(radius)

    sphere = ba.Particle(ag_mat, sphere_ff)
    position = R3(0, 0, -1*hmdso_thickness)
    sphere.translate(position)
#     ln_distr = ba.DistributionLogNormal(radius, sigma)
#     par_distr = ba.ParameterDistribution(
#         "/Particle/Sphere/Radius", ln_distr, nparticles, nfwhm,
#     part_coll = ba.ParticleDistribution(sphere, par_distr)

    # interference function
    layout = ba.RadialParacrystal(sphere, distance, 1e6*ba.nm)
    layout.setKappa(kappa)
    layout.setDomainSize(2e4*nm)
    profile = ba.Profile1DGauss(disorder)
    layout.setProbabilityDistribution(profile)

    # roughness
    autocorr_1 = ba.SelfAffineFractalModel(2.3*nm, 0.3, 5*nm)
    autocorr_2 = ba.SelfAffineFractalModel(1.1*nm, 0.3, 5*nm)
    transient = ba.TanhTransient()
    roughness_1 = ba.Roughness(autocorr_1, transient)
    roughness_2 = ba.Roughness(autocorr_2, transient)

    # layers
    vacuum_layer = ba.Layer(vacuum)
    hmdso_layer = ba.Layer(hmdso_mat, hmdso_thickness, roughness_1)
    hmdso_layer.deposit2D(layout)
    ptfe_layer = ba.Layer(ptfe_mat, ptfe_thickness, roughness_2)
    substrate_layer = ba.Layer(si_mat)

    # assembling sample
    sample = ba.Sample()
    sample.addLayer(vacuum_layer)
    sample.addLayer(hmdso_layer)
    sample.addLayer(ptfe_layer)
    sample.addLayer(substrate_layer)

    return sample

def get_simulation(P):
    """
    Creates and returns GISAS simulation.
    """
    beam = ba.Beam(1.2e7, wavelength, alpha_i)

    sample = get_sample(P)

    # Pilatus detector as used in experiment
    #  (pixel size 0.172 mm, distance from sample 1730.0 mm):
    detector = ba.SphericalDetector(981, -1.7*deg, 1.7*deg, 1043, -0.6*deg, 1.24*deg)

    simulation = ba.ScatteringSimulation(beam, sample, detector)

    # simulation.detector().setRegionOfInterest(85, 70, 120, 92.)
    # beamstop:
    # simulation.detector().addMask(ba.Rectangle(101.9, 82.1, 103.7, 85.2), True)

    return simulation

def load_data(filename):
    """
    Loads experimental data and returns numpy array.
    """
    filepath = os.path.join(datadir, filename)
    return ba.readData2D(filepath)


if __name__ == '__main__':
    data = load_data("scatter2d/galaxi_data.tif.gz")

    fit_objective = ba.FitObjective()
    fit_objective.addFitPair(get_simulation, data)
    fit_objective.initPrint(10)
    observer = ba_fitmonitor.PlotterGISAS()
    <%- if test_mode or figure_mode -%>
    plt.close() # hide plot
    <%- end -%>

    fit_objective.initPlot(10, observer)

    P = lmfit.Parameters()
    P.add("radius", value=5.*nm, min=4*nm, max=6*nm)
    P.add("sigma", value=0.55, min=0.2, max=0.8)
    P.add("distance", value=27.*nm, min=20*nm, max=70*nm)

    <%- if test_mode or figure_mode -%>
    result = lmfit.minimize(fit_objective.evaluate_residuals, P, max_nfev=2)
    <%- else -%>
    result = lmfit.minimize(fit_objective.evaluate_residuals, P)
    <%- end -%>
    print(lmfit.fit_report(result))

    plt.show()
