import matplotlib.pyplot as plt

# Fixed RCs:
plt.rcParams['font.style'] = 'normal'
plt.rcParams['mathtext.fontset'] = 'cm'
plt.rcParams['svg.fonttype'] = 'path'
plt.rc('text.latex', preamble=r'\usepackage{amsmath}')

# Context RCs:
show_figs = True
rc_tex = {
    'text.usetex': True,
    'font.family': 'computer modern roman',
}
rc_plt = {
    'text.usetex': False,
    'font.family': 'sans-serif',
}

# Instance configs:
title_props = dict(
    use_tex=False,
    figsize=(5, 2.5),
    fs=53,
    fw='normal',
)
letter_props = dict(
    use_tex=False,
    figsize=(1, 1),
    fs=70,
    fw='bold',
)
element_props = dict(
    use_tex=True,
    figsize=(3, 3),
    fs=80,
    fw='normal',
)
line_props = dict(
    use_tex=True,
    figsize=(3, 2),
    fs=75,
    fw='normal',
)

# Auxiliary configs:
fig_kwargs = dict(
    gridspec_kw={'left': 0, 'right': 1, 'top': 1, 'bottom': 0},
    facecolor='none',
    edgecolor='none',
    frameon=False,
)
ax_kwargs = dict(
    facecolor='none',
    frame_on=False,
)
text_kwargs = dict(
    color='k',
    x=0.5,
    y=0.5,
    ha='center',
    va='center',
)

# Targets:
texts = {
    ### PLAIN FONT ###

    # TITLES (NEURONAL CIRCUIT):
    'neuron_titles_tympanum': ('Tympanal\nMembrane', title_props),
    'neuron_titles_receptors': ('Receptor\nNeurons', title_props),
    'neuron_titles_interneurons': ('Local\nInterneurons', title_props),
    'neuron_titles_ascending': ('Ascending\nNeurons', title_props),
    'neuron_titles_brain': ('Central\nBrain', title_props),

    # TITLES (MODEL CIRCUIT):
    'model_titles_bandpass': ('Bandpass\nFiltering', title_props),
    'model_titles_envelope': ('Envelope\nExtraction', title_props),
    'model_titles_logarithm': ('Logarithmic\nCompression', title_props),
    'model_titles_adaptation': ('Intensity\nAdaptation', title_props),
    'model_titles_convolution': ('Convolutional\nFiltering', title_props),
    'model_titles_nonlinear': ('Threshold\nNonlinearity', title_props),
    'model_titles_integration': ('Temporal\nAveraging', title_props),
    'model_titles_readout': ('Weighted\nReadout', title_props),

    # SUBPLOT LETTERS:
    'subplot_a': ('a', letter_props),
    'subplot_b': ('b', letter_props),
    'subplot_c': ('c', letter_props),
    'subplot_d': ('d', letter_props),

    ### TEX FONT ###

    # ELEMENT LABELS (MODEL CIRCUIT):
    'model_elements_filt': (r'$x_{\text{filt}}$', element_props),
    'model_elements_env': (r'$x_{\text{env}}$', element_props),
    'model_elements_log': (r'$x_{\text{dB}}$', element_props),
    'model_elements_adapt': (r'$x_{\text{adapt}}$', element_props),
    'model_elements_c1': (r'$c_1$', element_props),
    'model_elements_c2': (r'$c_2$', element_props),
    'model_elements_c3': (r'$c_3$', element_props),
    'model_elements_b1': (r'$b_1$', element_props),
    'model_elements_b2': (r'$b_2$', element_props),
    'model_elements_b3': (r'$b_3$', element_props),
    'model_elements_f1': (r'$f_1$', element_props),
    'model_elements_f2': (r'$f_2$', element_props),
    'model_elements_f3': (r'$f_3$', element_props),
    'model_elements_out': (r'$\hat{y}$', element_props),

    # LINE LABELS (MODEL CIRCUIT):
    'model_lines_env': (r'$\lvert\cdot\lvert,h_{\text{LP}}$', line_props),
    'model_lines_log': (r'$\text{log}$', line_props),
    'model_lines_hp': (r'$h_{\text{HP}}$', line_props),
    'model_lines_k1': (r'$k_1$', line_props),
    'model_lines_k2': (r'$k_2$', line_props),
    'model_lines_k3': (r'$k_3$', line_props),
    'model_lines_t1': (r'$\Theta_1$', line_props),
    'model_lines_t2': (r'$\Theta_2$', line_props),
    'model_lines_t3': (r'$\Theta_3$', line_props),
    # 'model_lines_H1': (r'$H(c_1-\Theta_1)$', line_props),
    # 'model_lines_H2': (r'$H(c_2-\Theta_2)$', line_props),
    # 'model_lines_H3': (r'$H(c_3-\Theta_3)$', line_props),
    'model_lines_lp': (r'$h_{\text{LP}}$', line_props),
    'model_lines_w1': (r'$\omega_1$', line_props),
    'model_lines_w2': (r'$\omega_2$', line_props),
    'model_lines_w3': (r'$\omega_3$', line_props),
}

# Save each target string:
for name, (text, props) in texts.items():
    plt.rcParams.update(rc_tex if props['use_tex'] else rc_plt)
    fig, ax = plt.subplots(figsize=props['figsize'], **fig_kwargs)
    ax.set(**ax_kwargs)
    ax.axis('off')
    ax.text(s=text, size=props['fs'], weight=props['fw'], **text_kwargs)
    fig.savefig(f'../figures/{name}_text.svg', dpi=300,
                bbox_inches='tight', pad_inches=0)
    if show_figs:
        plt.show()
    plt.close(fig)