import paths
import numpy as np
from astropy import units as u
from astropy import constants
from astropy import log
from astropy import table
from astropy import modeling
from astropy.modeling.models import custom_model

import dust_emissivity

import pylab as pl

import lines
import salt_tables
from lte_modeling_tools import (rovib_lte_model_generator,
                                simple_lte_model_generator, kkms_of_nupper,
                                nupper_of_kkms)
from pyspeckit.spectrum.models import lte_molecule as plm
from pyspeckit.spectrum.models.lte_molecule import get_molecular_parameters

vib_constants = {'KCl': (281*u.cm**-1).to(u.eV, u.spectral()).to(u.K, u.temperature_energy()),
                 'K37Cl': (275.87497*u.cm**-1).to(u.eV, u.spectral()).to(u.K, u.temperature_energy()),
                 '41KCl': (276.639613*u.cm**-1).to(u.eV, u.spectral()).to(u.K, u.temperature_energy()),
                 'NaCl': (366*u.cm**-1).to(u.eV, u.spectral()).to(u.K, u.temperature_energy()),
                 # couldn't find numbers; assume it's reduced by mass (which
                 # isn't a good assumption)
                 'Na37Cl': (35+23)/(37+23)*(366*u.cm**-1).to(u.eV, u.spectral()).to(u.K, u.temperature_energy()),
                }

def get_vib_energies(moltbl):
    return {v: moltbl['E_L'][(moltbl['vu'] == v) & (moltbl['Jl'] == 0)].quantity[0].value
            for v in np.unique(moltbl['vu'])}

def get_rot_energies(moltbl):
    def matchJu(ju):
        rslt = moltbl['E_U'][(moltbl['vu'] == 0) & (moltbl['Ju'] == ju)].quantity
        if len(rslt) == 1:
            return rslt
    return {ju: matchJu(ju)[0].value
            for ju in np.unique(moltbl['Ju'])
            if matchJu(ju)
           }

def fit_multi_tex(eupper, nupperoverg, vstate, jstate, vibenergies,
                  rotenergies, verbose=False, plot=False, uplims=None,
                  errors=None,
                  min_nupper=1,
                  replace_errors_with_uplims=False,
                  #molecule=None,
                  partition_func=None,
                  colors='rgbcmyk',
                  molname=None,
                  collims=(np.log(1e9), np.log(1e17)),
                  rottemlims=(10,300),
                  vibtemlims=(500,16000),
                  marker='o', max_uplims='half'):
    """
    Fit the Boltzmann diagram with a vibrational and a rotational temperature

    Parameters
    ----------
    max_uplims: str or number
        The maximum number of upper limits before the fit is ignored completely
        and instead zeros are returned
    """

    nupperoverg_tofit = nupperoverg.copy().to(u.cm**-2).value
    if errors is not None:
        errors = errors.to(u.cm**-2).value

    if uplims is not None:
        upperlim_mask = nupperoverg < uplims

        # allow this magical keyword 'half'
        max_uplims = len(nupperoverg)/2. if max_uplims == 'half' else max_uplims

        if upperlim_mask.sum() > max_uplims:
            # too many upper limits = bad idea to fit.
            return 0*u.cm**-2, 0*u.K, 0, 0

        if errors is None:
            # if errors are not specified, we set the upper limits as values to
            # be fitted
            # (which gives a somewhat useful upper limit on the temperature)
            nupperoverg_tofit[upperlim_mask] = uplims[upperlim_mask]
        else:
            # otherwise, we set the values to zero-column and set the errors to
            # be whatever the upper limits are (hopefully a 1-sigma upper
            # limit)
            # 1.0 here becomes 0.0 in log and makes the relative errors meaningful
            nupperoverg_tofit[upperlim_mask] = 1.0
            if replace_errors_with_uplims:
                errors[upperlim_mask] = uplims[upperlim_mask]
    else:
        upperlim_mask = np.ones_like(nupperoverg_tofit, dtype='bool')

    # always ignore negatives & really low values
    good = nupperoverg_tofit > min_nupper
    # skip any fits that have fewer than 50% good values
    if good.sum() < len(nupperoverg_tofit)/2.:
        return 0*u.cm**-2, 0*u.K, 0, 0

    if errors is not None:
        rel_errors = errors / nupperoverg_tofit
        weights = 1. / rel_errors**2
        log.debug("Fitting with data = {0}, weights = {1}, errors = {2},"
                  "relative_errors = {3}"
                  .format(np.log(nupperoverg_tofit[good]),
                          np.log(weights[good]),
                          errors[good],
                          rel_errors[good],
                         ))
    else:
        # want log(weight) = 1
        weights = np.exp(np.ones_like(nupperoverg_tofit))

    model = rovib_lte_model_generator(vibenergies, rotenergies)

    #print('vibenergies:',vibenergies, '\nrotenergies:', rotenergies)

    fitter = modeling.fitting.LevMarLSQFitter()
    #fitter = modeling.fitting.LinearLSQFitter()

    model_to_fit = model
    model_to_fit.logcolumn.bounds = collims
    model_to_fit.rottem.bounds = rottemlims
    model_to_fit.vibtem.bounds = vibtemlims
    #model_to_fit.rottem.fixed = True

    # check that the model works
    #first_check = model_to_fit(eupper[good].to(u.K).value)
    #print("Result of first check: ",first_check)

    result = fitter(model_to_fit,
                    jstate[good].astype('int'),
                    vstate[good].astype('int'),
                    np.log(nupperoverg_tofit[good]),
                    weights=np.log(weights[good]))
    print(result)
    tex = result.rottem #u.Quantity(-1./result.slope, u.K)
    tvib = result.vibtem

    #partition_func = specmodel.calculate_partitionfunction(molecule.data['States'],
    #                                                       temperature=tex.value)
    #assert len(partition_func) == 1
    #Q_rot = tuple(partition_func.values())[0]
    #print("Q_rot:", Q_rot)
    Q_rot = partition_func(tex.value)

    Ntot = np.exp(result.logcolumn + np.log(Q_rot)) * u.cm**-2

    if verbose:
        print(("Tex={0}, Ntot={1}, log(Ntot)={4}, Q_rot={2}, "
               "nuplim={3}".format(tex, Ntot, Q_rot, upperlim_mask.sum(),
                                   np.log10(Ntot.value),
                                  )))

    if plot:
        import pylab as pl
        for vib,color in zip(np.arange(vstate.max()+1), colors):
            mask = vstate == vib
            if mask.sum() == 0:
                continue
            L, = pl.plot(eupper[mask], np.log10(nupperoverg_tofit[mask]),
                         marker=marker,
                         color=color, markeredgecolor='none', alpha=0.5,
                         linestyle='none',
                         #markersize=2,
                        )
            if uplims is not None:
                L, = pl.plot(eupper[upperlim_mask & mask],
                             np.log10(uplims)[upperlim_mask & mask], 'bv', alpha=0.2,
                             markersize=2)
                #L, = pl.plot(eupper[upperlim_mask],
                #             np.log10(nupperoverg)[upperlim_mask], 'bv', alpha=0.2)
            if errors is not None:
                yerr = np.array([np.log10(nupperoverg_tofit)-np.log10(nupperoverg_tofit-errors),
                                 np.log10(nupperoverg_tofit+errors)-np.log10(nupperoverg_tofit)])
                # if lower limit is nan, set to zero
                yerr[0,:] = np.nan_to_num(yerr[0,:])
                if np.any(np.isnan(yerr[1,:])):
                    print(ValueError("*** Some upper limits are NAN"))
                # use 'good' to exclude plotting errorbars for upper limits
                pl.errorbar(eupper.value[good & mask],
                            np.log10(nupperoverg_tofit)[good & mask],
                            yerr=yerr[:,good & mask],
                            linestyle='none',
                            linewidth=0.5,
                            color=L.get_color(),
                            marker='.', zorder=-5,
                            markersize=2)
            #xax = np.array([0, eupper.max().value+500])

            jstates = np.arange(1,np.max(list(rotenergies.keys()))+1)
            vstates = np.ones(np.max(list(rotenergies.keys())))*vib
            line = result(jstates, vstates)
            xax = np.array([rotenergies[ju] + vibenergies[vib]
                            for ju in jstates])

            for marker, mask in (('.', (jstates>-1) & (jstates <= 10)),
                                 ('s', (jstates>10) & (jstates <= 20)),
                                 ('o', (jstates>20) & (jstates <= 30)),
                                 ('D', (jstates>30) & (jstates <= 40)),
                                 ('p', (jstates>40) & (jstates <= 50))):
                pl.plot(xax[mask],
                        np.log10(np.exp(line))[mask],
                        marker=marker,
                        linestyle='none',
                        color=color,
                        alpha=0.3,
                        linewidth=1.0,
                        markersize=4,
                       )
                    #label='v={0}'.format(vib),
                    #label=('v={2}$T_R={0:0.1f}$ $T_v={3:0.1f}$ $\log(N)={1:0.1f}$'
                    #       .format(tex.value, np.log10(Ntot.value), vib,
                    #               tvib.value
                    #              ))

        for jj in np.unique(jstate):
            #jstates = np.arange(1,np.max(list(rotenergies.keys()))+1)
            #vstates = np.ones(np.max(list(rotenergies.keys())))*vib
            vstates = np.arange(9)
            jstates = np.ones(9)*jj
            line = result(jstates, vstates)
            xax = np.array([rotenergies[jj] + vibenergies[vv]
                            for vv in vstates])

            pl.plot(xax, np.log10(np.exp(line)), '--',
                    color='k',
                    alpha=0.1,
                    linewidth=1.0,
                   )
                    #label="Ju={0}".format(jj),
                    #label=('v={2}$T_R={0:0.1f}$ $T_v={3:0.1f}$ $\log(N)={1:0.1f}$'
                    #       .format(tex.value, np.log10(Ntot.value), vib,
                    #               tvib.value
                    #              ))

        for eu, nu, jj, vv in zip(eupper, nupperoverg_tofit, jstate, vstate):
            model_nu = result(jj, vv)
            #print("POINT MATCHING: ",eu, model_nu, np.log(nu), jj, vv)
            pl.plot(u.Quantity([eu, eu]),
                    np.log10(np.exp([np.log(nu), model_nu])),
                    linestyle='-', alpha=0.3,
                    color='k',
                    linewidth=0.5)

        pl.ylabel("log N$_u$/g (cm$^{-2}$)")
        pl.xlabel("E$_u$ (K)")

        pl.plot([], [], linestyle='-', color='k',
                label=('$T_R={0:0.1f}$ $T_v={2:0.1f}$ $\log(N)={1:0.1f}$'
                       .format(tex.value, np.log10(Ntot.value), tvib.value)))

        if (uplims is not None) and ((errors is None) or replace_errors_with_uplims):
            # if errors are specified, their errorbars will be better
            # representations of what's actually being fit
            pl.plot(eupper, np.log10(uplims), marker='_', alpha=0.5,
                    linestyle='none', color='k')

    return Ntot, tex, result.rottem, result.vibtem, result.logcolumn

