"""Overview plots for the Young's modulus look-up table

The Young's modulus is corrected for pixelation (pixel size 0.34 µm).
The temperature is set to 23°C.
The 20µm plot is used in the dclab documentation. The others
are used in the Shape-Out documentation.
"""
from dclab.features import emodulus
import numpy as np
import seaborn as sns
import matplotlib.pylab as plt
from matplotlib import cm, colors
from matplotlib.ticker import MultipleLocator

sns.set_style("whitegrid")

media = ["CellCarrier", "CellCarrier B", "water"]
matrix = {  # channel width: flow rates
    15: [0.016, 0.032, 0.048],
    20: [0.040, 0.080, 0.120],
    30: [0.160, 0.240, 0.320],
    40: [0.320, 0.400, 0.600],
}
ylim = (0, 0.2)
xlim = {  # channel width: xmax
    15: (10, 200),
    20: (20, 350),
    30: (50, 800),
    40: (120, 1400),
}

deform_max = 0.17
grid_size = 1000
emin = 4e-2
emax = 2e1
norm = colors.LogNorm(vmin=emin, vmax=emax)
cmap = "jet"


for lut_name in [
    "LE-2D-FEM-19",
    "HE-2D-FEM-22",
    "HE-3D-FEM-22"
]:

    for channel_width in matrix:
        flow_rates = matrix[channel_width]
        # setup figure
        fig, axes = plt.subplots(nrows=3, ncols=3, sharex=True, sharey=True,
                                 gridspec_kw={"left": 0.12,
                                              "right": .85,
                                              "wspace": .1,
                                              "hspace": .15,
                                              },
                                 figsize=(7.5, 6.3))
        fig.suptitle("{} µm channel".format(channel_width))
        axes[2][1].set_xlabel("Area [µm²]")
        axes[1][0].set_ylabel("Deformation")
        # column label (flow rate)
        for ii in range(3):
            axes[0][ii].set_title("{} µL/s".format(flow_rates[ii]),
                                  fontsize=10, fontweight="bold")
        # row label (medium)
        for jj in range(3):
            pos = axes[jj][0].get_position()
            x = pos.x0 - .11
            y = pos.y0 + pos.height/2
            fig.text(x, y, media[jj], ha='left', va='center', rotation=90,
                     fontsize=10, fontweight="bold")

        # plot the LUTs channel-wise
        lut_base, lut_meta = emodulus.load_lut(lut_name)
        deform = np.linspace(lut_base[:, 1].min(), lut_base[:, 1].max(),
                             grid_size, endpoint=True)
        area_um = np.linspace(lut_base[:, 0].min(), lut_base[:, 0].max(),
                              grid_size, endpoint=True)
        for jj, medium in enumerate(media):  # rows
            for ii, flow_rate in enumerate(matrix[channel_width]):  # columns
                area_um_scaled = emodulus.scale_linear.scale_area_um(
                    area_um=area_um,
                    channel_width_in=lut_meta["channel_width"],
                    channel_width_out=channel_width)
                aa, dd = np.meshgrid(area_um_scaled, deform)
                emod = emodulus.get_emodulus(area_um=aa,
                                             deform=dd,
                                             medium=medium,
                                             channel_width=channel_width,
                                             flow_rate=flow_rate,
                                             px_um=.34,
                                             temperature=23.0,
                                             visc_model="buyukurganci-2022",
                                             lut_data=lut_name)
                ax = axes[jj][ii]
                ax.set_aspect("auto")
                pc = ax.pcolormesh(aa, dd, emod, shading="nearest", cmap=cmap,
                                   norm=norm, zorder=11)
                ax.grid(zorder=10)
                # ticks every .5
                ax.yaxis.set_major_locator(MultipleLocator(0.05))

        axes[2][2].set_ylim(*ylim)
        axes[2][2].set_xlim(*xlim[channel_width])
        axes[2][2].grid(True, which="minor", zorder=10)

        # add colorbar
        cax = plt.axes([0.87, 0.1, 0.02, 0.8])

        mappable = cm.ScalarMappable(norm=norm, cmap=cmap)
        plt.colorbar(pc, orientation="vertical", cax=cax, extend="both",
                     label=f"apparent Young's modulus ({lut_name}) [kPa]")

        plt.savefig("_emodulus_{}um_{}.png".format(channel_width, lut_name),
                    dpi=150)
        plt.close()
