#!/usr/bin/env python3
"""
Generates correlated random rough surfaces for two sample interfaces.
"""

import random
import numpy as np
import bornagain as ba
from bornagain import ba_plot as bp, nm, nm2, nm3
from matplotlib import gridspec


FILM_THICKNESS = 25*nm
CORRELATION_SPLIT_FREQ = 0.035/nm


def correlation(a, b):
    return np.corrcoef(a.flatten(), b.flatten())[0, 1]


def frequency_filtered(h, cutoff, keep_low):
    fy = np.fft.fftfreq(h.shape[0], d=Ly/h.shape[0])
    fx = np.fft.fftfreq(h.shape[1], d=Lx/h.shape[1])
    qx, qy = np.meshgrid(fx, fy)
    q = np.hypot(qx, qy)
    mask = q <= cutoff if keep_low else q >= cutoff
    return np.fft.ifft2(np.fft.fft2(h) * mask).real


def psd_line(h):
    factor = X_points*Y_points/np.sqrt(Lx*Ly)
    n = np.floor(X_points/2 + 1).astype(int)
    psd = abs(np.fft.rfft2(h)/factor)**2
    psd_points_x = np.array([psd[0][i] for i in range(1, n)])
    f_points_x = np.array([i/Lx for i in range(1, n)])
    return f_points_x, psd_points_x


def add_map(ax, h, title, vlim):
    im = ax.imshow(h, extent=[0, Lx, 0, Ly], cmap='inferno',
                   vmin=-vlim, vmax=vlim)
    ax.set_title(title)
    ax.set_xlabel('x, nm')
    ax.set_ylabel('y, nm')
    bp.plt.colorbar(im, ax=ax, shrink=0.75)


def add_psd(ax, h):
    f, psd = psd_line(h)
    ax.plot(f, psd)
    ax.set_yscale('log')
    ax.set_xscale('log')
    ax.set_title('PSD')
    ax.set_xlabel('$\\nu, nm^{-1}$')


def add_histogram(ax, h):
    rms = np.std(h)
    ax.hist(h.flatten(), range=[-4*rms, 4*rms], bins=240, color='black')
    ax.set_title('Height histogram')
    ax.set_xlabel('h, nm')
    ax.tick_params(left=False, labelleft=False)


def plot(film_top, substrate_top):
    vlim = max(np.max(np.abs(film_top)), np.max(np.abs(substrate_top)))

    fig = bp.plt.figure(figsize=(11, 6.5))
    gs = gridspec.GridSpec(2, 3)
    gs.update(left=0.07, right=0.97, bottom=0.09, top=0.92,
              wspace=0.35, hspace=0.35)

    add_map(fig.add_subplot(gs[0, 0]), substrate_top,
            'Substrate interface', vlim)
    add_psd(fig.add_subplot(gs[0, 1]), substrate_top)
    add_histogram(fig.add_subplot(gs[0, 2]), substrate_top)

    add_map(fig.add_subplot(gs[1, 0]), film_top,
            'Film top interface', vlim)
    add_psd(fig.add_subplot(gs[1, 1]), film_top)
    add_histogram(fig.add_subplot(gs[1, 2]), film_top)


def get_sample():
    """
    A sample with one layer on a substrate, with partially
    correlated roughnesses.
    """
    # defining materials
    vacuum_color = (0.90, 0.93, 0.97)
    vacuum = ba.RefractiveMaterial("ambience", vacuum_color, 0, 0)
    layer_color = (0.05, 0.62, 0.55)
    layer_mat = ba.RefractiveMaterial("layer", layer_color, 5e-6, 0)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.RefractiveMaterial("substrate", substrate_color, 15e-6, 0)

    # defining roughness
    height_distribution = ba.ErfTransient()
    max_spat_freq = 0.5/nm

    # for layer
    omega = 1*nm3
    a1 = 1
    a2 = 10*nm
    a3 = 100*nm2
    a4 = 1000*nm3
    autocorr_layer = ba.LinearGrowthModel(omega, a1, a2, a3, a4,
                                          max_spat_freq)
    roughness_layer = ba.Roughness(autocorr_layer, height_distribution)

    # for substrate
    sigma = 0.9*nm
    alpha = 0.5
    xi = 35*nm
    autocorr_sub = ba.SelfAffineFractalModel(sigma, alpha, xi,
                                             max_spat_freq)
    roughness_sub = ba.Roughness(autocorr_sub, height_distribution)

    # defining layers
    l_ambience = ba.Layer(vacuum)
    l_layer = ba.Layer(layer_mat, FILM_THICKNESS, roughness_layer)
    l_substrate = ba.Layer(substrate_mat, roughness_sub)

    # adding layers
    my_sample = ba.Sample()
    my_sample.addLayer(l_ambience)
    my_sample.addLayer(l_layer)
    my_sample.addLayer(l_substrate)

    return my_sample


if __name__ == '__main__':

    # sample size
    Lx = 1000*nm
    Ly = 1000*nm
    X_points = 512
    Y_points = 512

    sample = get_sample()
    seed = random.randrange(2**32)  # random every run

    # generate both rough interfaces together
    maps = ba.RoughnessMap.generateInterfaceMaps(
        X_points, Y_points, Lx, Ly, sample, [1, 2], seed)
    film_top = maps[0]
    substrate_top = maps[1]

    film_low = frequency_filtered(film_top, CORRELATION_SPLIT_FREQ, True)
    substrate_low = frequency_filtered(substrate_top, CORRELATION_SPLIT_FREQ,
                                       True)
    film_high = frequency_filtered(film_top, 2*CORRELATION_SPLIT_FREQ, False)
    substrate_high = frequency_filtered(substrate_top,
                                        2*CORRELATION_SPLIT_FREQ, False)

    print("seed =", seed)
    print("correlation model = LinearGrowthModel")
    print("film thickness = {:.1f} nm".format(FILM_THICKNESS))
    print("frequency split = {:.3f} nm^-1".format(CORRELATION_SPLIT_FREQ))
    print("full-height correlation =",
          "{:.3}".format(correlation(film_top, substrate_top)))
    print("low-frequency correlation =",
          "{:.3}".format(correlation(film_low, substrate_low)))
    print("high-frequency correlation =",
          "{:.3}".format(correlation(film_high, substrate_high)))

    plot(film_top, substrate_top)
    plotargs = bp.parse_commandline()
    bp.export(**plotargs)