def fit_tex(eupper, nupperoverg, verbose=False, plot=False, uplims=None,
            errors=None, min_nupper=1,
            replace_errors_with_uplims=False,
            #molecule=None,
            partition_func=None,
            color='r',
            marker='o',
            max_uplims='half',
            label='',
            axis=None,
            labeltemonly=False,
            logy=True,
           ):
    """
    Fit the Boltzmann diagram

    Parameters
    ----------
    max_uplims: str or number
        The maximum number of upper limits before the fit is ignored completely
        and instead zeros are returned
    """

    nupperoverg_tofit = nupperoverg.copy().to(u.cm**-2).value
    if errors is not None:
        errors = errors.to(u.cm**-2).value

    if uplims is not None:
        upperlim_mask = nupperoverg < uplims

        # allow this magical keyword 'half'
        max_uplims = len(nupperoverg)/2. if max_uplims == 'half' else max_uplims

        if upperlim_mask.sum() > max_uplims:
            # too many upper limits = bad idea to fit.
            return 0*u.cm**-2, 0*u.K, 0, 0

        if errors is None:
            # if errors are not specified, we set the upper limits as values to
            # be fitted
            # (which gives a somewhat useful upper limit on the temperature)
            nupperoverg_tofit[upperlim_mask] = uplims[upperlim_mask]
        else:
            # otherwise, we set the values to zero-column and set the errors to
            # be whatever the upper limits are (hopefully a 1-sigma upper
            # limit)
            # 1.0 here becomes 0.0 in log and makes the relative errors meaningful
            nupperoverg_tofit[upperlim_mask] = 1.0
            if replace_errors_with_uplims:
                errors[upperlim_mask] = uplims[upperlim_mask]
    else:
        upperlim_mask = np.ones_like(nupperoverg_tofit, dtype='bool')

    # always ignore negatives & really low values
    good = nupperoverg_tofit > min_nupper
    # skip any fits that have fewer than 50% good values
    if good.sum() < len(nupperoverg_tofit)/2.:
        return 0*u.cm**-2, 0*u.K, 0, 0

    if errors is not None:
        rel_errors = errors / nupperoverg_tofit
        weights = 1. / rel_errors**2
        log.debug("Fitting with data = {0}, weights = {1}, errors = {2},"
                  "relative_errors = {3}"
                  .format(np.log(nupperoverg_tofit[good]),
                          np.log(weights[good]),
                          errors[good],
                          rel_errors[good],
                         ))
    else:
        # want log(weight) = 1
        weights = np.exp(np.ones_like(nupperoverg_tofit))

    #model = modeling.models.Linear1D()
    model = simple_lte_model_generator()
    fitter = modeling.fitting.LevMarLSQFitter()
    #fitter = modeling.fitting.LinearLSQFitter()
    #model.slope.bounds = [0, -1000]
    if good.sum() >= 2:
        result = fitter(model,
                        eupper[good].to(u.K).value,
                        np.log(nupperoverg_tofit[good]),
                        weights=np.log(weights[good]))
        #tex = u.Quantity(-1./result.slope, u.K)
        tex = u.Quantity(result.tem, u.K)
        if tex < 0*u.K:
            tex = 100 * u.K

        #partition_func = specmodel.calculate_partitionfunction(molecule.data['States'],
        #                                                       temperature=tex.value)
        #assert len(partition_func) == 1
        #Q_rot = tuple(partition_func.values())[0]
        Q_rot = partition_func(tex.value)

        #Ntot = np.exp(result.intercept + np.log(Q_rot)) * u.cm**-2
        Ntot = np.exp(result.logcolumn) * Q_rot * u.cm**-2
    else:
        Ntot = 1e10*u.cm**-2
        tex = 100*u.K
        Q_rot = 1
        result = model

    if verbose:
        print(("Tex={0}, Ntot={1}, log(Ntot)={4}, Q_rot={2}, "
               "nuplim={3}".format(tex, Ntot, Q_rot, upperlim_mask.sum(),
                                   np.log10(Ntot.value),
                                  )))

    if plot:
        import pylab as pl
        ax = pl.gca() if axis is None else axis
        yplot = np.log10(nupperoverg_tofit) if logy else nupperoverg_tofit
        L, = ax.plot(eupper, yplot, marker=marker,
                     color=color, markeredgecolor='none', alpha=0.5,
                     linestyle='none',
                     #markersize=2,
                    )
        if uplims is not None:
            yupplims = np.log10(uplims) if logy else uplims
            L, = ax.plot(eupper[upperlim_mask],
                         yuplims[upperlim_mask], 'bv', alpha=0.2,
                         markersize=2)
            #L, = ax.plot(eupper[upperlim_mask],
            #             np.log10(nupperoverg)[upperlim_mask], 'bv', alpha=0.2)
        if errors is not None:
            if logy:
                yerr = np.array([np.log10(nupperoverg_tofit)-np.log10(nupperoverg_tofit-errors),
                                 np.log10(nupperoverg_tofit+errors)-np.log10(nupperoverg_tofit)])
            else:
                yerr = errors
            # if lower limit is nan, set to zero
            yerr[0,:] = np.nan_to_num(yerr[0,:])
            if np.any(np.isnan(yerr[1,:])):
                #raise ValueError("*** Some upper limits are NAN")
                print(ValueError("*** Some upper limits are NAN"))
            # use 'good' to exclude plotting errorbars for upper limits
            ax.errorbar(eupper.value[good],
                        yplot[good],
                        yerr=yerr[:,good],
                        linestyle='none',
                        linewidth=0.5,
                        color='k',
                        marker='.', zorder=-5,
                        markersize=2)
        xax = np.array([0, eupper.max().value+500])
        line = (xax*(-1/result.tem.value) +
                result.logcolumn.value)
        if labeltemonly:
            labeltext = '{2}$T={0:0.1f}$'.format(tex, np.log10(Ntot.value), label)
        else:
            labeltext = '{2}$T={0:0.1f}$ $\log(N)={1:0.1f}$'.format(tex, np.log10(Ntot.value), label)

        if logy:
            pline = np.log10(np.exp(line))
        else:
            pline = np.exp(line)
        ax.plot(xax, pline, '--',
                color=color,
                alpha=0.6,
                linewidth=1.0,
                label=labeltext)
        ax.set_ylabel("log N$_u$/g (cm$^{-2}$)")
        ax.set_xlabel("E$_u$ (K)")

        if (uplims is not None) and ((errors is None) or replace_errors_with_uplims):
            # if errors are specified, their errorbars will be better
            # representations of what's actually being fit
            ax.plot(eupper, np.log10(uplims), marker='_', alpha=0.5,
                    linestyle='none', color='k')

    return Ntot, tex, result.tem, result.logcolumn

def get_deg_(frqs):
    def get_deg(freq):
        closest = np.nanargmin(np.abs(freq-frqs))
        assert frqs[closest]-freq < 1*u.MHz
        return degeneracies[closest]
    return get_deg

def get_Aul_(frqs):
    def get_Aul(freq):
        closest = np.nanargmin(np.abs(freq-frqs))
        assert frqs[closest]-freq < 1*u.MHz
        return einsteinAij[closest]
    return get_Aul

