#!/usr/bin/env python
# coding: utf-8

# In[1]:


import matplotlib.pyplot as plt
import matplotlib.colors as pltcolors
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
get_ipython().run_line_magic('config', "InlineBackend.figure_format='retina'")

from colorspacious import cspace_convert
import numpy as np
import pyart
from skimage.io import imread
import skimage.color as skcolor
import xarray as xr


# This colormap is built with four configurable segments:
# 
# 1) A grey point of configurable lightness linearly spaced in CIELab to start value of "Spectral_r"
# 2) The full "Spectral_r" map
# 3) A subset of the pink portion of the "PiYG" colormap
# 4) A subset of the purple portion of the "PRGn_r" colormap
# 
# Control points include:
# 
# - Value breakpoints between each segment (to be chosen by some meaningful reflectivity criteria)
# - Value of scale minimum (scale maximum is determined by assuming that segments 3 and 4 are of equal length)
# - Toggle if piecewise linearization of lightness combines segment 1 and the increasing portion of segment 2, or separately
# - Lightness of grey point

# In[2]:


# Control
min_value = -20
start_spectral = 10
start_pink = 60
start_purple = 70
linearize_greys_with_spectral = True
grey_point_lightness = 0
resolution_factor = 4

# Define initial colormaps
spectral_rgb_initial = plt.cm.Spectral_r(
    np.linspace(0, 1, (start_pink - start_spectral) * resolution_factor)
)[:, :3]
pink_rgb_initial = plt.cm.PiYG(
    np.linspace(0, .25, (start_purple - start_pink) * resolution_factor)
)[:, :3]
purple_rgb_initial = plt.cm.PRGn_r(
    np.linspace(0.75, 1, (start_purple - start_pink) * resolution_factor)
)[:, :3]
spectral_lab = skcolor.rgb2lab(spectral_rgb_initial)
pink_lab = skcolor.rgb2lab(pink_rgb_initial)
purple_lab = skcolor.rgb2lab(purple_rgb_initial)
greyscale_lab = np.stack([
    np.linspace(
        grey_point_lightness,
        spectral_lab[0, 0],
        (start_spectral - min_value) * resolution_factor
    ),
    np.linspace(0, spectral_lab[0, 1], (start_spectral - min_value) * resolution_factor),
    np.linspace(0, spectral_lab[0, 2], (start_spectral - min_value) * resolution_factor)
], axis=-1)

# Piecewise linearization of lightness
peak_idx = np.argmax(spectral_lab[:, 0])
peak_value = np.linspace(
    start_spectral, start_pink, (start_pink - start_spectral) * resolution_factor
)[peak_idx]
print(f"Spectral inflection point at value: {peak_value}")
if linearize_greys_with_spectral:
    greys_and_spectral_increasing_lab = np.concatenate([greyscale_lab, spectral_lab[:peak_idx]])
    greys_and_spectral_increasing_lab[:, 0] = np.linspace(
        greys_and_spectral_increasing_lab[0, 0],
        greys_and_spectral_increasing_lab[-1, 0],
        greys_and_spectral_increasing_lab.shape[0]
    )
else:
    spectral_increasing_lab = spectral_lab[:peak_idx].copy()
    spectral_increasing_lab[:, 0] = np.linspace(
        spectral_increasing_lab[0, 0],
        spectral_increasing_lab[-1, 0],
        spectral_increasing_lab.shape[0]
    )
    greys_and_spectral_increasing_lab = np.concatenate([greyscale_lab, spectral_increasing_lab])
spectral_decreasing_lab = spectral_lab[peak_idx:].copy()
spectral_decreasing_lab[:, 0] = np.linspace(
    spectral_decreasing_lab[0, 0],
    spectral_decreasing_lab[-1, 0],
    spectral_decreasing_lab.shape[0]
)
pink_lab[:, 0] = np.linspace(
    spectral_decreasing_lab[-1, 0],
    pink_lab[-1, 0],
    pink_lab.shape[0]
)
purple_lab[:, 0] = np.linspace(
    pink_lab[-1, 0],
    purple_lab[-1, 0],
    purple_lab.shape[0]
)
combined_lab = np.concatenate([
    greys_and_spectral_increasing_lab,
    spectral_decreasing_lab,
    pink_lab,
    purple_lab
])
combined_rgb = skcolor.lab2rgb(combined_lab)

# Save txt
np.savetxt("chase-spectral-202211091600.txt", combined_rgb)

# Colormap!
spectral_refl = pltcolors.LinearSegmentedColormap.from_list("ChaseSpectral", combined_rgb)
spectral_refl


# Now, let's explore some visualizations!

# In[3]:


