import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import os

out = '/mnt/user-data/outputs'
os.makedirs(out, exist_ok=True)

# Data
hp_cap   = [0, 1, 3.2, 13.2, 113.2]
hp_rise  = [1055.7, 860.09, 871.09, 901.59, 921.83]
hp_rise_std = [203.26, 41.857, 72.217, 49.143, 45.386]
hp_vpp   = [207.96, 461.03, 517.96, 548.20, 555.19]
hp_vpp_std = [8.8376, 7.4999, 10.208, 9.8535, 9.2046]

lp_cap   = [0, 1, 11, 111]
lp_rise  = [1055.7, 1118.7, 1757.2, 6850.3]
lp_rise_std = [203.26, 410.87, 529.05, 1038.7]
lp_vpp   = [207.96, 215.38, 221.2, 196.81]
lp_vpp_std = [9.5601, 9.5601, 12.229, 11.408]

BLUE  = '#2A6DB5'
GREEN = '#1A8A5A'

style = {
    'font.family': 'serif',
    'axes.spines.top': False,
    'axes.spines.right': False,
    'axes.linewidth': 0.8,
    'xtick.direction': 'in',
    'ytick.direction': 'in',
    'xtick.major.size': 4,
    'ytick.major.size': 4,
    'grid.linestyle': '--',
    'grid.linewidth': 0.4,
    'grid.alpha': 0.5,
    'figure.dpi': 150,
}

def make_plot(x, y, yerr, color, xlabel, ylabel, title, fname, xlog=False):
    with plt.rc_context(style):
        fig, ax = plt.subplots(figsize=(5.5, 4.0))
        ax.errorbar(x, y, yerr=yerr, fmt='o', color=color,
                    ecolor=color, elinewidth=1.2, capsize=4, capthick=1.2,
                    markersize=6, markerfacecolor='white', markeredgewidth=1.5,
                    zorder=3)
        ax.grid(True, axis='both')
        if xlog:
            ax.set_xscale('log')
        ax.set_xlabel(xlabel, fontsize=11)
        ax.set_ylabel(ylabel, fontsize=11)
        ax.set_title(title, fontsize=11, fontweight='normal', pad=8)
        # note on conditions
        ax.text(0.98, 0.97,
                '5 kHz, 1 V$_{pp}$, 13 µs pulse\nn ≈ 10 000 per point',
                transform=ax.transAxes, fontsize=8,
                ha='right', va='top', color='#555555')
        ax.tick_params(labelsize=10)
        fig.tight_layout()
        path = os.path.join(out, fname)
        fig.savefig(path, format='pdf', bbox_inches='tight')
        plt.close(fig)
        print(f'saved {path}')

# Plot 1: HP rise time vs capacitance
make_plot(hp_cap, hp_rise, hp_rise_std,
          BLUE,
          'Added capacitance to high-pass stage (nF)',
          'Rise time (ns)',
          'Rise time vs. high-pass added capacitance',
          'plot1_hp_risetime.pdf')

# Plot 2: HP Vpp vs capacitance
make_plot(hp_cap, hp_vpp, hp_vpp_std,
          GREEN,
          'Added capacitance to high-pass stage (nF)',
          'V$_{pp}$ (mV)',
          'Peak-to-peak voltage vs. high-pass added capacitance',
          'plot2_hp_vpp.pdf')

# Plot 3: LP rise time vs capacitance (log x for spread)
make_plot(lp_cap, lp_rise, lp_rise_std,
          BLUE,
          'Added capacitance to low-pass stage (nF)',
          'Rise time (ns)',
          'Rise time vs. low-pass added capacitance',
          'plot3_lp_risetime.pdf',
          xlog=False)

# Plot 4: LP Vpp vs capacitance
make_plot(lp_cap, lp_vpp, lp_vpp_std,
          GREEN,
          'Added capacitance to low-pass stage (nF)',
          'V$_{pp}$ (mV)',
          'Peak-to-peak voltage vs. low-pass added capacitance',
          'plot4_lp_vpp.pdf')

print('All done.')