if __name__ == "__main__":

    all_lines = {**lines.disk_lines, **lines.absorbers}

    frequencies = u.Quantity([all_lines[x] for x in all_lines
                              if 'Cl' in x])

    # from astroquery.vamdc import Vamdc
    # from vamdclib import specmodel

    # kcl35 = Vamdc.query_molecule('KCl-35')
    # rt35 = kcl35.data['RadiativeTransitions']
    # frqs = u.Quantity([(float(rt35[key].FrequencyValue)*u.MHz).to(u.GHz,
    #                                                               u.spectral())
    #                    for key in rt35])

    # upperStateRefs = [rt35[key].UpperStateRef for key in rt35]
    # degeneracies = [int(kcl35.data['States'][upperStateRef].TotalStatisticalWeight)
    #                 for upperStateRef in upperStateRefs]
    # einsteinAij = u.Quantity([float(rt35[key].TransitionProbabilityA) for key in rt35], 1/u.s)

    frqs, einsteinAij, degeneracies, EU, partfunc = get_molecular_parameters('KCl', fmin=85*u.GHz, fmax=360*u.GHz)

    tbl = table.Table.read(paths.tpath('fitted_stacked_lines.txt'), format='ascii.fixed_width')

    kcl35mask = np.array([(not hasattr(row['Species'], 'mask')) and
                         ('KCl' == row['Species'][:3] or
                          '39K-35Cl' in row['Species']) for row in tbl])

    bad = ((tbl['Line Name'] == 'KClv=4_13-12') |
           (tbl['Line Name'] == 'KClv=6_31-30') |
           (tbl['Line Name'] == 'KClv=8_27-26'))

    print("KCl: {0} in-band, {1} detected".format(kcl35mask.sum(),
                                                  (kcl35mask & (~bad)).sum()))

    kcl35tbl = tbl[kcl35mask]
    kcl35freqs = u.Quantity(kcl35tbl['Frequency'], u.GHz)
    kkms_kcl35 = (2*np.pi*kcl35tbl['Fitted Width']**2)**0.5 * kcl35tbl['Fitted Amplitude K']
    ekkms_kcl35 = (2*np.pi)**0.5 * (kcl35tbl['Fitted Width error']**2 *
                                    kcl35tbl['Fitted Amplitude K']**2 +
                                    kcl35tbl['Fitted Width']**2 *
                                    kcl35tbl['Fitted Amplitude error K']**2)**0.5
    #kkms_kcl35 = (2*np.pi*(4)**2)**0.5 * kcl35tbl['Fitted Amplitude K']
    #Aul = u.Quantity(list(map(get_Aul_(frqs), kcl35freqs)), u.Hz)
    #deg = u.Quantity(list(map(get_deg_(frqs), kcl35freqs)))
    #kcl35tbl.add_column(table.Column(name='Aul', data=Aul))
    #kcl35tbl.add_column(table.Column(name='Degeneracy', data=deg))
    Aul = kcl35tbl['Aij']
    deg = kcl35tbl['deg']
    kcl35_nug = nupper_of_kkms(kkms=kkms_kcl35,
                               freq=kcl35freqs,
                               Aul=Aul,
                               #degeneracies=deg,
                               )
    kcl35_nu = kcl35_nug / deg
    ekcl35_nug = nupper_of_kkms(kkms=ekkms_kcl35,
                                freq=kcl35freqs,
                                Aul=Aul,
                                #degeneracies=deg
                               )
    ekcl35_nu = ekcl35_nug / deg
    kcl35tbl.add_column(table.Column(name='N_U', data=kcl35_nu,
                                     ))

    v0 = np.array(['v=0' in row['Species'] for row in kcl35tbl])
    v1 = np.array(['v=1' in row['Species'] for row in kcl35tbl])
    v2 = np.array(['v=2' in row['Species'] for row in kcl35tbl])
    v3 = np.array(['v=3' in row['Species'] for row in kcl35tbl])
    v4 = np.array(['v=4' in row['Species'] for row in kcl35tbl])
    v5 = np.array(['v=5' in row['Species'] for row in kcl35tbl])
    v6 = np.array(['v=6' in row['Species'] for row in kcl35tbl])
    v7 = np.array(['v=7' in row['Species'] for row in kcl35tbl])
    v8 = np.array(['v=8' in row['Species'] for row in kcl35tbl])
    j13 = np.array(['J=13' in row['Species'] for row in kcl35tbl])
    j29 = np.array(['J=29' in row['Species'] for row in kcl35tbl])
    j31 = np.array(['J=31' in row['Species'] for row in kcl35tbl])

    pl.figure(1).clf()
    print("KCl")

    optdepth = (dust_emissivity.dust.kappa(kcl35freqs, beta=2) * 120*u.g/u.cm**2).decompose().value
    kcl35_nug_dustcorr = nupper_of_kkms(kkms=kkms_kcl35 * np.exp(optdepth),
                                        freq=kcl35freqs,
                                        Aul=Aul,
                                        #degeneracies=deg,
                                        )
    kcl35_nu_dustcorr = kcl35_nug_dustcorr / deg

    tex0 = fit_tex(u.Quantity(kcl35tbl['EU_K'][v0], u.K), kcl35_nu[v0],
                   errors=ekcl35_nu[v0], plot=True, verbose=True,
                   #partition_func=partfunc,
                   partition_func=partfunc,
                   marker='o', color='r', label='v=0 ')
    #tex0_dustc = fit_tex(u.Quantity(kcl35tbl['EU_K'][v0], u.K), kcl35_nu_dustcorr[v0],
    #                     errors=ekcl35_nu[v0], plot=True, verbose=True,
    #                     partition_func=partfunc, marker='o', color='b', label='v=0 ')
    #tex0_dustc_nob3 = fit_tex(u.Quantity(kcl35tbl['EU_K'][v0][1:], u.K), kcl35_nu_dustcorr[v0][1:],
    #                          errors=ekcl35_nu[v0][1:], plot=True, verbose=True,
    #                     partition_func=partfunc, marker='o', color='g', label='v=0 no b3')
    #tex1 = fit_tex(u.Quantity(kcl35tbl['EU_K'][v1], u.K), kcl35_nu[v1],
    #               errors=ekcl35_nu[v1], plot=True, verbose=True,
    #               partition_func=partfunc, marker='s', color='b', label='v=1 ')
    #tex2 = fit_tex(u.Quantity(kcl35tbl['EU_K'][v2], u.K), kcl35_nu[v2],
    #               errors=ekcl35_nu[v2], plot=True, verbose=True,
    #               partition_func=partfunc, marker='^', color='g', label='v=2 ')
    #tex3 = fit_tex(u.Quantity(kcl35tbl['EU_K'][v3], u.K), kcl35_nu[v3],
    #               errors=ekcl35_nu[v3], plot=True, verbose=True,
    #               partition_func=partfunc, marker='v', color='orange', label='v=3 ')
    #tex4 = fit_tex(u.Quantity(kcl35tbl['EU_K'][v4], u.K), kcl35_nu[v4],
    #               errors=ekcl35_nu[v4], plot=True, verbose=True,
    #               partition_func=partfunc, marker='d', color='m', label='v=4 ')
    texJ13 = fit_tex(u.Quantity(kcl35tbl['EU_K'][j13], u.K), kcl35_nu[j13],
                     errors=ekcl35_nu[j13], plot=True, verbose=True,
                     partition_func=partfunc, marker='h', color='k', label='J=13 ')
    #texJ13_dustc = fit_tex(u.Quantity(kcl35tbl['EU_K'][j13], u.K), kcl35_nu_dustcorr[j13],
    #                 errors=ekcl35_nu[j13], plot=True, verbose=True,
    #                 partition_func=partfunc, marker='h', color='k', label='J=13 dust')
    texJ29 = fit_tex(u.Quantity(kcl35tbl['EU_K'][j29], u.K), kcl35_nu[j29],
                     errors=ekcl35_nu[j29], plot=True, verbose=True,
                     partition_func=partfunc, marker='o', color='b', label='J=29 ')
    texJ31 = fit_tex(u.Quantity(kcl35tbl['EU_K'][j31], u.K), kcl35_nu[j31],
                     errors=ekcl35_nu[j31], plot=True, verbose=True,
                     partition_func=partfunc, marker='P', color='navy', label='J=31 ')

    pl.legend(loc='lower right')
    pl.title("KCl")
    pl.axis((0, 2500, 7.75, 10.5))
    pl.savefig(paths.fpath("KCl_rotational_diagrams.pdf"))

    pl.figure(1).clf()
    vstate = 0*v0 + 1*v1 + 2*v2 + 3*v3 + 4*v4 + 5*v5 + 6*v6 + 7*v7 + 8*v8

    print(fit_multi_tex(eupper=u.Quantity(kcl35tbl['EU_K'], u.K),
                        nupperoverg=kcl35_nu,
                        vstate=vstate,
                        jstate=np.array(kcl35tbl['J$_u$'], dtype='int'),
                        vibenergies=get_vib_energies(salt_tables.KCl),
                        rotenergies=get_rot_energies(salt_tables.KCl),
                        errors=ekcl35_nu,
                        rottemlims=(30,150),
                        vibtemlims=(2500,5000),
                        collims=(np.log(1e8), np.log(1e16)),
                        plot=True, verbose=True, partition_func=partfunc, marker='o',
                        molname='KCl',
                        colors=('r','g','b','orange','m','c','darkred','darkgreen',
                                'purple'), )
         )
    pl.legend(loc='lower right')
    pl.title("KCl")
    pl.axis((0, 4000, 8.0, 10.5))

    pl.savefig(paths.fpath("KCl_rotational-vibrational_fit_diagrams.pdf"))

    #kcl37 = Vamdc.query_molecule('KCl-37')
    #rt37 = kcl37.data['RadiativeTransitions']
    #frqs = u.Quantity([(float(rt37[key].FrequencyValue)*u.MHz).to(u.GHz,
    #                                                              u.spectral())
    #                   for key in rt37])

    #upperStateRefs = [rt37[key].UpperStateRef for key in rt37]
    #degeneracies = [int(kcl37.data['States'][upperStateRef].TotalStatisticalWeight)
    #                for upperStateRef in upperStateRefs]
    #einsteinAij = u.Quantity([float(rt37[key].TransitionProbabilityA) for key in rt37], 1/u.s)

    frqs, einsteinAij, degeneracies, EU, partfunc = get_molecular_parameters('KCl-37, v=0-15', catalog='CDMS', fmin=85*u.GHz, fmax=360*u.GHz)

    tbl = table.Table.read(paths.tpath('fitted_stacked_lines.txt'), format='ascii.fixed_width')

    v0mask = np.array([(not hasattr(row['Species'], 'mask')) and ('v=0' in row['Species'])
                       for row in tbl])
    v1mask = np.array([(not hasattr(row['Species'], 'mask')) and ('v=1' in row['Species'])
                       for row in tbl])
    v2mask = np.array([(not hasattr(row['Species'], 'mask')) and ('v=2' in row['Species'])
                       for row in tbl])
    v3mask = np.array([(not hasattr(row['Species'], 'mask')) and ('v=3' in row['Species'])
                       for row in tbl])
    v4mask = np.array([(not hasattr(row['Species'], 'mask')) and ('v=4' in row['Species'])
                       for row in tbl])

    kcl37mask = np.array([(not hasattr(row['Species'], 'mask')) and
                          ('K37Cl' == row['Species'][:5] or
                           '39K-37Cl' in row['Species'])
                          for row in tbl])

    # mask out a bad fit
    bad = ((tbl['Line Name'] == 'K37Clv=1_29-28') |
           (tbl['Line Name'] == '41KClv=0 46-45'))

    print("K37Cl: {0} in-band, {1} detected".format(kcl37mask.sum(),
                                                    (kcl37mask & (~bad)).sum()))

    kcl37tbl = tbl[kcl37mask & (~bad)]
    kcl37freqs = u.Quantity(kcl37tbl['Frequency'], u.GHz)
    kkms_kcl37 = (2*np.pi*kcl37tbl['Fitted Width']**2)**0.5 * kcl37tbl['Fitted Amplitude K']
    ekkms_kcl37 = (2*np.pi)**0.5 * (kcl37tbl['Fitted Width error']**2 *
                                    kcl37tbl['Fitted Amplitude K']**2 +
                                    kcl37tbl['Fitted Width']**2 *
                                    kcl37tbl['Fitted Amplitude error K']**2)**0.5
    #Aul = u.Quantity(list(map(get_Aul_(frqs), kcl37freqs)), u.Hz)
    #deg = u.Quantity(list(map(get_deg_(frqs), kcl37freqs)))
    Aul = kcl37tbl['Aij']
    deg = kcl37tbl['deg']
    kcl37_nu = nupper_of_kkms(kkms=kkms_kcl37,
                              freq=kcl37freqs,
                              Aul=Aul)/deg
    ekcl37_nu = nupper_of_kkms(kkms=ekkms_kcl37,
                               freq=kcl37freqs,
                               Aul=Aul)/deg

    v0 = np.array(['v=0' in row['Species'] for row in kcl37tbl])
    v1 = np.array(['v=1' in row['Species'] for row in kcl37tbl])

    pl.figure(2).clf()
    print("K 37Cl")
    tex0 = fit_tex(u.Quantity(kcl37tbl['EU_K'][v0], u.K), kcl37_nu[v0],
                   errors=ekcl37_nu[v0], plot=True, verbose=True,
                   partition_func=partfunc, marker='o', color='r', label='v=0 ')
    tex1 = fit_tex(u.Quantity(kcl37tbl['EU_K'][v1], u.K), kcl37_nu[v1],
                   errors=ekcl37_nu[v1], plot=True, verbose=True,
                   partition_func=partfunc, marker='s', color='b', label='v=1 ')
    pl.legend(loc='lower right')
    pl.title("K$^{37}$Cl")
    pl.savefig(paths.fpath("KCl37_rotational_diagrams.pdf"))

    pl.figure(2).clf()
    vstate = 0*v0 + 1*v1

    print(fit_multi_tex(eupper=u.Quantity(kcl37tbl['EU_K'], u.K),
                        nupperoverg=kcl37_nu,
                        vstate=vstate,
                        jstate=np.array(kcl37tbl['J$_u$'], dtype='int'),
                        vibenergies=get_vib_energies(salt_tables.K37Cl),
                        rotenergies=get_rot_energies(salt_tables.K37Cl),
                        errors=ekcl37_nu,
                        plot=True, verbose=True, partition_func=partfunc, marker='o',
                        molname='K37Cl',
                        colors=('r','g','b','orange','m','c'), )
         )
    pl.legend(loc='lower right')
    pl.title("K$^{37}$Cl")

    pl.savefig(paths.fpath("KCl37_rotational-vibrational_fit_diagrams.pdf"))

    #nacl = Vamdc.query_molecule('NaCl$')
    #rt_nacl = nacl.data['RadiativeTransitions']
    #frqs = u.Quantity([(float(rt_nacl[key].FrequencyValue)*u.MHz).to(u.GHz,
    #                                                                 u.spectral())
    #                   for key in rt_nacl])

    #upperStateRefs = [rt_nacl[key].UpperStateRef for key in rt_nacl]
    #degeneracies = [int(nacl.data['States'][upperStateRef].TotalStatisticalWeight)
    #                for upperStateRef in upperStateRefs]
    #einsteinAij = u.Quantity([float(rt_nacl[key].TransitionProbabilityA) for key in rt_nacl], 1/u.s)

    frqs, einsteinAij, degeneracies, EU, partfunc = get_molecular_parameters('NaCl, v=0-15', catalog='CDMS', fmin=85*u.GHz, fmax=360*u.GHz)

    tbl = table.Table.read(paths.tpath('fitted_stacked_lines.txt'), format='ascii.fixed_width')

    naclmask = np.array([(not hasattr(row['Species'], 'mask')) and
                         ('NaCl' == row['Species'][:4] or
                         '23Na-35Cl' in row['Species'])
                         for row in tbl])

    # mask out a bad fit
    # on edge of absorption feature
    bad = (tbl['Line Name'] == 'NaClv=4_17-16')

    print("NaCl: {0} in-band, {1} detected".format(naclmask.sum(),
                                                   (naclmask & (~bad)).sum()))

    nacltbl = tbl[naclmask & (~bad)]
    naclfreqs = u.Quantity(nacltbl['Frequency'], u.GHz)
    kkms_nacl = (2*np.pi*nacltbl['Fitted Width']**2)**0.5 * nacltbl['Fitted Amplitude K']
    ekkms_nacl = (2*np.pi)**0.5 * (nacltbl['Fitted Width error']**2 *
                                   nacltbl['Fitted Amplitude K']**2 +
                                   nacltbl['Fitted Width']**2 *
                                   nacltbl['Fitted Amplitude error K']**2)**0.5
    #Aul = u.Quantity(list(map(get_Aul_(frqs), naclfreqs)), u.Hz)
    #deg = u.Quantity(list(map(get_deg_(frqs), naclfreqs)))
    Aul = nacltbl['Aij']
    deg = nacltbl['deg']
    nacl_nu = nupper_of_kkms(kkms=kkms_nacl,
                             freq=naclfreqs,
                             Aul=Aul,)/deg
    enacl_nu = nupper_of_kkms(kkms=ekkms_nacl,
                              freq=naclfreqs,
                              Aul=Aul)/deg

    v0 = np.array(['v=0' in row['Species'] for row in nacltbl])
    v1 = np.array(['v=1' in row['Species'] for row in nacltbl])
    v2 = np.array(['v=2' in row['Species'] for row in nacltbl])
    v3 = np.array(['v=3' in row['Species'] for row in nacltbl])
    v4 = np.array(['v=4' in row['Species'] for row in nacltbl])
    v5 = np.array(['v=5' in row['Species'] for row in nacltbl])
    v6 = np.array(['v=6' in row['Species'] for row in nacltbl])
    v7 = np.array(['v=7' in row['Species'] for row in nacltbl])
    v8 = np.array(['v=8' in row['Species'] for row in nacltbl])
    j7 = np.array(['J=7' in row['Species'] for row in nacltbl])
    j8 = np.array(['J=8' in row['Species'] for row in nacltbl])

    pl.figure(3).clf()
    print("NaCl")
    #tex0 = fit_tex(u.Quantity(nacltbl['EU_K'][v0], u.K), nacl_nu[v0], plot=True,
    #               verbose=True, partition_func=partfunc, marker='o', color='r')
    tex1 = fit_tex(u.Quantity(nacltbl['EU_K'][v1], u.K), nacl_nu[v1],
                   errors=enacl_nu[v1], plot=True,
                   verbose=True, partition_func=partfunc, marker='s', color='b', label='v=1 ')
    tex2 = fit_tex(u.Quantity(nacltbl['EU_K'][v2], u.K), nacl_nu[v2],
                   errors=enacl_nu[v2], plot=True,
                   verbose=True, partition_func=partfunc, marker='^', color='g', label='v=2 ')
    #tex3 = fit_tex(u.Quantity(nacltbl['EU_K'][v3], u.K), nacl_nu[v3],
    #               errors=enacl_nu[v3], plot=True,
    #               verbose=True, partition_func=partfunc, marker='o', color='r', label='v=3 ')
    #tex4 = fit_tex(u.Quantity(nacltbl['EU_K'][v4], u.K), nacl_nu[v4],
    #               errors=enacl_nu[v4], plot=True,
    #               verbose=True, partition_func=partfunc, marker='d', color='orange', label='v=4 ')
    #tex5 = fit_tex(u.Quantity(nacltbl['EU_K'][v5], u.K), nacl_nu[v5],
    #               errors=enacl_nu[v5], plot=True,
    #               verbose=True, partition_func=partfunc, marker='v', color='m', label='v=5 ')
    #tex6 = fit_tex(u.Quantity(nacltbl['EU_K'][v6], u.K), nacl_nu[v6],
    #               errors=enacl_nu[v6], plot=True,
    #               verbose=True, partition_func=partfunc, marker='v', color='m', label='v=6 ')
    texJ8 = fit_tex(u.Quantity(nacltbl['EU_K'][j8], u.K), nacl_nu[j8],
                    errors=enacl_nu[j8], plot=True, verbose=True,
                    partition_func=partfunc, marker='h', color='k', label='J=8 ')
    texJ7 = fit_tex(u.Quantity(nacltbl['EU_K'][j7], u.K), nacl_nu[j7],
                    errors=enacl_nu[j7], plot=True, verbose=True,
                    partition_func=partfunc, marker='>', color='c', label='J=7 ')
    pl.legend(loc='upper right')
    pl.axis([300,3300,9.0,13])
    pl.title("NaCl")
    pl.savefig(paths.fpath("NaCl_rotational_diagrams.pdf"))

    pl.figure(3).clf()
    vstate = 0*v0 + 1*v1 + 2*v2 + 3*v3 + 4*v4 + 5*v5 + 6*v6 + 7*v7 + 8*v8

    print(fit_multi_tex(eupper=u.Quantity(nacltbl['EU_K'], u.K),
                        nupperoverg=nacl_nu,
                        vstate=vstate,
                        jstate=np.array(nacltbl['J$_u$'], dtype='int'),
                        vibenergies=get_vib_energies(salt_tables.NaCl),
                        rotenergies=get_rot_energies(salt_tables.NaCl),
                        errors=enacl_nu,
                        rottemlims=(20,175),
                        vibtemlims=(500,8000),
                        collims=(np.log(1e10), np.log(1e16)),
                        plot=True, verbose=True, partition_func=partfunc, marker='o',
                        molname='NaCl',
                        colors=('r','g','b','orange','m','c','darkred','darkgreen',
                                'purple'), )
         )
    pl.legend(loc='upper right')
    pl.axis([300,4000,8.8,10.8])
    pl.title("NaCl")

    pl.savefig(paths.fpath("NaCl_rotational-vibrational_fit_diagrams.pdf"))

    #nacl37 = Vamdc.query_molecule('NaCl-37')
    #rt_nacl37 = nacl37.data['RadiativeTransitions']
    #frqs = u.Quantity([(float(rt_nacl37[key].FrequencyValue)*u.MHz).to(u.GHz,
    #                                                                   u.spectral())
    #                   for key in rt_nacl37])

    #upperStateRefs = [rt_nacl37[key].UpperStateRef for key in rt_nacl37]
    #degeneracies = [int(nacl37.data['States'][upperStateRef].TotalStatisticalWeight)
    #                for upperStateRef in upperStateRefs]
    #einsteinAij = u.Quantity([float(rt_nacl37[key].TransitionProbabilityA) for key in rt_nacl37], 1/u.s)

    frqs, einsteinAij, degeneracies, EU, partfunc = get_molecular_parameters('NaCl-37, v=0-15', catalog='CDMS', fmin=85*u.GHz, fmax=360*u.GHz)

    tbl = table.Table.read(paths.tpath('fitted_stacked_lines.txt'), format='ascii.fixed_width')

    nacl37mask = np.array([(not hasattr(row['Species'], 'mask')) and
                          ('Na37Cl' == row['Species'][:6] or
                           '23Na-37Cl' in row['Species'])
                          for row in tbl])

    bad = (tbl['Line Name'] == 'Na37Clv=6_8-7') # tabulated velocity too far off

    print("Na37Cl: {0} in-band, {1} detected".format(nacl37mask.sum(),
                                                     (nacl37mask & (~bad)).sum()))

    nacl37tbl = tbl[nacl37mask & (~bad)]
    nacl37freqs = u.Quantity(nacl37tbl['Frequency'], u.GHz)
    kkms_nacl37 = (2*np.pi*nacl37tbl['Fitted Width']**2)**0.5 * nacl37tbl['Fitted Amplitude K']
    ekkms_nacl37 = (2*np.pi)**0.5 * (nacl37tbl['Fitted Width error']**2 *
                                     nacl37tbl['Fitted Amplitude K']**2 +
                                     nacl37tbl['Fitted Width']**2 *
                                     nacl37tbl['Fitted Amplitude error K']**2)**0.5
    #Aul = u.Quantity(list(map(get_Aul_(frqs), nacl37freqs)), u.Hz)
    #deg = u.Quantity(list(map(get_deg_(frqs), nacl37freqs)))
    Aul = nacl37tbl['Aij']
    deg = nacl37tbl['deg']
    nacl37_nu = nupper_of_kkms(kkms=kkms_nacl37,
                               freq=nacl37freqs,
                               Aul=Aul)/deg
    enacl37_nu = nupper_of_kkms(kkms=ekkms_nacl37,
                                freq=nacl37freqs,
                                Aul=Aul)/deg

    v0 = np.array(['v=0' in row['Species'] for row in nacl37tbl])
    v1 = np.array(['v=1' in row['Species'] for row in nacl37tbl])
    v2 = np.array(['v=2' in row['Species'] for row in nacl37tbl])
    v3 = np.array(['v=3' in row['Species'] for row in nacl37tbl])
    v4 = np.array(['v=4' in row['Species'] for row in nacl37tbl])
    v5 = np.array(['v=5' in row['Species'] for row in nacl37tbl])
    v6 = np.array(['v=6' in row['Species'] for row in nacl37tbl])
    j7 = np.array(['J=7' in row['Species'] for row in nacl37tbl])
    j8 = np.array(['J=8' in row['Species'] for row in nacl37tbl])

    pl.figure(4).clf()
    print("Na37Cl")
    tex0 = fit_tex(u.Quantity(nacl37tbl['EU_K'][v0], u.K), nacl37_nu[v0],
                   errors=enacl37_nu[v0], plot=True,
                   verbose=True, partition_func=partfunc, marker='o', color='r')
    tex1 = fit_tex(u.Quantity(nacl37tbl['EU_K'][v1], u.K), nacl37_nu[v1],
                   errors=enacl37_nu[v1], plot=True,
                   verbose=True, partition_func=partfunc, marker='s', color='b', label='v=1 ')
    tex2 = fit_tex(u.Quantity(nacl37tbl['EU_K'][v2], u.K), nacl37_nu[v2], plot=True,
                   verbose=True, partition_func=partfunc, marker='^', color='g', label='v=2 ')
    #tex3 = fit_tex(u.Quantity(nacl37tbl['EU_K'][v3], u.K), nacl37_nu[v3], plot=True,
    #               verbose=True, partition_func=partfunc, marker='o', color='r', label='v=3 ')
    #tex4 = fit_tex(u.Quantity(nacl37tbl['EU_K'][v4], u.K), nacl37_nu[v4], plot=True,
    #               verbose=True, partition_func=partfunc, marker='d', color='orange', label='v=4 ')
    #tex5 = fit_tex(u.Quantity(nacl37tbl['EU_K'][v5], u.K), nacl37_nu[v5],
    #               errors=enacl37_nu[v5], plot=True,
    #               verbose=True, partition_func=partfunc, marker='v', color='m', label='v=5 ')
    texJ8 = fit_tex(u.Quantity(nacl37tbl['EU_K'][j8], u.K), nacl37_nu[j8],
                    errors=enacl37_nu[j8], plot=True, verbose=True,
                    partition_func=partfunc, marker='h', color='k', label='J=8 ')
    texJ7 = fit_tex(u.Quantity(nacl37tbl['EU_K'][j7], u.K), nacl37_nu[j7],
                    errors=enacl37_nu[j7], plot=True, verbose=True,
                    partition_func=partfunc, marker='>', color='c', label='J=7 ')

    pl.legend(loc='lower right')
    pl.title("Na$^{37}$Cl")
    pl.axis([-50,3000,9.0,12.25])
    pl.savefig(paths.fpath("NaCl37_rotational_diagrams.pdf"))

    pl.figure(4).clf()
    vstate = 0*v0 + 1*v1 + 2*v2 + 3*v3 + 4*v4 + 5*v5 + 6*v6

    print(fit_multi_tex(eupper=u.Quantity(nacl37tbl['EU_K'], u.K),
                        nupperoverg=nacl37_nu,
                        vstate=vstate,
                        jstate=np.array(nacl37tbl['J$_u$'], dtype='int'),
                        vibenergies=get_vib_energies(salt_tables.Na37Cl),
                        rotenergies=get_rot_energies(salt_tables.Na37Cl),
                        errors=enacl37_nu,
                        rottemlims=(20,225),
                        plot=True, verbose=True, partition_func=partfunc, marker='o',
                        molname='Na37Cl',
                        colors=('r','g','b','orange','m','c'), )
         )
    pl.legend(loc='lower right')
    pl.title("Na$^{37}$Cl")
    pl.axis((-50,3700,8,11.0))

    pl.savefig(paths.fpath("NaCl37_rotational-vibrational_fit_diagrams.pdf"))

    #k41cl = Vamdc.query_molecule('K-41-Cl')
    #rt41 = k41cl.data['RadiativeTransitions']
    #frqs = u.Quantity([(float(rt41[key].FrequencyValue)*u.MHz).to(u.GHz,
    #                                                              u.spectral())
    #                   for key in rt41])

    #upperStateRefs = [rt41[key].UpperStateRef for key in rt41]
    #degeneracies = [int(k41cl.data['States'][upperStateRef].TotalStatisticalWeight)
    #                for upperStateRef in upperStateRefs]
    #einsteinAij = u.Quantity([float(rt41[key].TransitionProbabilityA) for key in rt41], 1/u.s)

    frqs, einsteinAij, degeneracies, EU, partfunc = get_molecular_parameters('K-41-Cl, v=0-15', catalog='CDMS', fmin=85*u.GHz, fmax=360*u.GHz)

    tbl = table.Table.read(paths.tpath('fitted_stacked_lines.txt'), format='ascii.fixed_width')

    k41clmask = np.array([(not hasattr(row['Species'], 'mask')) and
                          ('41KCl' == row['Species'][:5] or
                           '41K-35Cl' in row['Species'])
                          for row in tbl])

    bad = (((tbl['Species'] == '41KClv=2') & (tbl['QNs'] == '29-28')) | # absorption
           ((tbl['Species'] == '41KClv=0') & (tbl['QNs'] == '31-30')) | # wing of NaCl
           ((tbl['Species'] == '41KClv=0') & (tbl['QNs'] == '46-45')) | # absorption
           (tbl['Line Name'] == 'NaClv=4_7-6') |
           (tbl['Line Name'] == '41KClv=2_29-28') |
           (tbl['Line Name'] == '41KClv=2 29-28') |
           (tbl['Line Name'] == '41KClv=0_31-30') |
           (tbl['Line Name'] == '41KClv=0 31-30') |
           (tbl['Line Name'] == '41KClv=0_46-45')
          )

    print("41KCl: {0} in-band, {1} detected".format(k41clmask.sum(),
                                                    (k41clmask & (~bad)).sum()))

    k41cltbl = tbl[k41clmask & (~bad)]
    k41clfreqs = u.Quantity(k41cltbl['Frequency'], u.GHz)
    kkms_k41cl = (2*np.pi*k41cltbl['Fitted Width']**2)**0.5 * k41cltbl['Fitted Amplitude K']
    ekkms_k41cl = (2*np.pi)**0.5 * (k41cltbl['Fitted Width error']**2 *
                                    k41cltbl['Fitted Amplitude K']**2 +
                                    k41cltbl['Fitted Width']**2 *
                                    k41cltbl['Fitted Amplitude error K']**2)**0.5
    #Aul = u.Quantity(list(map(get_Aul_(frqs), k41clfreqs)), u.Hz)
    #deg = u.Quantity(list(map(get_deg_(frqs), k41clfreqs)))
    Aul = k41cltbl['Aij']
    deg = k41cltbl['deg']
    k41cl_nu = nupper_of_kkms(kkms=kkms_k41cl,
                              freq=k41clfreqs,
                              Aul=Aul)/deg
    ek41cl_nu = nupper_of_kkms(kkms=ekkms_k41cl,
                               freq=k41clfreqs,
                               Aul=Aul)/deg

    v0 = np.array(['v=0' in row['Species'] for row in k41cltbl])
    v1 = np.array(['v=1' in row['Species'] for row in k41cltbl])
    v2 = np.array(['v=2' in row['Species'] for row in k41cltbl])
    v3 = np.array(['v=3' in row['Species'] for row in k41cltbl])

    pl.figure(5).clf()
    print("41KCl")
    tex0 = fit_tex(u.Quantity(k41cltbl['EU_K'][v0], u.K), k41cl_nu[v0],
                   errors=ek41cl_nu[v0], plot=True,
                   verbose=True, partition_func=partfunc, marker='o', color='r', label='v=0 ')
    #tex1 = fit_tex(u.Quantity(k41cltbl['EU_K'][v1], u.K), k41cl_nu[v1],
    #               errors=ek41cl_nu[v1], plot=True,
    #               verbose=True, partition_func=partfunc, marker='s', color='b', label='v=1 ')
    pl.legend(loc='lower right')
    pl.title("$^{41}$KCl")
    pl.savefig(paths.fpath("K41Cl_rotational_diagrams.pdf"))

    pl.figure(5).clf()
    vstate = 0*v0 + 1*v1 + 2*v2 + 3*v3
    print(fit_multi_tex(eupper=u.Quantity(k41cltbl['EU_K'], u.K),
                        nupperoverg=k41cl_nu,
                        vstate=vstate,
                        jstate=np.array(k41cltbl['J$_u$'], dtype='int'),
                        vibenergies=get_vib_energies(salt_tables.K41Cl),
                        rotenergies=get_rot_energies(salt_tables.K41Cl),
                        errors=ek41cl_nu,
                        plot=True, verbose=True, partition_func=partfunc, marker='o',
                        molname='41KCl',
                        colors=('r','g','b','orange','m','c',), )
         )
    pl.legend(loc='lower right')
    pl.title("$^{41}$KCl")

    pl.savefig(paths.fpath("41KCl_rotational-vibrational_fit_diagrams.pdf"))


