import glob
import numpy as np
import matplotlib.pyplot as plt
from thunderhopper.modeltools import load_data
from thunderhopper.filtertools import find_kern_specs
from misc_functions import unsort_unique
from color_functions import sample_cmap
from IPython import embed

# Settings:
target = 'Omocestus_rufipes'
data_path = glob.glob(f'../data/inv/full/{target}*.npz')[0]
stages = ['conv', 'feat']
load_kwargs = dict(
    files=stages,
    keywords=['scales', 'measure']
)

# Subset settings:
all_types = np.array([1, -1, 2, -2, 3, -3, 4, -4, 5, -5,
                      6, -6, 7, -7, 8, -8, 9, -9, 10, -10]).astype(float)
all_sigmas = np.array([0.001, 0.002, 0.004, 0.008, 0.016, 0.032])
kerns = None
types = np.array([1, -1, 2, -2, 3, -3, 4, -4, 5, -5,
                  6, -6, 7, -7, 8, -8, 9, -9, 10, -10])
sigmas = np.array([0.008, 0.016])

# Plot settings:
line_kwargs = dict(
    linewidth=2,
)
median_kwargs = dict(
    linewidth=4,
    c='k',
    linestyle='--'
)

# Load invariance data:
data, config = load_data(data_path, **load_kwargs)
scales = data['scales']

# Reduce to kernel subset:
if any(var is not None for var in [kerns, types, sigmas]):
    subset_inds = find_kern_specs(config['k_specs'], kerns, types, sigmas)
    data['measure_conv'] = data['measure_conv'][:, subset_inds]
    data['measure_feat'] = data['measure_feat'][:, subset_inds]
    config['kernels'] = config['kernels'][:, subset_inds]
    config['k_specs'] = config['k_specs'][subset_inds, :]
kern_types = unsort_unique(config['k_specs'][:, 0])
kern_sigmas = unsort_unique(config['k_specs'][:, 1])

# Prepare colors:
type_colors = {t: c for t, c in zip(all_types, sample_cmap('turbo', all_types.size))}
sigma_colors = {s: c for s, c in zip(all_sigmas, sample_cmap('turbo', all_sigmas.size))}

# Prepare graph:
fig, axes = plt.subplots(2, 4, figsize=(16, 16), layout='constrained', sharex=True)
axes[0, 0].set_xlim(scales[0], scales[-1])
axes[0, 0].set_xscale('log')
axes[0, 0].set_ylabel('conv')
axes[1, 0].set_ylabel('feat')

# Condense across kernels:
median_conv = np.median(data['measure_conv'], axis=1)
median_feat = np.median(data['measure_feat'], axis=1)

# Coded by type:
leg_handles = []
for kern_type in kern_types:
    color = type_colors[kern_type]
    inds = find_kern_specs(config['k_specs'], types=kern_type)
    leg_handles.append(axes[0, 0].plot(scales, data['measure_conv'][:, inds],
                       c=color, label=f'{kern_type}', **line_kwargs)[0])
    axes[0, 0].plot(scales, median_conv, **median_kwargs)
    axes[1, 0].plot(scales, data['measure_feat'][:, inds], c=color, **line_kwargs)
    axes[1, 0].plot(scales, median_feat, **median_kwargs)
axes[0, 0].legend(handles=leg_handles, loc='upper left')
axes[0, 0].set_title('Coded by type')

# Coded by sigma:
leg_handles = []
for kern_sigma in kern_sigmas:
    color = sigma_colors[kern_sigma]
    inds = find_kern_specs(config['k_specs'], sigmas=kern_sigma)
    leg_handles.append(axes[0, 1].plot(scales, data['measure_conv'][:, inds],
                       c=color, label=f'{kern_sigma}', **line_kwargs)[0])
    axes[0, 1].plot(scales, median_conv, **median_kwargs)
    axes[1, 1].plot(scales, data['measure_feat'][:, inds], c=color, **line_kwargs)
    axes[1, 1].plot(scales, median_feat, **median_kwargs)
axes[0, 1].legend(handles=leg_handles, loc='upper left')
axes[0, 1].set_title('Coded by sigma')

# Normalize measures:
data['measure_conv'] -= data['measure_conv'].min(axis=0)
data['measure_conv'] /= data['measure_conv'].max(axis=0)
data['measure_feat'] -= data['measure_feat'].min(axis=0)
data['measure_feat'] /= data['measure_feat'].max(axis=0)

# Condense across kernels:
median_conv = np.median(data['measure_conv'], axis=1)
median_feat = np.median(data['measure_feat'], axis=1)

# Coded by type:
leg_handles = []
for kern_type in kern_types:
    color = type_colors[kern_type]
    inds = find_kern_specs(config['k_specs'], types=kern_type)
    leg_handles.append(axes[0, 2].plot(scales, data['measure_conv'][:, inds],
                       c=color, label=f'{kern_type}', **line_kwargs)[0])
    axes[0, 2].plot(scales, median_conv, **median_kwargs)
    axes[1, 2].plot(scales, data['measure_feat'][:, inds], c=color, **line_kwargs)
    axes[1, 2].plot(scales, median_feat, **median_kwargs)
axes[0, 2].legend(handles=leg_handles, loc='upper left')
axes[0, 2].set_title('Coded by type')


# Coded by sigma:
leg_handles = []
for kern_sigma in kern_sigmas:
    color = sigma_colors[kern_sigma]
    inds = find_kern_specs(config['k_specs'], sigmas=kern_sigma)
    leg_handles.append(axes[0, 3].plot(scales, data['measure_conv'][:, inds],
                       c=color, label=f'{kern_sigma}', **line_kwargs)[0])
    axes[0, 3].plot(scales, median_conv, **median_kwargs)
    axes[1, 3].plot(scales, data['measure_feat'][:, inds], c=color, **line_kwargs)
    axes[1, 3].plot(scales, median_feat, **median_kwargs)
axes[0, 3].legend(handles=leg_handles, loc='upper left')
axes[0, 3].set_title('Coded by sigma')

plt.show()
embed()