"""Plot reference state profiles."""

import matplotlib.pyplot as plt

from . import conf, _helpers
from .phyvars import REFSTATE
from .stagyydata import StagyyData


def plot_ref(sdat, var):
    """Plot one reference state.

    Args:
        sdat (:class:`~stagpy.stagyydata.StagyyData`): a StagyyData instance.
        var (str): refstate variable, a key of
            :data:`~stagpy.phyvars.REFSTATE`.
    """
    fig, axis = plt.subplots()
    adbts = sdat.refstate.adiabats
    if len(adbts) > 2:
        for iad, adia in enumerate(adbts[:-1], 1):
            axis.plot(adia[var], adia['z'],
                      conf.refstate.style,
                      label=f'System {iad}')
    axis.plot(adbts[-1][var], adbts[-1]['z'],
              conf.refstate.style, color='k',
              label='Combined profile')
    if var == 'Tcond':
        axis.set_xscale('log')
    axis.set_xlabel(REFSTATE[var].description)
    axis.set_ylabel('z Position')
    if len(adbts) > 2:
        axis.legend()
    _helpers.saveplot(fig, f'refstate_{var}')


def cmd():
    """Implementation of refstate subcommand.

    Other Parameters:
        conf.core
        conf.plot
    """
    sdat = StagyyData()

    lov = conf.refstate.plot.split(',')
    if not lov:
        return
    for var in lov:
        plot_ref(sdat, var)

@@ -31,16 +31,52 @@
 es.%22%22%22%0A%0A
+from __future__ import annotations%0A%0A
 import m
@@ -213,20 +213,37 @@
 ref(sdat
-, va
+: StagyyData, var: st
 r):%0A    
@@ -294,18 +294,20 @@
     sdat
+: a
  
-(
 :class:%60
@@ -340,23 +340,8 @@
 ata%60
-): a StagyyData
  ins
@@ -362,14 +362,8 @@
  var
- (str)
 : re
@@ -387,28 +387,16 @@
 a key of
-%0A           
  :data:%60

Add type annotations to refstate module