#Question 1

def Affiche(c):
	print("X = ", str(c[0]).replace("'",''))
	print("G = ", str(c[1]).replace("'",''))
	print("Y = ", str(c[2]).replace("'",''),end='\n\n')

def Lecture(nomfichier):
	fichier = open(nomfichier,"r")
	liste = fichier.readlines()
	print(liste)
	#Ensemble de départ est lettres
	X=set()
	#Ensemble d'arrivée est numéros
	Y=set()
	#Correspondance est ?
	G=()
	#On transforme la liste en une chaine pour pouvoir la parcourir avec la commande rstrip
	chaine=str(liste)
	couple=chaine.rstrip().split("\n",">")

	i=0
	while i < len(liste):
		if chaine[i].isalpha() :
			X.add(chaine[i])
		elif chaine[i].isdigit() :
			Y.add(chaine[i])
		i+=1
	#Triplet (X,G,Y)
	c=(X,G,Y)

	return c

nomfichier=input("Saisisez le nom du fichier : ")
print(Affiche(Lecture(nomfichier)))


import requests
import sys
import os

from PyQt5.QtGui import QPixmap
from PyQt5.QtWidgets import QApplication, QWidget
from PyQt5.QtCore import Qt
from UI import Ui_Form

SCREEN_SIZE = [600, 450]