def colormap_evaluation(sample, cmap, vmin=0, vmax=80):
    # Sample image in original RGB
    fig = Figure()
    canvas = FigureCanvas(fig)
    ax = fig.gca()

    sample.plot.imshow(
        vmin=vmin,
        vmax=vmax,
        cmap=cmap,
        add_labels=False,
        ax=ax
    )

    fig.tight_layout(pad=0)
    ax.margins(0)

    canvas.draw()       # draw the canvas, cache the renderer

    image_from_plot = np.frombuffer(canvas.tostring_rgb(), dtype=np.uint8)
    image_from_plot = image_from_plot.reshape(canvas.get_width_height()[::-1] + (3,))
    image_from_plot = image_from_plot / 255.
    
    # Convert to greyscale
    img_greyscale_JCh = cspace_convert(image_from_plot, "sRGB1", "JCh")
    img_greyscale_JCh[..., 1] = 0
    img_greyscale_sRGB = cspace_convert(img_greyscale_JCh, "JCh", "sRGB1")
    
    # Deuteranomaly
    cvd_space = {
        "name": "sRGB1+CVD",
        "cvd_type": "deuteranomaly",
        "severity": 50
    }
    img_deuteranomaly_sRGB = cspace_convert(image_from_plot, cvd_space, "sRGB1")
    
    # Deuteranopia
    cvd_space = {
        "name": "sRGB1+CVD",
        "cvd_type": "deuteranomaly",
        "severity": 100
    }
    img_deuteranopia_sRGB = cspace_convert(image_from_plot, cvd_space, "sRGB1")
    
    # Protanomaly
    cvd_space = {
        "name": "sRGB1+CVD",
        "cvd_type": "protanomaly",
        "severity": 50
    }
    img_protanomaly_sRGB = cspace_convert(image_from_plot, cvd_space, "sRGB1")
    
    # Deuteranopia
    cvd_space = {
        "name": "sRGB1+CVD",
        "cvd_type": "protanomaly",
        "severity": 100
    }
    img_protanopia_sRGB = cspace_convert(image_from_plot, cvd_space, "sRGB1")
    
    # Overall Figure
    fig, axes = plt.subplot_mosaic(
        [
            ['l', 'l', 'l'],
            ['orig', 'd_anom', 'd_anop'],
            ['bw', 'p_anom', 'p_anop']
        ],
        figsize=(12, 9)
    )
    
    # Color lightness scale
    x = np.linspace(vmin, vmax, int((vmax - vmin) * 4))
    color_range_sRGB = cmap(np.linspace(0, 1, len(x)))[:, :3]
    color_range_Lab = cspace_convert(color_range_sRGB, "sRGB1", "CIELab")
    axes['l'].scatter(
        x,
        color_range_Lab[:, 0],
        c = [tuple(row) for row in color_range_sRGB],
        s = 100
    )
    axes['l'].set_title(f"Radar colormap test for '{cmap.name}'")
    axes['l'].set_ylabel("Lightness")
    
    # Sample images
    axes['orig'].imshow(image_from_plot)
    axes['orig'].axis('off')
    axes['orig'].set_title("Full Color")
    
    axes['bw'].imshow(img_greyscale_sRGB)
    axes['bw'].axis('off')
    axes['bw'].set_title("Greyscale")
    
    axes['d_anom'].imshow(img_deuteranomaly_sRGB)
    axes['d_anom'].axis('off')
    axes['d_anom'].set_title("Deuteranomaly")
    
    axes['d_anop'].imshow(img_deuteranopia_sRGB)
    axes['d_anop'].axis('off')
    axes['d_anop'].set_title("Deuteranopia")
    
    axes['p_anom'].imshow(img_protanomaly_sRGB)
    axes['p_anom'].axis('off')
    axes['p_anom'].set_title("Protanomaly")
    
    axes['p_anop'].imshow(img_protanopia_sRGB)
    axes['p_anop'].axis('off')
    axes['p_anop'].set_title("Protanopia")
    
    return fig


# In[4]:


ds = refl = xr.open_zarr(
    "s3://svrimgdm/v0/svrimg_dm_v0_testing_full.zarr/",
    storage_options={
        'anon': True,
        'client_kwargs': {
            'endpoint_url': 'https://sfo3.digitaloceanspaces.com'
        }
    },
    consolidated=True
)
ds


# In[5]:


# Choose a mode, get a random sample with that mode label
mode = "BE"
idx = np.random.choice(np.flatnonzero(ds['mode'].values == mode))
sample = ds['reflectivity'].isel(sample=idx, time_offset=2).load()
sample


# In[6]:


colormap_evaluation(sample, spectral_refl, vmin=min_value, vmax=2 * start_purple - start_pink)
plt.show()


# In[ ]: