#!/usr/bin/env python3
"""Original two-compartment pressure-step response; output PNG in cwd."""
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


def main():
    t = np.linspace(0, 6, 1201)
    fast = .1 * (1 - np.exp(-t / .2))
    slow = .1 * (1 - np.exp(-t / 2))
    fig, ax = plt.subplots(figsize=(8, 3.5), dpi=120, facecolor='white')
    ax.plot(t, fast, color='#2367a1', label='Fast region: time constant 0.2 s')
    ax.plot(t, slow, color='#b85224', label='Slow region: time constant 2 s')
    ax.plot(t, fast + slow, color='#24845e', lw=2.2, label='Total finite-time response')
    ax.axhline(.2, color='0.4', ls=':', label='Equilibrium total: 0.20')
    for when, label_x, label_y in [(.5, 1.1, .133), (2., 2.6, .15)]:
        value = .1 * (1 - np.exp(-when / .2)) + .1 * (1 - np.exp(-when / 2))
        ax.scatter([when], [value], color='#24845e', zorder=5, s=24)
        ax.annotate(f'{when:g} s: {value:.3f}', (when, value), xytext=(label_x, label_y),
                    fontsize=9, arrowprops={'arrowstyle': '-', 'color': '0.4'})
    ax.set(xlim=(0, 6), ylim=(0, .22), xlabel='Time after pressure step (s)',
           ylabel='Volume response / pressure (L per cm H2O)')
    ax.set_facecolor('white')
    ax.grid(alpha=.18)
    ax.legend(loc='lower right', fontsize=8)
    ax.set_title('Unequal filling times hide equilibrium compliance', fontsize=13)
    fig.text(.5, .015, 'Constructed passive model; not measured patient data', ha='center', fontsize=9)
    fig.tight_layout(rect=(0, .055, 1, 1))
    fig.savefig('lung-compartment-filling-times.png', dpi=120, facecolor='white', transparent=False)
    plt.close(fig)


if __name__ == '__main__':
    main()