class Map(QWidget, Ui_Form):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.setupUi(self)

        self.show_map_button.clicked.connect(self.show_map)

    # Вывод изображения карты в зависимости от введённых координат
    def show_map(self):
        if self.sender():
            self.get_coords()
        self.getImage()

        self.pixmap = QPixmap(self.map_file)
        self.map_label.setPixmap(self.pixmap)

    # Получение введённых пользователем координат и масштаба
    def get_coords(self):
        self.width = self.width_edit.text().strip()
        self.longitude = self.longitude_edit.text().strip()
        self.scale = int(self.scale_edit.text().strip())

    # Получение изображения по запросу
    def getImage(self):
        map_request = f"http://static-maps.yandex.ru/1.x/?ll" \
                      f"={self.longitude},{self.width}8&spn" \
                      f"={1 / (self.scale * 100)},{1 / (self.scale * 100)}&l=map"
        response = requests.get(map_request)
        if not response:
            print("Ошибка выполнения запроса:")
            print(map_request)
            print("Http статус:", response.status_code, "(", response.reason, ")")
            sys.exit(1)

        # Запишем полученное изображение в файл.
        self.map_file = "map.png"
        with open(self.map_file, "wb") as file:
            file.write(response.content)

    # Обработка клавиатуры
    def keyPressEvent(self, event):
        if event.key() == Qt.Key_PageUp:
            if self.scale < 7:
                self.scale += 1
                self.show_map()
        if event.key() == Qt.Key_PageDown:
            if self.scale > 1:
                self.scale -= 1
                self.show_map()

    def closeEvent(self, event):
        """При закрытии формы подчищаем за собой"""
        os.remove(self.map_file)

def except_hook(cls, exception, traceback):
    sys.__excepthook__(cls, exception, traceback)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    map = Map()
    map.show()
    sys.excepthook = except_hook
    sys.exit(app.exec())


"""
Test density estimation techniques
"""
import pytest
import numpy as np
from numpy.testing import assert_allclose
from scipy.stats import norm
from astroML.density_estimation import KNeighborsDensity, GaussianMixture1D

classifiers = [KNeighborsDensity(method='simple', n_neighbors=250),
               KNeighborsDensity(method='bayesian', n_neighbors=250)]

@pytest.mark.parametrize("clf", classifiers)
def test_1D_density(clf, atol=100):
    np.random.seed(0)
    dist = norm(0, 1)

    X = dist.rvs((5000, 1))
    X2 = np.linspace(-5, 5, 10).reshape((10, 1))
    true_dens = dist.pdf(X2[:, 0]) * X.shape[0]

    clf.fit(X)
    dens = clf.eval(X2)

    assert_allclose(dens, true_dens, atol=atol)

def test_gaussian1d():
    x = np.linspace(-6, 10, 1000)
    means = np.array([-1.5, 0.0, 2.3])
    sigmas = np.array([1, 0.25, 3.8])
    weights = np.array([1, 1, 1])

    gauss = GaussianMixture1D(means=means, sigmas=sigmas, weights=weights)
    y = gauss.pdf(x)

    # Check whether sampling works
    gauss.sample(10)

    dx = x[1] - x[0]
    integral = np.sum(y*dx)

    assert_allclose(integral, 1., atol=0.02)


# -*- coding: utf-8 -*-
from linebot.models import (
    CarouselTemplate, TemplateSendMessage, CarouselColumn,
    PostbackTemplateAction, MessageTemplateAction, URITemplateAction, BubbleContainer, BoxComponent, ButtonComponent,
    ImageComponent, MessageAction, URIAction, CarouselContainer, FlexSendMessage
)

from Bot.Youtube_API import get_videos, Order

def get_carousel_template_message(video_sort=Order, carousel_column_amount=5):
    if carousel_column_amount > 9:
        return "過多的 columns (必須小於10)"
    columns = []
    for video in get_videos("UC-sM_PLqzgktdUcW2LEKKkQ", video_sort, carousel_column_amount):
        columns.append(CarouselColumn(
            thumbnail_image_url=video,
            title='標題',
            text='敘述',
            actions=[
                PostbackTemplateAction(
                    label='postback1',
                    text='postback text1',
                    data='action=buy&itemid=1'
                ),
                MessageTemplateAction(
                    label='message1',
                    text='message text1'
                ),
                URITemplateAction(
                    label='uri1',
                    uri='http://example.com/1'
                )
            ]
        ))
    columns.append(get_video_menu())

    return TemplateSendMessage(
        alt_text='Carousel template',
        template=CarouselTemplate(
            columns=columns
        )
    )

def get_video_menu():
    return BubbleContainer(
        direction='ltr',
        #hero=ImageComponent(
        #    url='https://previews.123rf.com/images/galinadarla/galinadarla1610/galinadarla161000124/67174940-web-site-menu-icon-flat-style-made-in-vector-black-and-white-icon.jpg',
        #    size='full',
        #    aspectRatio='16:9',
        #    aspectMode='cover'
        #),
        body=BoxComponent(
            layout='vertical',
            contents=[
                ButtonComponent(
                    action=MessageAction(label='最新上傳', text='最新上傳'),
                    style='primary',
                    color='#29D90C',
                    margin='lg'
                ),
                ButtonComponent(
                    action=MessageAction(label='熱門影片', text='熱門影片'),
                    style='primary',
                    color='#29D90C',
                    margin='lg'
                ),
                ButtonComponent(
                    action=MessageAction(label='最多喜歡', text='最多喜歡'),
                    style='primary',
                    color='#29D90C',
                    margin='lg'
                ),
                ButtonComponent(
                    action=MessageAction(label='相關影片', text='相關影片'),
                    style='primary',
                    color='#29D90C',
                    margin='lg'
                ),
                ButtonComponent(
                    action=MessageAction(label='標題順序', text='標題順序'),
                    style='primary',
                    color='#29D90C',
                    margin='lg'
                ),
                ButtonComponent(
                    action=MessageAction(label='影片上傳次數', text='影片上傳次數'),
                    style='primary',
                    color='#29D90C',
                    margin='lg'
                )
            ]
        )
    )

def create_flex_template(order=Order.DATE, amount=5):
    thumbnails_url, video_url = get_videos('UC-sM_PLqzgktdUcW2LEKKkQ', order, amount)
    bubble_containers = []
    for i in range(len(video_url)):
        bubble_containers.append(create_bubble_container(thumbnails_url[i], video_url[i]))
    bubble_containers.append(get_video_menu())
    return FlexSendMessage(
        alt_text='test',
        contents=CarouselContainer(
            contents=bubble_containers
        )
    )

def create_bubble_container(thumbnail_image_url, video_url):
    return BubbleContainer(
        direction='ltr',
        hero=ImageComponent(
            url=thumbnail_image_url,
            size='full',
            aspectRatio='16:9',
            aspectMode='cover'
        ),
        body=BoxComponent(
            layout='vertical',
            contents=[
                ButtonComponent(
                    action=URIAction(label='觀看影片', uri=video_url),
                    style='primary',
                    color='#29D90C',
                    margin='lg'
                ),
                ButtonComponent(
                    action=URIAction(label='分享影片', uri='https://line.me/R/msg/text/?' + video_url),
                    style='primary',
                    color='#29D90C',
                    margin='lg'
                )
            ]
        )
    )


# -*- coding: utf-8 -*-
from datetime import datetime

from Acquisition import aq_base

from pytz import timezone
from seantis.reservation.tests import IntegrationTestCase
from seantis.reservation import utils
from seantis.reservation import exports
from seantis.reservation.export import ExportView, prepare_record

class TestExports(IntegrationTestCase):

    def test_reservations_export(self):
        self.login_manager()

        resource = self.create_resource()
        sc = resource.scheduler()

        start = datetime(2012, 2, 1, 12, 0, tzinfo=timezone('UTC'))
        end = datetime(2012, 2, 1, 16, 0, tzinfo=timezone('UTC'))
        some_date = datetime(2014, 1, 30, 13, 37, tzinfo=timezone('UTC'))
        dates = (start, end)

        sc.allocate(dates, approve_manually=False, quota=2)[0]

        token1 = sc.reserve(
            u'a@example.com', dates,
            data=utils.mock_data_dictionary(
                {
                    'stop': u'hammertime!',
                    'bust': u'a move!',
                    'when': some_date
                }
            )
        )

        token2 = sc.reserve(
            u'b@example.com', dates,
            data=utils.mock_data_dictionary(
                {
                    'never': u'gonna',
                    'give': u'you up',
                    'when': some_date
                }
            )
        )

        dataset = exports.reservations.dataset(
            {resource.uuid(): resource}, 'en', 'all', 'all'
        )

        self.assertEqual(len(dataset), 2)

        existing_columns = 11
        token_1_unique = 2
        token_2_unique = 2
        token_common = 1

        self.assertEqual(
            len(dataset.headers),
            existing_columns + token_1_unique + token_2_unique + token_common
        )

        self.assertEqual(dataset.dict[0]['Token'], utils.string_uuid(token1))
        self.assertEqual(dataset.dict[0]['Mocktest.stop'], u'hammertime!')
        self.assertEqual(dataset.dict[0]['Mocktest.bust'], u'a move!')
        self.assertEqual(dataset.dict[0]['Mocktest.never'], None)
        self.assertEqual(dataset.dict[0]['Mocktest.give'], None)

        self.assertEqual(dataset.dict[1]['Token'], utils.string_uuid(token2))
        self.assertEqual(dataset.dict[1]['Mocktest.stop'], None)
        self.assertEqual(dataset.dict[1]['Mocktest.bust'], None)
        self.assertEqual(dataset.dict[1]['Mocktest.never'], u'gonna')
        self.assertEqual(dataset.dict[1]['Mocktest.give'], u'you up')

        self.assertEqual(dataset.dict[0]['Start'], start)
        self.assertEqual(dataset.dict[1]['Start'], start)
        self.assertEqual(dataset.dict[0]['End'], end)
        self.assertEqual(dataset.dict[1]['End'], end)

        self.assertEqual(dataset.dict[0]['Mocktest.when'], some_date)
        self.assertEqual(dataset.dict[1]['Mocktest.when'], some_date)

        # just make sure these don't raise exceptions
        for format in ('xls', 'xlsx', 'json', 'csv'):
            transform_record = lambda r: prepare_record(r, format)
            dataset = exports.reservations.dataset(
                {resource.uuid(): resource},
                'en', 'all', 'all', transform_record
            )
            getattr(dataset, format)

    def test_reservations_export_date_filter(self):
        self.login_manager()

        resource = self.create_resource()
        sc = resource.scheduler()

        this_year = (
            datetime(2013, 1, 1, 12, 0), datetime(2013, 1, 1, 13, 0)
        )

        next_year = (
            datetime(2014, 1, 1, 12, 0), datetime(2014, 1, 1, 13, 0)
        )

        reservation_email = u'test@example.com'
        sc.allocate(this_year, approve_manually=False)
        sc.allocate(next_year, approve_manually=False)

        sc.reserve(reservation_email, this_year)
        sc.reserve(reservation_email, next_year)

        dataset = exports.reservations.dataset(
            {resource.uuid(): resource}, 'en', year='all', month='all'
        )
        self.assertEqual(len(dataset), 2)

        dataset = exports.reservations.dataset(
            {resource.uuid(): resource}, 'en', year='all', month='1'
        )
        self.assertEqual(len(dataset), 2)

        dataset = exports.reservations.dataset(
            {resource.uuid(): resource}, 'en', year='2013', month='all'
        )
        self.assertEqual(len(dataset), 1)

        dataset = exports.reservations.dataset(
            {resource.uuid(): resource}, 'en', year='2010', month='all'
        )
        self.assertEqual(len(dataset), 0)

    def test_reservations_export_title(self):
        self.login_manager()

        resource = self.create_resource()
        sc = resource.scheduler()

        start = datetime(2012, 2, 1, 12, 0)
        end = datetime(2012, 2, 1, 16, 0)
        dates = (start, end)

        reservation_email = u'test@example.com'
        sc.allocate(dates, approve_manually=False, quota=2)[0]

        sc.reserve(reservation_email, dates)

        # with title
        resource.aq_inner.aq_parent.title = 'testtitel'

        dataset = exports.reservations.dataset(
            {resource.uuid(): resource}, 'en', 'all', 'all'
        )

        self.assertEqual(len(dataset), 1)
        self.assertEqual(dataset[0][0], 'testtitel')

        # without title
        resource = aq_base(resource)

        dataset = exports.reservations.dataset(
            {resource.uuid(): resource}, 'en', 'all', 'all'
        )

        self.assertEqual(len(dataset), 1)
        self.assertEqual(dataset[0][0], None)

    def test_export_view(self):

        class MyExportView(ExportView):
            pass

        view = MyExportView(self.portal, self.request())

        self.assertRaises(NotImplementedError, lambda: view.content_type)
        self.assertRaises(NotImplementedError, lambda: view.file_extension)

        view.request['source'] = 'reservations'
        self.assertTrue(hasattr(view.source, '__call__'))

        view.request['source'] = 'united-reservations'
        self.assertTrue(hasattr(view.source, '__call__'))

        view.request['source'] = 'nonexistant'
        self.assertRaises(NotImplementedError, lambda: view.source)

        class MyJsonExportView(ExportView):
            content_type = 'application/json'
            file_extension = 'json'

        view = MyJsonExportView(self.portal, self.request())

        view.request['source'] = 'reservations'
        self.assertEqual('[]', view.render())

        response = view.request.RESPONSE
        self.assertEqual(response.headers['content-length'], '2')
        self.assertEqual(
            response.headers['content-disposition'],
            'filename="Plone site.Reservations (Normal).json"'
        )
        self.assertEqual(
            response.headers['content-type'], 'application/json;charset=utf-8'
        )


from Tkinter import *

def create_blob(w,bbox,fill):
  x1,y1,x2,y2 = bbox
  h = y2 - y1
  h2 = h / 2
  items = (
    # Fill in bg colour
    w.create_arc(x1-h2,y1,x1+h2,y2,style=PIESLICE,start=90,extent=180,fill=fill,outline=''),
    w.create_arc(x2-h2,y1,x2+h2,y2,style=PIESLICE,start=270,extent=180,fill=fill,outline=''),
    w.create_rectangle(x1,y1,x2,y2,fill=fill,outline=''),
    # Add outlines
    w.create_arc(x1-h2,y1,x1+h2,y2,style=ARC,start=90,extent=180),
    w.create_arc(x2-h2,y1,x2+h2,y2,style=ARC,start=270,extent=180),
    w.create_line(x1,y1,x2,y1),
    w.create_line(x1,y2,x2,y2)
  )
  return items

def create_text_blob(w,pos,text,fill):
  text = w.create_text(pos,text=text)
  items = w.create_blob(w.bbox(text),fill=fill)
  w.tag_raise(text)
  return (text, items)

Canvas.create_blob = create_blob
Canvas.create_text_blob = create_text_blob


import json

from django.utils import unittest

from iml_common.lib.name_value_list import NameValueList

class TestNameValueList(unittest.TestCase):
    def setUp(self):
        self.raw_values = [
            {"resolve": False},
            {"ping": True},
            {"auth": False},
            {"hostname_valid": False},
            {"fqdn_resolves": False},
            {"fqdn_matches": False},
            {"reverse_resolve": False},
            {"reverse_ping": False},
            {"yum_valid_repos": False},
            {"yum_can_update": False},
            {"openssl": False},
        ]

        self.values = NameValueList(self.raw_values)

    def test_creation_correction(self):
        for value, raw_value in zip(self.values, self.raw_values):
            self.assertEqual(value.name, raw_value.keys()[0])
            self.assertEqual(value.value, raw_value.values()[0])

    def test_name_value_create(self):
        self.values["execute"] = True
        self.assertEqual(self.values["execute"], True)
        pass

    def test_update(self):
        self.values["resolve"] = True
        self.assertEqual(self.values["resolve"], True)

    def test_type_change(self):
        self.values["reverse_resolve"] = "string"
        self.assertEqual(self.values["reverse_resolve"], "string")

    def test_index_constant(self):
        self.values["reverse_resolve"] = True
        self.assertEqual(self.values.keys().index("reverse_resolve"), 6)

    def test_collection_via_json(self):
        json_string = json.dumps(self.values.collection())
        loaded_values = NameValueList(json.loads(json_string))
        self.assertEqual(self.values, loaded_values)

    def test_collection_len(self):
        self.assertEqual(len(self.values), len(self.values.collection()))

    def test_collection_values(self):
        for entry in self.values.collection():
            self.assertEqual(entry["value"], self.values[entry["name"]])


from app import app
from app.models.emotion_predictor import predictor, response, emotion_color
import numpy as np
from flask import request, render_template

@app.route("/",methods=['POST'])
def predict_emo():
    sentence = [x for x in request.form.values()][0]
    quoted = '"' + sentence + '"'
    opening = "Your emotion is dominated by..."
    prediction = predictor(sentence)
    index = np.random.randint(0,3)
    reply = response(prediction,index)
    color_config = emotion_color(prediction)
    return render_template('index.html', sentence_text=quoted, opening_text=opening, prediction_text= prediction, reply_text = reply, color_config = color_config )


from django.contrib.auth import get_user_model
from django.db import models

class Group(models.Model):
    title = models.CharField(
        verbose_name='Заголовок',
        max_length=200, help_text='Кратко опишите заголовок'
    )
    slug = models.SlugField(
        verbose_name='Адрес для страницы с задачей',
        max_length=200, unique=True,
        help_text='Укажите адрес для страницы задачи'
    )
    description = models.TextField(
        verbose_name='Описание группы',
        null=True, blank=True,
        help_text='Краткое описание группы'
    )

    def __str__(self):
        return self.title

User = get_user_model()

class Post(models.Model):
    author = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        related_name='posts'
    )
    group = models.ForeignKey(
        Group,
        blank=True, null=True,
        on_delete=models.SET_NULL,
        related_name='posts'
    )
    text = models.TextField(verbose_name='Текст', help_text='Напишите статью')
    image = models.ImageField(
        verbose_name='Картинка',
        upload_to='posts/',
        blank=True,
    )
    pub_date = models.DateTimeField(auto_now_add=True)

    class Meta:
        ordering = ['-pub_date']

    def __str__(self):
        return self.text[:15]

class Comment(models.Model):
    post = models.ForeignKey(
        Post,
        verbose_name='Комментируемый пост.',
        blank=True, null=True,
        on_delete=models.SET_NULL,
        related_name='comments'
    )
    author = models.ForeignKey(
        User,
        blank=True, null=True,
        verbose_name='Автор комментария.',
        on_delete=models.CASCADE,
        related_name='comments',
    )
    text = models.TextField(
        verbose_name='Текст комментария.',
        max_length=2000,
        help_text='Введите текст комментария.'
    )
    created = models.DateTimeField(auto_now_add=True)

class Follow(models.Model):
    user = models.ForeignKey(
        User,
        verbose_name='Подписчик',
        on_delete=models.CASCADE,
        related_name='follower',
        blank=True, null=True,
    )
    author = models.ForeignKey(
        User,
        verbose_name='На кого подписываются',
        on_delete=models.CASCADE,
        related_name='following',
        blank=True, null=True,
    )

    class Meta:
        constraints = [
            models.UniqueConstraint(
                fields=['user', 'author'],
                name='unique_follow')
        ]


#!/usr/bin/env python3

import re
import sys
import collections

infile = 'day21.txt' if len(sys.argv)<2 else sys.argv[1]
input = open(infile,'r').read().split('\n')

start = ['.#.','..#','###']

rules = dict()

def part(a):
    for y in a:
        print(''.join(y))

def pluck_piece(a,x,y,s):
    p = []
    for py in range(y,y+s):
        p.append(a[py][x:x+s])
    return p

def create_art(a):
    step = 2 if len(a) % 2 == 0 else 3
    nstep = step+1
    na = ['']*(len(a)//step * nstep)
    ny = 0
    #print('step:',step,'nstep:',nstep,'ns:',ns,'len a',len(a))
    for y in range(0,len(a),step):
        nx = 0
        for x in range(0,len(a),step):
            p = pluck_piece(a,x,y,step)
            #print('plucked:',p,"at:",x,y)
            pf = '/'.join(p)
            if pf in rules:
                pe = rules[pf].split('/')
                for r,e in enumerate(pe):
                    na[ny+r] += e
                #print('matched enh:',pe)
            else:
                print("uh... now what")
                quit()
            nx += nstep
        ny += nstep
    return na

print("> Reading Rules:")
for l in input:
    r,e = l.split(' => ')[0].split('/'),l.split(' => ')[1]
    s = len(r)
    # find all permutations of the rule (rotate/mirror/etc)

    print(r)
    rules['/'.join(r)] = e
    rules['/'.join(r[::-1])] = e
    rules['/'.join([x[::-1] for x in r])] = e
    rules['/'.join([x[::-1] for x in r[::-1]])] = e

    # Transpose, then flip
    t = list(map(list,zip(*r)))
    rt = [''.join(x) for x in t]

    rules['/'.join(rt)] = e
    rules['/'.join(rt[::-1])] = e
    rules['/'.join([x[::-1] for x in rt])] = e
    rules['/'.join([x[::-1] for x in rt[::-1]])] = e
print("< Done with Rules")

def create_art_loop(art,loop):
    for i in range(loop):
        nart = create_art(art)
        #print(f"new art [{i}]:")
        #part(nart)
        art = nart
    return art

print("Part 1:")
print(sum(x.count('#') for x in create_art_loop(start,5)))

print("Part 2:")
print(sum(x.count('#') for x in create_art_loop(start,18)))


def make_step(x, y, x_velo, y_velo):
    # Velocity
    x += x_velo
    y += y_velo
    # Drag
    if x_velo > 0:
        x_velo -= 1
    elif x_velo < 0:
        x_velo += 1
    # Gravity
    y_velo -= 1
    return (x, y, x_velo, y_velo)

def launch(x_velo, y_velo, x1, x2, y1, y2, steps=1000):
    x, y = 0, 0  # start
    high_y = 0
    for _ in range(steps):  # try to reach the target area within 1000 steps
        x, y, x_velo, y_velo = make_step(x, y, x_velo, y_velo)
        high_y = max(y, high_y)
        if x1 <= x <= x2 and y1 <= y <= y2:
            return high_y
    return -1

_, area = open("17.txt").read().strip().split(":")
x, y = area.strip().split(",")
x, y = x.split("=")[1], y.split("=")[1]
(x1, x2), (y1, y2) = x.strip().split(".."), y.strip().split("..")
x1, x2, y1, y2 = int(x1), int(x2), int(y1), int(y2)

highest_y = []
for y_velo in range(y1, 200):
    for x_velo in range(x2 + 1):
        high_y = launch(x_velo, y_velo, x1, x2, y1, y2)
        if high_y >= 0:
            highest_y.append(high_y)

print(f"P1: {max(highest_y)}")
print(f"P2: {len(highest_y)}")


salirProceso = False
examingreso = {}
numEstudiantes = 0
opcion = ''

while not salirProceso:
    print("\tMENU\n")
    print("i) Registrar estudiantes")
    print("ii) Consulta de estudiante con mayor puntaje")
    print("iii) Consulta de estudiante con menor puntaje")
    print("iv) Consulta de estudiante a través de la identificacion ")
    print("v) Puntaje promedio obtenido")
    print("vi) Consulta de estudiantes que aprobaron el examen")
    print("vii) Salir")
    opcion = (input("Opcion ->  "))

    if (opcion == 'i'):
        # registra estudiantes
        clave = int((input("Ingresa tu identificacion: ")))
        name = ((input("Ingresa tu nombre: ")))
        score = int((input("Ingresa tu puntaje: ")))
        examingreso.update({clave: [name, score]})
        numEstudiantes = numEstudiantes+1
        print("Diccionario  - ", examingreso)

    if (opcion == 'ii'):
        mayor = 0
        aux = 0
        mayor = examingreso[clave][1]
        index = 0
        #print("nota tomada para prueba : ",mayor)

        for iden in examingreso:
            aux = examingreso[iden][1]
         #   print("nota tomada : ",mayor)
            if(mayor <= aux):
                mayor = aux
                index = iden
        print("IDENTIFICACION\tNOMBRE\tPUNTAJE")
        print(index, "\t\t", examingreso[index]
              [0], "    ", examingreso[index][1])

    if (opcion == 'iii'):
        mayor = 0
        aux = 0
        menor = examingreso[clave][1]
        index = 0
        print("nota menor para prueba : ", menor)

        for iden in examingreso:
            aux = examingreso[iden][1]
            if(menor >= aux):
                menor = aux
                index = iden

        print("IDENTIFICACION\tNOMBRE\tPUNTAJE")
        print(index, "\t\t", examingreso[index]
              [0], "    ", examingreso[index][1])

    if (opcion == 'iv'):
        encontrado = False
        indentificador = 0

        while not encontrado:
            indentificador = int((input("Ingresa tu identificacion: ")))
            valor = []
            try:
                valor = examingreso[indentificador]
            except:
                print("No existe el identificador")
                continue
            if (valor != ''):
                print("IDENTIFICACION\tNOMBRE\tPUNTAJE")
                print(indentificador, "\t\t",
                      examingreso[indentificador][0], "    ", examingreso[indentificador][1])
                encontrado = True

    if (opcion == 'v'):
        aux = 0
        promedio = 0

        for iden in examingreso:
            aux = examingreso[iden][1]
            promedio = promedio + aux

        print("Promedio obtenido : ", round((promedio / numEstudiantes), 2))

    if (opcion == 'vi'):
        aux = 0
        print("IDENTIFICACION\tNOMBRE\tPUNTAJE")
        for iden in examingreso:
            aux = examingreso[iden][1]
            if aux >= 80:
                print(iden, "\t\t", examingreso[iden]
                      [0], "    ", examingreso[iden][1])

    if (opcion == 'vii'):
        print("SALIR DEL PROGRAMA")
        salirProceso = True

if __name__ == '__main__':
    pass