"""
This is a set of classes for reading budget information out of MODFLOW-style
listing files.  Cumulative and incremental budgets are returned as numpy
recarrays, which can then be easily plotted.

"""

import collections
import os
import re
import sys
from datetime import timedelta
import numpy as np

from ..utils.utils_def import totim_to_datetime

class ListBudget(object):
    """
    MODFLOW family list file handling

    Parameters
    ----------
    file_name : str
        the list file name
    budgetkey : str
        the text string identifying the budget table. (default is None)
    timeunit : str
        the time unit to return in the recarray. (default is 'days')

    Notes
    -----
    The ListBudget class should not be instantiated directly.  Access is
    through derived classes: MfListBudget (MODFLOW), SwtListBudget (SEAWAT)
    and SwrListBudget (MODFLOW with the SWR process)

    Examples
    --------
    >>> mf_list = MfListBudget("my_model.list")
    >>> incremental, cumulative = mf_list.get_budget()
    >>> df_in, df_out = mf_list.get_dataframes(start_datetime="10-21-2015")

    """

    def __init__(self, file_name, budgetkey=None, timeunit='days'):

        # Set up file reading
        assert os.path.exists(file_name)
        self.file_name = file_name
        if sys.version_info[0] == 2:
            self.f = open(file_name, 'r')
        elif sys.version_info[0] == 3:
            self.f = open(file_name, 'r', encoding='ascii', errors='replace')

        self.tssp_lines = 0

        # Assign the budgetkey, which should have been overriden
        if budgetkey is None:
            self.set_budget_key()
        else:
            self.budgetkey = budgetkey

        self.totim = []
        self.timeunit = timeunit
        self.idx_map = []
        self.entries = []
        self.null_entries = []

        self.time_line_idx = 20
        if timeunit.upper() == 'SECONDS':
            self.timeunit = 'S'
            self.time_idx = 0
        elif timeunit.upper() == 'MINUTES':
            self.timeunit = 'M'
            self.time_idx = 1
        elif timeunit.upper() == 'HOURS':
            self.timeunit = 'H'
            self.time_idx = 2
        elif timeunit.upper() == 'DAYS':
            self.timeunit = 'D'
            self.time_idx = 3
        elif timeunit.upper() == 'YEARS':
            self.timeunit = 'Y'
            self.time_idx = 4
        else:
            raise Exception('need to reset time_idxs attribute to '
                            'use units other than days and check usage of '
                            'timedelta')

        # Fill budget recarrays
        self._load()
        self._isvalid = False
        if len(self.idx_map) > 0:
            self._isvalid = True

        # Close the open file
        self.f.close()

        # return
        return

    def set_budget_key(self):
        raise Exception('Must be overridden...')

    def isvalid(self):
        """
        Get a boolean indicating if budget data are available in the file.

        Returns
        -------
        out : boolean
            Boolean indicating if budget data are available in the file.

        Examples
        --------
        >>> mf_list = MfListBudget('my_model.list')
        >>> valid = mf_list.isvalid()

        """
        return self._isvalid

    def get_record_names(self):
        """
        Get a list of water budget record names in the file.

        Returns
        -------
        out : list of strings
            List of unique text names in the binary file.

        Examples
        --------
        >>> mf_list = MfListBudget('my_model.list')
        >>> names = mf_list.get_record_names()

        """
        if not self._isvalid:
            return None
        return self.inc.dtype.names

    def get_times(self):
        """
        Get a list of unique water budget times in the list file.

        Returns
        -------
        out : list of floats
            List contains unique water budget simulation times (totim) in list file.

        Examples
        --------
        >>> mf_list = MfListBudget('my_model.list')
        >>> times = mf_list.get_times()

        """
        if not self._isvalid:
            return None
        return self.inc['totim'].tolist()

    def get_kstpkper(self):
        """
        Get a list of unique stress periods and time steps in the list file
        water budgets.

        Returns
        ----------
        out : list of (kstp, kper) tuples
            List of unique kstp, kper combinations in list file.  kstp and
            kper values are zero-based.

        Examples
        --------
        >>> mf_list = MfListBudget("my_model.list")
        >>> kstpkper = mf_list.get_kstpkper()

        """
        if not self._isvalid:
            return None
        kstpkper = []
        for kstp, kper in zip(self.inc['time_step'],
                              self.inc['stress_period']):
            kstpkper.append((kstp, kper))
        return kstpkper

    def get_incremental(self, names=None):
        """
        Get a recarray with the incremental water budget items in the list file.

        Parameters
        ----------
        names : str or list of strings
            Selection of column names to return.  If names is not None then
            totim, time_step, stress_period, and selection(s) will be returned.
            (default is None).

        Returns
        -------
        out : recarray
            Numpy recarray with the water budget items in list file. The
            recarray also includes totim, time_step, and stress_period.

        Examples
        --------
        >>> mf_list = MfListBudget("my_model.list")
        >>> incremental = mf_list.get_incremental()

        """
        if not self._isvalid:
            return None
        if names is None:
            return self.inc
        else:
            if not isinstance(names, list):
                names = [names]
            names.insert(0, 'stress_period')
            names.insert(0, 'time_step')
            names.insert(0, 'totim')
            return self.inc[names].view(np.recarray)

    def get_cumulative(self, names=None):
        """
        Get a recarray with the cumulative water budget items in the list file.

        Parameters
        ----------
        names : str or list of strings
            Selection of column names to return.  If names is not None then
            totim, time_step, stress_period, and selection(s) will be returned.
            (default is None).

        Returns
        -------
        out : recarray
            Numpy recarray with the water budget items in list file. The
            recarray also includes totim, time_step, and stress_period.

        Examples
        --------
        >>> mf_list = MfListBudget("my_model.list")
        >>> cumulative = mf_list.get_cumulative()

       """
        if not self._isvalid:
            return None
        if names is None:
            return self.cum
        else:
            if not isinstance(names, list):
                names = [names]
            names.insert(0, 'stress_period')
            names.insert(0, 'time_step')
            names.insert(0, 'totim')
            return self.cum[names].view(np.recarray)

    def get_budget(self, names=None):
        """
        Get the recarrays with the incremental and cumulative water budget items
        in the list file.

        Parameters
        ----------
        names : str or list of strings
            Selection of column names to return.  If names is not None then
            totim, time_step, stress_period, and selection(s) will be returned.
            (default is None).

        Returns
        -------
        out : recarrays
            Numpy recarrays with the water budget items in list file. The
            recarray also includes totim, time_step, and stress_period. A
            separate recarray is returned for the incremental and cumulative
            water budget entries.

        Examples
        --------
        >>> mf_list = MfListBudget("my_model.list")
        >>> budget = mf_list.get_budget()

        """
        if not self._isvalid:
            return None
        if names is None:
            return self.inc, self.cum
        else:
            if not isinstance(names, list):
                names = [names]
            names.insert(0, 'stress_period')
            names.insert(0, 'time_step')
            names.insert(0, 'totim')
            return self.inc[names].view(np.recarray), self.cum[names].view(
                    np.recarray)

    def get_data(self, kstpkper=None, idx=None, totim=None, incremental=False):
        """
        Get water budget data from the list file for the specified conditions.

        Parameters
        ----------
        idx : int
            The zero-based record number.  The first record is record 0.
            (default is None).
        kstpkper : tuple of ints
            A tuple containing the time step and stress period (kstp, kper).
            These are zero-based kstp and kper values. (default is None).
        totim : float
            The simulation time. (default is None).
        incremental : bool
            Boolean flag used to determine if incremental or cumulative water
            budget data for the specified conditions will be returned. If
            incremental=True, incremental water budget data will be returned.
            If incremental=False, cumulative water budget data will be
            returned. (default is False).

        Returns
        -------
        data : numpy recarray
            Array has size (number of budget items, 3). Recarray names are 'index',
            'value', 'name'.

        See Also
        --------

        Notes
        -----
        if both kstpkper and totim are None, will return the last entry

        Examples
        --------
        >>> import matplotlib.pyplot as plt
        >>> import flopy
        >>> mf_list = flopy.utils.MfListBudget("my_model.list")
        >>> data = mf_list.get_data(kstpkper=(0,0))
        >>> plt.bar(data['index'], data['value'])
        >>> plt.xticks(data['index'], data['name'], rotation=45, size=6)
        >>> plt.show()

        """
        if not self._isvalid:
            return None
        ipos = None
        if kstpkper is not None:
            try:
                ipos = self.get_kstpkper().index(kstpkper)
            except:
                pass
        elif totim is not None:
            try:
                ipos = self.get_times().index(totim)
            except:
                pass
        elif idx is not None:
            ipos = idx
        else:
            ipos = -1

        if ipos is None:
            print('Could not find specified condition.')
            print('  kstpkper = {}'.format(kstpkper))
            print('  totim = {}'.format(totim))
            return None

        if incremental:
            t = self.inc[ipos]
        else:
            t = self.cum[ipos]

        dtype = np.dtype(
                [('index', np.int32), ('value', np.float32), ('name', '|S25')])
        v = np.recarray(shape=(len(self.inc.dtype.names[3:])), dtype=dtype)
        for i, name in enumerate(self.inc.dtype.names[3:]):
            mult = 1.
            if '_OUT' in name:
                mult = -1.
            v[i]['index'] = i
            v[i]['value'] = mult * t[name]
            v[i]['name'] = name
        return v

    def get_dataframes(self, start_datetime='1-1-1970',diff=False):
        """
        Get pandas dataframes with the incremental and cumulative water budget
        items in the list file.

        Parameters
        ----------
        start_datetime : str
            If start_datetime is passed as None, the rows are indexed on totim.
            Otherwise, a DatetimeIndex is set. (default is 1-1-1970).

        Returns
        -------
        out : panda dataframes
            Pandas dataframes with the incremental and cumulative water budget
            items in list file. A separate pandas dataframe is returned for the
            incremental and cumulative water budget entries.

        Examples
        --------
        >>> mf_list = MfListBudget("my_model.list")
        >>> incrementaldf, cumulativedf = mf_list.get_dataframes()

        """

        try:
            import pandas as pd
        except Exception as e:
            raise Exception(
                    "ListBudget.get_dataframe() error import pandas: " + \
                    str(e))

        if not self._isvalid:
            return None
        totim = self.get_times()
        if start_datetime is not None:
            totim = totim_to_datetime(totim,
                                      start=pd.to_datetime(start_datetime),
                                      timeunit=self.timeunit)

        df_flux = pd.DataFrame(self.inc, index=totim).loc[:, self.entries]
        df_vol = pd.DataFrame(self.cum, index=totim).loc[:, self.entries]

        if not diff:
            return df_flux, df_vol

        else:
            in_names = [col for col in df_flux.columns if col.endswith("_IN")]
            out_names = [col for col in df_flux.columns if col.endswith("_OUT")]
            #print(in_names,out_names)
            #print(df_flux.columns)
            base_names = [name.replace("_IN",'') for name in in_names]
            for name in base_names:
                in_name = name + "_IN"
                out_name = name + "_OUT"
                df_flux.loc[:,name.lower()] = df_flux.loc[:,in_name] - df_flux.loc[:,out_name]
                df_flux.pop(in_name)
                df_flux.pop(out_name)
                df_vol.loc[:,name.lower()] = df_vol.loc[:,in_name] - df_vol.loc[:,out_name]
                df_vol.pop(in_name)
                df_vol.pop(out_name)
            cols = list(df_flux.columns)
            cols.sort()
            cols = [col.lower() for col in cols]
            df_flux.columns = cols
            df_vol.columns = cols
            return df_flux, df_vol
    def _build_index(self, maxentries):
        self.idx_map = self._get_index(maxentries)
        return

    def _get_index(self, maxentries):
        # --parse through the file looking for matches and parsing ts and sp
        idxs = []
        l_count = 1
        while True:
            seekpoint = self.f.tell()
            line = self.f.readline()
            if line == '':
                break
            if self.budgetkey in line:
                for l in range(self.tssp_lines):
                    line = self.f.readline()
                try:
                    ts, sp = self._get_ts_sp(line)
                except:
                    print('unable to cast ts,sp on line number', l_count,
                          ' line: ', line)
                    break
                # print('info found for timestep stress period',ts,sp)

                idxs.append([ts, sp, seekpoint])

                if maxentries and len(idxs) >= maxentries:
                    break

        return idxs

    def _seek_to_string(self, s):
        """
        Parameters
        ----------
        s : str
            Seek through the file to the next occurrence of s.  Return the
            seek location when found.

        Returns
        -------
        seekpoint : int
            Next location of the string

        """
        while True:
            seekpoint = self.f.tell()
            line = self.f.readline()
            if line == '':
                break
            if s in line:
                break
        return seekpoint

    def _get_ts_sp(self, line):
        """
        From the line string, extract the time step and stress period numbers.

        """

        # Old method.  Was not generic enough.
        # ts = int(line[self.ts_idxs[0]:self.ts_idxs[1]])
        # sp = int(line[self.sp_idxs[0]:self.sp_idxs[1]])

        # Get rid of nasty things
        line = line.replace(',', '')

        searchstring = 'TIME STEP'
        idx = line.index(searchstring) + len(searchstring)
        ll = line[idx:].strip().split()
        ts = int(ll[0])

        searchstring = 'STRESS PERIOD'
        idx = line.index(searchstring) + len(searchstring)
        ll = line[idx:].strip().split()
        sp = int(ll[0])

        return ts, sp

    def _set_entries(self):
        if len(self.idx_map) < 1:
            return None, None
        if len(self.entries) > 0:
            raise Exception('entries already set:' + str(self.entries))
        if not self.idx_map:
            raise Exception('must call build_index before call set_entries')
        try:
            incdict, cumdict = self._get_sp(self.idx_map[0][0],
                                            self.idx_map[0][1],
                                            self.idx_map[0][2])
        except:
            raise Exception('unable to read budget information from first '
                            'entry in list file')
        self.entries = incdict.keys()
        null_entries = collections.OrderedDict()
        incdict = collections.OrderedDict()
        cumdict = collections.OrderedDict()
        for entry in self.entries:
            incdict[entry] = []
            cumdict[entry] = []
            null_entries[entry] = np.NaN
        self.null_entries = [null_entries, null_entries]
        return incdict, cumdict

    def _load(self, maxentries=None):
        self._build_index(maxentries)
        incdict, cumdict = self._set_entries()
        if incdict is None and cumdict is None:
            return
        totim = []
        for ts, sp, seekpoint in self.idx_map:
            tinc, tcum = self._get_sp(ts, sp, seekpoint)
            for entry in self.entries:
                incdict[entry].append(tinc[entry])
                cumdict[entry].append(tcum[entry])

            # Get the time for this record
            seekpoint = self._seek_to_string('TIME SUMMARY AT END')
            tslen, sptim, tt = self._get_totim(ts, sp, seekpoint)
            totim.append(tt)

        # get kstp and kper
        idx_array = np.array(self.idx_map)

        # build dtype for recarray
        dtype_tups = [('totim', np.float32), ("time_step", np.int32),
                      ("stress_period", np.int32)]
        for entry in self.entries:
            dtype_tups.append((entry, np.float32))
        dtype = np.dtype(dtype_tups)

        # create recarray
        nentries = len(incdict[entry])
        self.inc = np.recarray(shape=(nentries,), dtype=dtype)
        self.cum = np.recarray(shape=(nentries,), dtype=dtype)

        # fill each column of the recarray
        for entry in self.entries:
            self.inc[entry] = incdict[entry]
            self.cum[entry] = cumdict[entry]

        # file the totim, time_step, and stress_period columns for the
        # incremental and cumulative recarrays (zero-based kstp,kper)
        self.inc['totim'] = np.array(totim)[:]
        self.inc["time_step"] = idx_array[:, 0] - 1
        self.inc["stress_period"] = idx_array[:, 1] - 1

        self.cum['totim'] = np.array(totim)[:]
        self.cum["time_step"] = idx_array[:, 0] - 1
        self.cum["stress_period"] = idx_array[:, 1] - 1

        return

    def _get_sp(self, ts, sp, seekpoint):
        self.f.seek(seekpoint)
        # --read to the start of the "in" budget information
        while True:
            line = self.f.readline()
            if line == '':
                print(
                        'end of file found while seeking budget information for ts,sp',
                        ts, sp)
                return self.null_entries

            # --if there are two '=' in this line, then it is a budget line
            if len(re.findall('=', line)) == 2:
                break

        tag = 'IN'
        incdict = collections.OrderedDict()
        cumdict = collections.OrderedDict()
        while True:

            if line == '':
                # raise Exception('end of file found while seeking budget information')
                print(
                        'end of file found while seeking budget information for ts,sp',
                        ts, sp)
                return self.null_entries
            if len(re.findall('=', line)) == 2:
                try:
                    entry, flux, cumu = self._parse_budget_line(line)
                except e:
                    print('error parsing budget line in ts,sp', ts, sp)
                    return self.null_entries
                if flux is None:
                    print(
                            'error casting in flux for', entry,
                            ' to float in ts,sp',
                            ts, sp)
                    return self.null_entries
                if cumu is None:
                    print(
                            'error casting in cumu for', entry,
                            ' to float in ts,sp',
                            ts, sp)
                    return self.null_entries
                if entry.endswith(tag.upper()):
                    if ' - ' in entry.upper():
                        key = entry.replace(' ', '')
                    else:
                        key = entry.replace(' ', '_')
                elif 'PERCENT DISCREPANCY' in entry.upper():
                    key = entry.replace(' ', '_')
                else:
                    key = '{}_{}'.format(entry.replace(' ', '_'), tag)
                incdict[key] = flux
                cumdict[key] = cumu
            else:
                if 'OUT:' in line.upper():
                    tag = 'OUT'
            line = self.f.readline()
            if entry.upper() == 'PERCENT DISCREPANCY':
                break

        return incdict, cumdict

    def _parse_budget_line(self, line):

        # get the budget item name
        entry = line.strip().split('=')[0].strip()

        # get the cumulative string
        idx = line.index('=') + 1
        line2 = line[idx:]
        ll = line2.strip().split()
        cu_str = ll[0]

        idx = line2.index('=') + 1
        fx_str = line2[idx:].strip()

        #
        # cu_str = line[self.cumu_idxs[0]:self.cumu_idxs[1]]
        # fx_str = line[self.flux_idxs[0]:self.flux_idxs[1]]

        flux, cumu = None, None
        try:
            cumu = float(cu_str)
        except:
            if 'NAN' in cu_str.strip().upper():
                cumu = np.NaN
        try:
            flux = float(fx_str)
        except:
            if 'NAN' in fx_str.strip().upper():
                flux = np.NaN
        return entry, flux, cumu

    def _get_totim(self, ts, sp, seekpoint):
        self.f.seek(seekpoint)
        # --read header lines
        ihead = 0
        while True:
            line = self.f.readline()
            ihead += 1
            if line == '':
                print(
                        'end of file found while seeking time information for ts,sp',
                        ts, sp)
                return np.NaN, np.NaN, np.Nan
            elif ihead == 2 and 'SECONDS     MINUTES      HOURS       DAYS        YEARS' not in line:
                break
            elif '-----------------------------------------------------------' in line:
                line = self.f.readline()
                break
        tslen = self._parse_time_line(line)
        if tslen == None:
            print('error parsing tslen for ts,sp', ts, sp)
            return np.NaN, np.NaN, np.Nan

        sptim = self._parse_time_line(self.f.readline())
        if sptim == None:
            print('error parsing sptim for ts,sp', ts, sp)
            return np.NaN, np.NaN, np.Nan

        totim = self._parse_time_line(self.f.readline())
        if totim == None:
            print('error parsing totim for ts,sp', ts, sp)
            return np.NaN, np.NaN, np.Nan
        return tslen, sptim, totim

    def _parse_time_line(self, line):
        if line == '':
            print('end of file found while parsing time information')
            return None
        try:
            time_str = line[self.time_line_idx:]
            raw = time_str.split()
            idx = self.time_idx
            # catch case where itmuni is undefined
            # in this case, the table format is different
            try:
                v = float(raw[0])
            except:
                time_str = line[45:]
                raw = time_str.split()
                idx = 0
            tval = float(raw[idx])
        except:
            print('error parsing tslen information', time_str)
            return None
        return tval

class SwtListBudget(ListBudget):
    """

    """

    def set_budget_key(self):
        self.budgetkey = 'MASS BUDGET FOR ENTIRE MODEL'
        return

class MfListBudget(ListBudget):
    """

    """

    def set_budget_key(self):
        self.budgetkey = 'VOLUMETRIC BUDGET FOR ENTIRE MODEL'
        return

class MfusgListBudget(ListBudget):
    """

    """

    def set_budget_key(self):
        self.budgetkey = 'VOLUMETRIC BUDGET FOR ENTIRE MODEL'
        return

class SwrListBudget(ListBudget):
    """

    """

    def set_budget_key(self):
        self.budgetkey = 'VOLUMETRIC SURFACE WATER BUDGET FOR ENTIRE MODEL'
        self.tssp_lines = 1
        return


"""
========================================================
Gaussian process regression (GPR) on Mauna Loa CO2 data.
========================================================

This example is based on Section 5.4.3 of "Gaussian Processes for Machine
Learning" [RW2006]. It illustrates an example of complex kernel engineering and
hyperparameter optimization using gradient ascent on the
log-marginal-likelihood. The data consists of the monthly average atmospheric
CO2 concentrations (in parts per million by volume (ppmv)) collected at the
Mauna Loa Observatory in Hawaii, between 1958 and 1997. The objective is to
model the CO2 concentration as a function of the time t.

The kernel is composed of several terms that are responsible for explaining
different properties of the signal:

- a long term, smooth rising trend is to be explained by an RBF kernel. The
  RBF kernel with a large length-scale enforces this component to be smooth;
  it is not enforced that the trend is rising which leaves this choice to the
  GP. The specific length-scale and the amplitude are free hyperparameters.

- a seasonal component, which is to be explained by the periodic
  ExpSineSquared kernel with a fixed periodicity of 1 year. The length-scale
  of this periodic component, controlling its smoothness, is a free parameter.
  In order to allow decaying away from exact periodicity, the product with an
  RBF kernel is taken. The length-scale of this RBF component controls the
  decay time and is a further free parameter.

- smaller, medium term irregularities are to be explained by a
  RationalQuadratic kernel component, whose length-scale and alpha parameter,
  which determines the diffuseness of the length-scales, are to be determined.
  According to [RW2006], these irregularities can better be explained by
  a RationalQuadratic than an RBF kernel component, probably because it can
  accommodate several length-scales.

- a "noise" term, consisting of an RBF kernel contribution, which shall
  explain the correlated noise components such as local weather phenomena,
  and a WhiteKernel contribution for the white noise. The relative amplitudes
  and the RBF's length scale are further free parameters.

Maximizing the log-marginal-likelihood after subtracting the target's mean
yields the following kernel with an LML of -83.214::

   34.4**2 * RBF(length_scale=41.8)
   + 3.27**2 * RBF(length_scale=180) * ExpSineSquared(length_scale=1.44,
                                                      periodicity=1)
   + 0.446**2 * RationalQuadratic(alpha=17.7, length_scale=0.957)
   + 0.197**2 * RBF(length_scale=0.138) + WhiteKernel(noise_level=0.0336)

Thus, most of the target signal (34.4ppm) is explained by a long-term rising
trend (length-scale 41.8 years). The periodic component has an amplitude of
3.27ppm, a decay time of 180 years and a length-scale of 1.44. The long decay
time indicates that we have a locally very close to periodic seasonal
component. The correlated noise has an amplitude of 0.197ppm with a length
scale of 0.138 years and a white-noise contribution of 0.197ppm. Thus, the
overall noise level is very small, indicating that the data can be very well
explained by the model. The figure shows also that the model makes very
confident predictions until around 2015.
"""
print(__doc__)

# Authors: Jan Hendrik Metzen <jhm@informatik.uni-bremen.de>
#
# License: BSD 3 clause

import numpy as np

from matplotlib import pyplot as plt

from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels \
    import RBF, WhiteKernel, RationalQuadratic, ExpSineSquared
from sklearn.datasets import fetch_mldata

data = fetch_mldata('mauna-loa-atmospheric-co2').data
X = data[:, [1]]
y = data[:, 0]

# Kernel with parameters given in GPML book
k1 = 66.0**2 * RBF(length_scale=67.0)  # long term smooth rising trend
k2 = 2.4**2 * RBF(length_scale=90.0) \
    * ExpSineSquared(length_scale=1.3, periodicity=1.0)  # seasonal component
# medium term irregularity
k3 = 0.66**2 \
    * RationalQuadratic(length_scale=1.2, alpha=0.78)
k4 = 0.18**2 * RBF(length_scale=0.134) \
    + WhiteKernel(noise_level=0.19**2)  # noise terms
kernel_gpml = k1 + k2 + k3 + k4

gp = GaussianProcessRegressor(kernel=kernel_gpml, alpha=0,
                              optimizer=None, normalize_y=True)
gp.fit(X, y)

print("GPML kernel: %s" % gp.kernel_)
print("Log-marginal-likelihood: %.3f"
      % gp.log_marginal_likelihood(gp.kernel_.theta))

# Kernel with optimized parameters
k1 = 50.0**2 * RBF(length_scale=50.0)  # long term smooth rising trend
k2 = 2.0**2 * RBF(length_scale=100.0) \
    * ExpSineSquared(length_scale=1.0, periodicity=1.0,
                     periodicity_bounds="fixed")  # seasonal component
# medium term irregularities
k3 = 0.5**2 * RationalQuadratic(length_scale=1.0, alpha=1.0)
k4 = 0.1**2 * RBF(length_scale=0.1) \
    + WhiteKernel(noise_level=0.1**2,
                  noise_level_bounds=(1e-3, np.inf))  # noise terms
kernel = k1 + k2 + k3 + k4

gp = GaussianProcessRegressor(kernel=kernel, alpha=0,
                              normalize_y=True)
gp.fit(X, y)

print("\nLearned kernel: %s" % gp.kernel_)
print("Log-marginal-likelihood: %.3f"
      % gp.log_marginal_likelihood(gp.kernel_.theta))

X_ = np.linspace(X.min(), X.max() + 30, 1000)[:, np.newaxis]
y_pred, y_std = gp.predict(X_, return_std=True)

# Illustration
plt.scatter(X, y, c='k')
plt.plot(X_, y_pred)
plt.fill_between(X_[:, 0], y_pred - y_std, y_pred + y_std,
                 alpha=0.5, color='k')
plt.xlim(X_.min(), X_.max())
plt.xlabel("Year")
plt.ylabel(r"CO$_2$ in ppm")
plt.title(r"Atmospheric CO$_2$ concentration at Mauna Loa")
plt.tight_layout()
plt.show()


#!/usr/bin/env python

def main():

    import os
    import argparse
    from qap.script_utils import check_csv_missing_subs, csv_to_pandas_df, \
        write_inputs_dict_to_yaml_file, read_yml_file
    from qap.qap_utils import raise_smart_exception

    parser = argparse.ArgumentParser()
    parser.add_argument("output_csv", type=str,
                        help="the main output directory of the QAP run "
                        "which contains the participant directories")
    parser.add_argument("data_config", type=str,
                        help="the main output directory of the QAP run "
                        "which contains the participant directories")
    parser.add_argument("data_type", type=str,
                        help="the main output directory of the QAP run "
                        "which contains the participant directories")

    args = parser.parse_args()

    csv_df = csv_to_pandas_df(args.output_csv)
    data_dict = read_yml_file(args.data_config)

    new_dict = check_csv_missing_subs(csv_df, data_dict, args.data_type)

    if new_dict:
        out_file = os.path.join(os.getcwd(),
                                "missing_%s_data.yml" % args.data_type)
        write_inputs_dict_to_yaml_file(new_dict, out_file)

if __name__ == "__main__":
    main()


'''
.. module:: skrf.io.general
========================================
general (:mod:`skrf.io.general`)
========================================

General io functions for reading and writing skrf objects

.. autosummary::
    :toctree: generated/

    read
    read_all
    read_all_networks
    write
    write_all
    save_sesh

Writing output to spreadsheet

.. autosummary::
    :toctree: generated/

    network_2_spreadsheet
    networkset_2_spreadsheet

'''
import sys

import six.moves.cPickle as pickle
from six.moves.cPickle import UnpicklingError

import inspect
import os
import zipfile
import warnings
import sys

from ..util import get_extn, get_fid
from ..network import Network
from ..frequency import Frequency
from ..media import  Media
from ..networkSet import NetworkSet
from ..calibration.calibration import Calibration

from copy import copy
dir_ = copy(dir)

# delayed import: from pandas import DataFrame, Series for ntwk_2_spreadsheet

# file extension conventions for skrf objects.
global OBJ_EXTN
OBJ_EXTN = [
    [Frequency, 'freq'],
    [Network, 'ntwk'],
    [NetworkSet, 'ns'],
    [Calibration, 'cal'],
    [Media, 'med'],
    [object, 'p'],
    ]

def read(file, *args, **kwargs):
    '''
    Read  skrf object[s] from a pickle file

    Reads a skrf object that is written with :func:`write`, which uses
    the :mod:`pickle` module.

    Parameters
    ------------
    file : str or file-object
        name of file, or  a file-object
    \*args, \*\*kwargs : arguments and keyword arguments
        passed through to pickle.load

    Examples
    -------------
    >>> n = rf.Network(f=[1,2,3],s=[1,1,1],z0=50)
    >>> n.write('my_ntwk.ntwk')
    >>> n_2 = rf.read('my_ntwk.ntwk')

    See Also
    ----------
    read : read a skrf object
    write : write skrf object[s]
    read_all : read all skrf objects in a directory
    write_all : write dictionary of skrf objects to a directory

    Notes
    -------
    if `file` is a file-object it is left open, if it is a filename then
    a file-object is opened and closed. If file is a file-object
    and reading fails, then the position is reset back to 0 using seek
    if possible.
    '''
    fid = get_fid(file, mode='rb')
    try:
        obj = pickle.load(fid, *args, **kwargs)
    except (UnpicklingError, UnicodeDecodeError) as e:
        # if fid is seekable then reset to beginning of file
        fid.seek(0)

        if isinstance(file, str):
            # we created the fid so close it
            fid.close()
        raise

    if isinstance(file, str):
        # we created the fid so close it
        fid.close()

    return obj

def write(file, obj, overwrite = True):
    '''
    Write skrf object[s] to a file

    This uses the :mod:`pickle` module to write skrf objects to a file.
    Note that you can write any pickl-able python object. For example,
    you can write a list or dictionary of :class:`~skrf.network.Network`
    objects
    or :class:`~skrf.calibration.calibration.Calibration` objects. This
    will write out a single file. If you would like to write out a
    seperate file for each object, use :func:`write_all`.

    Parameters
    ------------
    file : file or string
        File or filename to which the data is saved.  If file is a
        file-object, then the filename is unchanged.  If file is a
        string, an appropriate extension will be appended to the file
        name if it does not already have an extension.

    obj : an object, or list/dict of objects
        object or list/dict of objects to write to disk

    overwrite : Boolean
        if file exists, should it be overwritten?

    Notes
    -------

    If `file` is a str, but doesnt contain a suffix, one is chosen
    automatically. Here are the extensions

    ====================================================  ===============
    skrf object                                           extension
    ====================================================  ===============
    :class:`~skrf.frequency.Frequency`                    '.freq'
    :class:`~skrf.network.Network`                        '.ntwk'
    :class:`~skrf.networkSet.NetworkSet`                  '.ns'
    :class:`~skrf.calibration.calibration.Calibration`    '.cal'
    :class:`~skrf.media.media.Media`                      '.med'
    other                                                 '.p'
    ====================================================  ===============

    To make the file written by this method cross-platform, the pickling
    protocol 2 is used. See :mod:`pickle` for more info.

    Examples
    -------------
    Convert a touchstone file to a pickled Network,

    >>> n = rf.Network('my_ntwk.s2p')
    >>> rf.write('my_ntwk',n)
    >>> n_red = rf.read('my_ntwk.ntwk')

    Writing a list of different objects

    >>> n = rf.Network('my_ntwk.s2p')
    >>> ns = rf.NetworkSet([n,n,n])
    >>> rf.write('out',[n,ns])
    >>> n_red = rf.read('out.p')

    See Also
    ------------
    read : read a skrf object
    write : write skrf object[s]
    read_all : read all skrf objects in a directory
    write_all : write dictionary of skrf objects to a directory
    skrf.network.Network.write : write method of Network
    skrf.calibration.calibration.Calibration.write : write method of Calibration

    '''
    if isinstance(file, str):
        extn = get_extn(file)
        if extn is None:
            # if there is not extension add one
            for obj_extn in OBJ_EXTN:
                if isinstance(obj, obj_extn[0]):
                    extn = obj_extn[1]
                    break
            file = file + '.' + extn

        if os.path.exists(file):
            if not overwrite:
                warnings.warn('file exists, and overwrite option is False. Not writing.')
                return

        with open(file, 'wb') as fid:
            pickle.dump(obj, fid, protocol=2)

    else:
        fid = file
        pickle.dump(obj, fid, protocol=2)
        fid.close()

def read_all(dir='.', contains = None, f_unit = None, obj_type=None):
    '''
    Read all skrf objects in a directory

    Attempts to load all files in `dir`, using :func:`read`. Any file
    that is not readable by skrf is skipped. Optionally, simple filtering
    can be achieved through the use of `contains` argument.

    Parameters
    --------------
    dir : str, optional
        the directory to load from, default  \'.\'
    contains : str, optional
        if not None, only files containing this substring will be loaded
    f_unit : ['hz','khz','mhz','ghz','thz']
        for all :class:`~skrf.network.Network` objects, set their
        frequencies's :attr:`~skrf.frequency.Frequency.f_unit`
    obj_type : str
        Name of skrf object types to read (ie 'Network')

    Returns
    ---------
    out : dictionary
        dictionary containing all loaded skrf objects. keys are the
        filenames without extensions, and the values are the objects

    Examples
    ----------
    >>> rf.read_all('skrf/data/')
    {'delay_short': 1-Port Network: 'delay_short',  75-110 GHz, 201 pts, z0=[ 50.+0.j],
    'line': 2-Port Network: 'line',  75-110 GHz, 201 pts, z0=[ 50.+0.j  50.+0.j],
    'ntwk1': 2-Port Network: 'ntwk1',  1-10 GHz, 91 pts, z0=[ 50.+0.j  50.+0.j],
    'one_port': one port Calibration: 'one_port', 500-750 GHz, 201 pts, 4-ideals/4-measured,
    ...

    >>> rf.read_all('skrf/data/', obj_type = 'Network')
    {'delay_short': 1-Port Network: 'delay_short',  75-110 GHz, 201 pts, z0=[ 50.+0.j],
    'line': 2-Port Network: 'line',  75-110 GHz, 201 pts, z0=[ 50.+0.j  50.+0.j],
    'ntwk1': 2-Port Network: 'ntwk1',  1-10 GHz, 91 pts, z0=[ 50.+0.j  50.+0.j],
    ...

    See Also
    ----------
    read : read a skrf object
    write : write skrf object[s]
    read_all : read all skrf objects in a directory
    write_all : write dictionary of skrf objects to a directory
    '''

    out={}
    for filename in os.listdir(dir):
        if contains is not None and contains not in filename:
            continue
        fullname = os.path.join(dir,filename)
        keyname = os.path.splitext(filename)[0]
        try:
            out[keyname] = read(fullname)
            continue
        except:
            pass

        try:
            out[keyname] = Network(fullname)
            continue
        except:
            pass

    if f_unit is not None:
        for keyname in out:
            try:
                out[keyname].frequency.unit = f_unit
            except:
                pass

    if obj_type is not None:
        out = dict([(k, out[k]) for k in out if
            isinstance(out[k],sys.modules[__name__].__dict__[obj_type])])

    return out

def read_all_networks(*args, **kwargs):
    '''
    Read all networks in a directory.

    This is a convenience function. It just calls::

        read_all(*args,obj_type='Network', **kwargs)

    See Also
    ----------
    read_all
    '''
    if 'f_unit' not in kwargs:
        kwargs.update({'f_unit':'ghz'})
    return read_all(*args,obj_type='Network', **kwargs)

ran = read_all_networks

def write_all(dict_objs, dir='.', *args, **kwargs):
    '''
    Write a dictionary of skrf objects individual files in `dir`.

    Each object is written to its own file. The filename used for each
    object is taken from its key in the dictionary. If no extension
    exists in the key, then one is added. See :func:`write` for a list
    of extensions. If you would like to write the dictionary to a single
    output file use :func:`write`.

    Notes
    -------
    Any object in  dict_objs that is pickl-able will be written.

    Parameters
    ------------
    dict_objs : dict
        dictionary of skrf objects
    dir : str
        directory to save skrf objects into
    \*args, \*\*kwargs :
            passed through to :func:`~skrf.io.general.write`. `overwrite`
            option may be of use.

    See Also
    -----------
    read : read a skrf object
    write : write skrf object[s]
    read_all : read all skrf objects in a directory
    write_all : write dictionary of skrf objects to a directory

    Examples
    ----------
    Writing a diction of different skrf objects

    >>> from skrf.data import line, short
    >>> d = {'ring_slot':ring_slot, 'one_port_cal':one_port_cal}
    >>> rf.write_all(d)

    '''
    if not os.path.exists('.'):
        raise OSError('No such directory: %s'%dir)

    for k in dict_objs:
        filename = k
        obj = dict_objs[k]

        extn = get_extn(filename)
        if extn is None:
            # if there is not extension add one
            for obj_extn in OBJ_EXTN:
                if isinstance(obj, obj_extn[0]):
                    extn = obj_extn[1]
                    break
            filename = filename + '.' + extn
        try:
            with open(os.path.join(dir+'/', filename), 'wb') as fid:
                write(fid, obj,*args, **kwargs)
        except Exception as inst:
            print(inst)
            warnings.warn('couldnt write %s: %s'%(k,str(inst)))

            pass

def save_sesh(dict_objs, file='skrfSesh.p', module='skrf', exclude_prefix='_'):
    '''
    Save all `skrf` objects in the local namespace.

    This is used to save current workspace in a hurry, by passing it the
    output of :func:`locals` (see Examples). Note this can be
    used for other modules as well by passing a different `module` name.

    Parameters
    ------------
    dict_objs : dict
        dictionary containing `skrf` objects. See the Example.
    file : str or file-object, optional
        the file to save all objects to
    module : str, optional
        the module name to grep for.
    exclude_prefix: str, optional
        dont save objects which have this as a prefix.

    See Also
    ----------
    read : read a skrf object
    write : write skrf object[s]
    read_all : read all skrf objects in a directory
    write_all : write dictionary of skrf objects to a directory

    Examples
    ---------
    Write out all skrf objects in current namespace.

    >>> rf.write_all(locals(), 'mysesh.p')

    '''
    objects = {}
    print('pickling: ')
    for k in dict_objs:
        try:
            if module  in inspect.getmodule(dict_objs[k]).__name__:
                try:
                    pickle.dumps(dict_objs[k])
                    if k[0] != '_':
                        objects[k] = dict_objs[k]
                        print(k+', ')
                finally:
                    pass

        except(AttributeError, TypeError):
            pass
    if len (objects ) == 0:
        print('nothing')

    write(file, objects)

def load_all_touchstones(dir = '.', contains=None, f_unit=None):
    '''
    Loads all touchtone files in a given dir into a dictionary.

    Notes
    -------
    Alternatively you can use the :func:`read_all` function.

    Parameters
    -----------
    dir :   string
            the path
    contains :      string
            a string the filenames must contain to be loaded.
    f_unit  : ['hz','mhz','ghz']
            the frequency unit to assign all loaded networks. see
            :attr:`frequency.Frequency.unit`.

    Returns
    ---------
    ntwkDict : a dictonary with keys equal to the file name (without
            a suffix), and values equal to the corresponding ntwk types

    Examples
    ----------
    >>> ntwk_dict = rf.load_all_touchstones('.', contains ='20v')

    See Also
    -----------
    read_all
    '''
    ntwkDict = {}

    for f in os.listdir (dir):
        if contains is not None and contains not in f:
            continue
        fullname = os.path.join(dir,f)
        keyname,extn = os.path.splitext(f)
        extn = extn.lower()
        try:
            if extn[1]== 's' and extn[-1]=='p':
                ntwkDict[keyname]=(Network(dir +'/'+f))
                if f_unit is not None: ntwkDict[keyname].frequency.unit=f_unit
        except:
            pass
    return ntwkDict

def write_dict_of_networks(ntwkDict, dir='.'):
    '''
    Saves a dictionary of networks touchstone files in a given directory

    The filenames assigned to the touchstone files are taken from
    the keys of the dictionary.

    Parameters
    -----------
    ntwkDict : dictionary
            dictionary of :class:`Network` objects
    dir : string
            directory to write touchstone file to

    '''
    warnings.warn('Deprecated. use write_all.', DeprecationWarning)
    for ntwkKey in ntwkDict:
        ntwkDict[ntwkKey].write_touchstone(filename = dir+'/'+ntwkKey)

def read_csv(filename):
    '''
    Read a 2-port s-parameter data from a csv file.

    Specifically, this reads a two-port csv file saved from a Rohde Shcwarz
    ZVA-40, and possibly other network analyzers. It returns into a
    :class:`Network` object.

    Parameters
    ------------
    filename : str
            name of file

    Returns
    --------
    ntwk : :class:`Network` object
            the network representing data in the csv file
    '''

    ntwk = Network(name=filename[:-4])
    try:
        data = npy.loadtxt(filename, skiprows=3,delimiter=',',\
                usecols=range(9))
        s11 = data[:,1] +1j*data[:,2]
        s21 = data[:,3] +1j*data[:,4]
        s12 = data[:,5] +1j*data[:,6]
        s22 = data[:,7] +1j*data[:,8]
        ntwk.s = npy.array([[s11, s21],[s12,s22]]).transpose().reshape(-1,2,2)
    except(IndexError):
        data = npy.loadtxt(filename, skiprows=3,delimiter=',',\
                usecols=range(3))
        ntwk.s = data[:,1] +1j*data[:,2]

    ntwk.frequency.f = data[:,0]
    ntwk.frequency.unit='ghz'

    return ntwk

## file conversion
def statistical_2_touchstone(file_name, new_file_name=None,\
        header_string='# GHz S RI R 50.0'):
    '''
    Converts Statistical file to a touchstone file.

    Converts the file format used by Statistical and other Dylan Williams
    software to standard touchstone format.

    Parameters
    ------------
    file_name : string
            name of file to convert
    new_file_name : string
            name of new file to write out (including extension)
    header_string : string
            touchstone header written to first beginning of file

    '''
    if new_file_name is None:
        new_file_name = 'tmp-'+file_name
        remove_tmp_file = True

    # This breaks compatibility with python 2.6 and older
    with file(file_name, 'r') as old_file, open(new_file_name, 'w') as new_file:
        new_file.write('%s\n'%header_string)
        for line in old_file:
            new_file.write(line)

    if remove_tmp_file is True:
        os.rename(new_file_name,file_name)

def network_2_spreadsheet(ntwk, file_name =None, file_type= 'excel', form='db',
    *args, **kwargs):
    '''
    Write a Network object to a spreadsheet, for your boss

    Write the s-parameters  of a network to a spreadsheet, in a variety
    of forms. This functions makes use of the pandas module, which in
    turn makes use of the xlrd module. These are imported during this
    function call. For more details about the file-writing functions
    see the pandas.DataFrom.to_?? functions.

    Notes
    ------
    The frequency unit used in the spreadsheet is take from
    `ntwk.frequency.unit`

    Parameters
    -----------
    ntwk :  :class:`~skrf.network.Network` object
        the network to write
    file_name : str, None
        the file_name to write. if None,  ntwk.name is used.
    file_type : ['csv','excel','html']
        the type of file to write. See pandas.DataFrame.to_??? functions.
    form : 'db','ma','ri'
        format to write data,
        * db = db, deg
        * ma = mag, deg
        * ri = real, imag
    \*args, \*\*kwargs :
        passed to pandas.DataFrame.to_??? functions.

    See Also
    ---------
    networkset_2_spreadsheet : writes a spreadsheet for many networks
    '''
    from pandas import DataFrame, Series # delayed because its not a requirement
    file_extns = {'csv':'csv','excel':'xls','html':'html'}

    form = form.lower()
    if form not in ['db','ri','ma']:
        raise ValueError('`form` must be either `db`,`ma`,`ri`')

    file_type = file_type.lower()
    if file_type not in file_extns.keys():
        raise ValueError('file_type must be `csv`,`html`,`excel` ')
    if ntwk.name is None and file_name is None:
        raise ValueError('Either ntwk must have name or give a file_name')

    if file_name is None and 'excel_writer' not in kwargs.keys():
        file_name = ntwk.name + '.'+file_extns[file_type]

    d = {}
    index =ntwk.frequency.f_scaled

    if form =='db':
        for m,n in ntwk.port_tuples:
            d['S%i%i Log Mag(dB)'%(m+1,n+1)] = \
                Series(ntwk.s_db[:,m,n], index = index)
            d[u'S%i%i Phase(deg)'%(m+1,n+1)] = \
                Series(ntwk.s_deg[:,m,n], index = index)
    elif form =='ma':
        for m,n in ntwk.port_tuples:
            d['S%i%i Mag(lin)'%(m+1,n+1)] = \
                Series(ntwk.s_mag[:,m,n], index = index)
            d[u'S%i%i Phase(deg)'%(m+1,n+1)] = \
                Series(ntwk.s_deg[:,m,n], index = index)
    elif form =='ri':
        for m,n in ntwk.port_tuples:
            d['S%i%i Real'%(m+1,n+1)] = \
                Series(ntwk.s_re[:,m,n], index = index)
            d[u'S%i%i Imag'%(m+1,n+1)] = \
                Series(ntwk.s_im[:,m,n], index = index)

    df = DataFrame(d)
    df.__getattribute__('to_%s'%file_type)(file_name,
        index_label='Freq(%s)'%ntwk.frequency.unit, *args, **kwargs)

def network_2_dataframe(ntwk, attrs=['s_db'], ports = None):
    '''
    Convert one or more attributes of a network to a pandas DataFrame

    Parameters
    --------------
    ntwk :  :class:`~skrf.network.Network` object
        the network to write
    attrs : list Network attributes
        like ['s_db','s_deg']
    ports : list of tuples
        list of port pairs to write. defaults to ntwk.port_tuples
        (like [[0,0]])

    Returns
    ----------
    df : pandas DataFrame Object
    '''
    from pandas import DataFrame, Series # delayed because its not a requirement
    d = {}
    index =ntwk.frequency.f_scaled

    if ports is None:
        ports = ntwk.port_tuples

    for attr in attrs:
        for m,n in ports:
            d['%s %i%i'%(attr, m+1,n+1)] = \
                Series(ntwk.__getattribute__(attr)[:,m,n], index = index)

    return DataFrame(d)

def networkset_2_spreadsheet(ntwkset, file_name=None, file_type= 'excel',
    *args, **kwargs):
    '''
    Write a NetworkSet object to a spreadsheet, for your boss

    Write  the s-parameters  of a each network in the networkset to a
    spreadsheet. If the `excel` file_type is used, then each network,
    is written to its own sheet, with the sheetname taken from the
    network `name` attribute.
    This functions makes use of the pandas module, which in turn makes
    use of the xlrd module. These are imported during this function

    Notes
    ------
    The frequency unit used in the spreadsheet is take from
    `ntwk.frequency.unit`

    Parameters
    -----------
    ntwkset :  :class:`~skrf.networkSet.NetworkSet` object
        the network to write
    file_name : str, None
        the file_name to write. if None,  ntwk.name is used.
    file_type : ['csv','excel','html']
        the type of file to write. See pandas.DataFrame.to_??? functions.
    form : 'db','ma','ri'
        format to write data,
        * db = db, deg
        * ma = mag, deg
        * ri = real, imag
    \*args, \*\*kwargs :
        passed to pandas.DataFrame.to_??? functions.

    See Also
    ---------
    networkset_2_spreadsheet : writes a spreadsheet for many networks
    '''
    from pandas import DataFrame, Series, ExcelWriter # delayed because its not a requirement
    if ntwkset.name is None and file_name is None:
        raise(ValueError('Either ntwkset must have name or give a file_name'))

    if file_type == 'excel':
        writer = ExcelWriter(file_name)
        [network_2_spreadsheet(k, writer, sheet_name =k.name, *args, **kwargs) for k in ntwkset]
        writer.save()
    else:
        [network_2_spreadsheet(k,*args, **kwargs) for k in ntwkset]

# Provide a StringBuffer that let's me work with Python2 strings and Python3 unicode strings without thinking
if sys.version_info < (3, 0):
    import StringIO

    class StringBuffer(StringIO.StringIO):
        def __enter__(self):
            return self

        def __exit__(self, *args):
            self.close()
else:
    import io
    StringBuffer = io.StringIO


import os
from collections import namedtuple
import pandas as pd

from bcbio import utils
import bcbio.pipeline.datadict as dd
import bcbio.rnaseq.gtf as gtf
from bcbio.distributed.transaction import file_transaction
from bcbio.provenance import do
from bcbio.utils import (file_exists, safe_makedir, is_gzipped,
                         R_package_path, Rscript_cmd)
from bcbio.pipeline import config_utils, disambiguate
from bcbio.bam import fastq
from bcbio import bam

def run_sailfish(data):
    samplename = dd.get_sample_name(data)
    files = dd.get_input_sequence_files(data)
    work_dir = dd.get_work_dir(data)
    if len(files) == 2:
        fq1, fq2 = files
    else:
        fq1, fq2 = files[0], None
    if not fastq.is_fastq(fq1):
        return [[data]]
    sailfish_dir = os.path.join(work_dir, "sailfish", samplename)
    gtf_file = dd.get_gtf_file(data)
    assert file_exists(gtf_file), "%s was not found, exiting." % gtf_file
    fasta_file = dd.get_ref_file(data)
    assert file_exists(fasta_file), "%s was not found, exiting." % fasta_file
    stranded = dd.get_strandedness(data).lower()
    out_file = sailfish(fq1, fq2, sailfish_dir, gtf_file, fasta_file, stranded, data)
    data = dd.set_sailfish(data, out_file)
    data = dd.set_sailfish_dir(data, sailfish_dir)
    return [[data]]

def sailfish(fq1, fq2, sailfish_dir, gtf_file, ref_file, strandedness, data):
    safe_makedir(sailfish_dir)
    samplename = dd.get_sample_name(data)
    quant_dir = os.path.join(sailfish_dir, "quant")
    out_file = os.path.join(quant_dir, "quant.sf")
    if file_exists(out_file):
        return out_file
    build_string = get_build_string(data)
    sailfish_idx = sailfish_index(ref_file, gtf_file, data, build_string)
    num_cores = dd.get_num_cores(data)
    sailfish = config_utils.get_program("sailfish", data["config"])
    cmd = "{sailfish} quant -i {sailfish_idx} -p {num_cores} "
    cmd += _libtype_string(fq1, fq2, strandedness)
    fq1_cmd = "{fq1}" if not is_gzipped(fq1) else "<(gzip -cd {fq1})"
    fq1_cmd = fq1_cmd.format(fq1=fq1)
    if not fq2:
        cmd += " -r {fq1_cmd} "
    else:
        fq2_cmd = "{fq2}" if not is_gzipped(fq2) else "<(gzip -cd {fq2})"
        fq2_cmd = fq2_cmd.format(fq2=fq2)
        cmd += " -1 {fq1_cmd} -2 {fq2_cmd} "
    cmd += "--useVBOpt --numBootstraps 30 "
    cmd += "-o {tx_out_dir}"
    message = "Quantifying transcripts in {fq1} and {fq2}."
    with file_transaction(data, quant_dir) as tx_out_dir:
        do.run(cmd.format(**locals()), message.format(**locals()), None)
        sleuthify_sailfish(tx_out_dir)
    return out_file

def sleuthify_sailfish(sailfish_dir):
    """
    if installed, use wasabi to create abundance.h5 output for use with
    sleuth
    """
    if not R_package_path("wasabi"):
        return None
    else:
        rscript = Rscript_cmd()
        cmd = """{rscript} --vanilla -e 'library("wasabi"); prepare_fish_for_sleuth(c("{sailfish_dir}"))'"""
        do.run(cmd.format(**locals()), "Converting Sailfish to Sleuth format.")
    return os.path.join(sailfish_dir, "abundance.h5")

def create_combined_fasta(data):
    """
    if there are genomes to be disambiguated, create a FASTA file of
    all of the transcripts for all genomes
    """
    out_dir = os.path.join(dd.get_work_dir(data), "inputs", "transcriptome")
    items = disambiguate.split([data])
    fasta_files = []
    for i in items:
        odata = i[0]
        gtf_file = dd.get_gtf_file(odata)
        ref_file = dd.get_ref_file(odata)
        out_file = os.path.join(out_dir, dd.get_genome_build(odata) + ".fa")
        if file_exists(out_file):
            fasta_files.append(out_file)
        else:
            out_file = gtf.gtf_to_fasta(gtf_file, ref_file, out_file=out_file)
            fasta_files.append(out_file)
    out_stem = os.path.join(out_dir, dd.get_genome_build(data))
    if dd.get_disambiguate(data):
        out_stem = "-".join([out_stem] + (dd.get_disambiguate(data) or []))
    combined_file = out_stem + ".fa"
    if file_exists(combined_file):
        return combined_file

    fasta_file_string = " ".join(fasta_files)
    cmd = "cat {fasta_file_string} > {tx_out_file}"
    with file_transaction(data, combined_file) as tx_out_file:
        do.run(cmd.format(**locals()), "Combining transcriptome FASTA files.")
    return combined_file

def create_combined_tx2gene(data):
    out_dir = os.path.join(dd.get_work_dir(data), "inputs", "transcriptome")
    items = disambiguate.split([data])
    tx2gene_files = []
    for i in items:
        odata = i[0]
        gtf_file = dd.get_transcriptome_gtf(odata)
        if not gtf_file:
            gtf_file = dd.get_gtf_file(odata)
        out_file = os.path.join(out_dir, dd.get_genome_build(odata) + "-tx2gene.csv")
        if file_exists(out_file):
            tx2gene_files.append(out_file)
        else:
            out_file = gtf.tx2genefile(gtf_file, out_file, tsv=False)
            tx2gene_files.append(out_file)
    combined_file = os.path.join(out_dir, "tx2gene.csv")
    if file_exists(combined_file):
        return combined_file

    tx2gene_file_string = " ".join(tx2gene_files)
    cmd = "cat {tx2gene_file_string} > {tx_out_file}"
    with file_transaction(data, combined_file) as tx_out_file:
        do.run(cmd.format(**locals()), "Combining tx2gene CSV files.")
    return combined_file

def get_build_string(data):
    build_string = dd.get_genome_build(data)
    if dd.get_disambiguate(data):
        build_string = "-".join([build_string] + (dd.get_disambiguate(data) or []))
    return build_string

def run_sailfish_index(*samples):
    samples = [utils.to_single_data(x) for x in samples]
    Build = namedtuple('Build', ['build', 'ref', 'gtf'])
    builds = {Build(get_build_string(x), dd.get_ref_file(x), dd.get_gtf_file(x))
              for x in samples}
    data = samples[0]
    indexdirs = {}
    for build in builds:
        indexdirs[build.build] = sailfish_index(build.ref, build.gtf, data,
                                                build.build)
    return [[x] for x in samples]

def sailfish_index(gtf_file, ref_file, data, build):
    work_dir = dd.get_work_dir(data)
    out_dir = os.path.join(work_dir, "sailfish", "index", build)
    sailfish = config_utils.get_program("sailfish", data["config"])
    num_cores = dd.get_num_cores(data)
    gtf_fa = create_combined_fasta(data)
    if file_exists(os.path.join(out_dir, "versionInfo.json")):
        return out_dir
    with file_transaction(data, out_dir) as tx_out_dir:
        fq1, _ = dd.get_input_sequence_files(data)
        kmersize = pick_kmersize(fq1)
        cmd = ("{sailfish} index -p {num_cores} -t {gtf_fa} -o {tx_out_dir} "
               "-k {kmersize}")
        message = "Creating sailfish index for {gtf_fa} with {kmersize} bp kmers."
        do.run(cmd.format(**locals()), message.format(**locals()), None)
    return out_dir

def _libtype_string(fq1, fq2, strandedness):
    """
    supports just the Tophat unstranded/firstrand/secondstrand
    """
    libtype = "-l I" if fq2 else "-l "
    strand = _sailfish_strand_string(strandedness)
    return libtype + strand

def _sailfish_strand_string(strandedness):
    return {'unstranded': "U",
            'firststrand': "SR",
            'secondstrand': "SF"}.get(strandedness, "U")

def _sailfish_expression_parser(sailfish_file, samplename):
    col_names = ["name", "length", "effectiveLength", "tpm", "numreads"]
    df = pd.read_csv(sailfish_file, comment="#", header=None, skiprows=1, index_col=0,
                     names=col_names, sep="\t")
    df["sample"] = samplename
    return df

def pick_kmersize(fq):
    """
    pick an appropriate kmer size based off of https://www.biostars.org/p/201474/
    tl;dr version: pick 31 unless the reads are very small, if not then guess
    that readlength / 2 is about right.
    """
    if bam.is_bam(fq):
        readlength = bam.estimate_read_length(fq)
    else:
        readlength = fastq.estimate_read_length(fq)
    halfread = int(round(readlength / 2))
    if halfread >= 31:
        kmersize = 31
    else:
        kmersize = halfread
    if kmersize % 2 == 0:
        kmersize += 1
    return kmersize


import argparse
import numpy as np
from sklearn.metrics import precision_score, recall_score, f1_score

_ALL_ = ["random", "decision_trees", "linear_svm", "gaussian_naive_bayes"]

def precision_recall_f1score(GT, Y_pred):
	precisions, recalls, f1scores = [], [], []
	for i in range(GT.shape[0]):
		precisions.append(precision_score(GT[i][:], Y_pred[i][:]))
		recalls.append(recall_score(GT[i][:], Y_pred[i][:]))
		f1scores.append(f1_score(GT[i][:], Y_pred[i][:]))
	precisions, recalls, f1scores = np.array(precisions), np.array(recalls), np.array(f1scores)
	precision = (np.mean(precisions), np.median(precisions), np.std(precisions))
	recall = (np.mean(recalls), np.median(recalls), np.std(recalls))
	f1score = (np.mean(f1scores), np.median(f1scores), np.std(f1scores))
	return precision, recall, f1score

def annotation_evaluation(collection, algorithm, model="track_based"):
	print "\n\n####Evaluating %s annotation predictions by algorithm '%s'####" % (model, algorithm)
	evaluation = {}
	filename = "classification/%s/confusion_matrix__%s.out" % (collection, algorithm)
	Y_pred = np.loadtxt(open(filename))
	GT = np.loadtxt(open("classification/%s/ground_truth.out" % collection))
	if model == "tag_based":
		GT = GT.transpose()
		Y_pred = Y_pred.transpose()
	precision, recall, f1score = precision_recall_f1score(GT, Y_pred)
	print "-----------------------------------------------------------"
	print "Precision: %s" % str(precision)
	print "Recall: %s" % str(recall)
	print "F1score: %s" % str(f1score)
	print "-----------------------------------------------------------"

if __name__ == '__main__':
	parser = argparse.ArgumentParser(description='Evaluate annotation predictions using track_based and tag_based measures')
	parser.add_argument('collection', help='Collection name (e.g.: majorminer)')
	parser.add_argument('-a', '--algorithm', nargs='?', default="all", help='algorithm to evaluate (default="all")')
	args = parser.parse_args()

	if args.algorithm == "all":
		for algorithm in _ALL_:
			annotation_evaluation(args.collection, algorithm, "track_based")
			annotation_evaluation(args.collection, algorithm, "tag_based")
	else:
		annotation_evaluation(args.collection, args.algorithm, "track_based")
		annotation_evaluation(args.collection, args.algorithm, "tag_based")


"""
==================================================
Probability Calibration for 3-class classification
==================================================

This example illustrates how sigmoid calibration changes predicted
probabilities for a 3-class classification problem. Illustrated is the
standard 2-simplex, where the three corners correspond to the three classes.
Arrows point from the probability vectors predicted by an uncalibrated
classifier to the probability vectors predicted by the same classifier after
sigmoid calibration on a hold-out validation set. Colors indicate the true
class of an instance (red: class 1, green: class 2, blue: class 3).

The base classifier is a random forest classifier with 25 base estimators
(trees). If this classifier is trained on all 800 training datapoints, it is
overly confident in its predictions and thus incurs a large log-loss.
Calibrating an identical classifier, which was trained on 600 datapoints, with
method='sigmoid' on the remaining 200 datapoints reduces the confidence of the
predictions, i.e., moves the probability vectors from the edges of the simplex
towards the center. This calibration results in a lower log-loss. Note that an
alternative would have been to increase the number of base estimators which
would have resulted in a similar decrease in log-loss.
"""
print(__doc__)

# Author: Jan Hendrik Metzen <jhm@informatik.uni-bremen.de>
# License: BSD Style.

import matplotlib.pyplot as plt

import numpy as np

from sklearn.datasets import make_blobs
from sklearn.ensemble import RandomForestClassifier
from sklearn.calibration import CalibratedClassifierCV
from sklearn.metrics import log_loss

np.random.seed(0)

# Generate data
X, y = make_blobs(n_samples=1000, n_features=2, random_state=42,
                  cluster_std=5.0)
X_train, y_train = X[:600], y[:600]
X_valid, y_valid = X[600:800], y[600:800]
X_train_valid, y_train_valid = X[:800], y[:800]
X_test, y_test = X[800:], y[800:]

# Train uncalibrated random forest classifier on whole train and validation
# data and evaluate on test data
clf = RandomForestClassifier(n_estimators=25)
clf.fit(X_train_valid, y_train_valid)
clf_probs = clf.predict_proba(X_test)
score = log_loss(y_test, clf_probs)

# Train random forest classifier, calibrate on validation data and evaluate
# on test data
clf = RandomForestClassifier(n_estimators=25)
clf.fit(X_train, y_train)
clf_probs = clf.predict_proba(X_test)
sig_clf = CalibratedClassifierCV(clf, method="sigmoid", cv="prefit")
sig_clf.fit(X_valid, y_valid)
sig_clf_probs = sig_clf.predict_proba(X_test)
sig_score = log_loss(y_test, sig_clf_probs)

# Plot changes in predicted probabilities via arrows
plt.figure(0)
colors = ["r", "g", "b"]
for i in range(clf_probs.shape[0]):
    plt.arrow(clf_probs[i, 0], clf_probs[i, 1],
              sig_clf_probs[i, 0] - clf_probs[i, 0],
              sig_clf_probs[i, 1] - clf_probs[i, 1],
              color=colors[y_test[i]], head_width=1e-2)

# Plot perfect predictions
plt.plot([1.0], [0.0], 'ro', ms=20, label="Class 1")
plt.plot([0.0], [1.0], 'go', ms=20, label="Class 2")
plt.plot([0.0], [0.0], 'bo', ms=20, label="Class 3")

# Plot boundaries of unit simplex
plt.plot([0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], 'k', label="Simplex")

# Annotate points on the simplex
plt.annotate(r'($\frac{1}{3}$, $\frac{1}{3}$, $\frac{1}{3}$)',
             xy=(1.0/3, 1.0/3), xytext=(1.0/3, .23), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.plot([1.0/3], [1.0/3], 'ko', ms=5)
plt.annotate(r'($\frac{1}{2}$, $0$, $\frac{1}{2}$)',
             xy=(.5, .0), xytext=(.5, .1), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.annotate(r'($0$, $\frac{1}{2}$, $\frac{1}{2}$)',
             xy=(.0, .5), xytext=(.1, .5), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.annotate(r'($\frac{1}{2}$, $\frac{1}{2}$, $0$)',
             xy=(.5, .5), xytext=(.6, .6), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.annotate(r'($0$, $0$, $1$)',
             xy=(0, 0), xytext=(.1, .1), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.annotate(r'($1$, $0$, $0$)',
             xy=(1, 0), xytext=(1, .1), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
plt.annotate(r'($0$, $1$, $0$)',
             xy=(0, 1), xytext=(.1, 1), xycoords='data',
             arrowprops=dict(facecolor='black', shrink=0.05),
             horizontalalignment='center', verticalalignment='center')
# Add grid
plt.grid("off")
for x in [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]:
    plt.plot([0, x], [x, 0], 'k', alpha=0.2)
    plt.plot([0, 0 + (1-x)/2], [x, x + (1-x)/2], 'k', alpha=0.2)
    plt.plot([x, x + (1-x)/2], [0, 0 + (1-x)/2], 'k', alpha=0.2)

plt.title("Change of predicted probabilities after sigmoid calibration")
plt.xlabel("Probability class 1")
plt.ylabel("Probability class 2")
plt.xlim(-0.05, 1.05)
plt.ylim(-0.05, 1.05)
plt.legend(loc="best")

print("Log-loss of")
print(" * uncalibrated classifier trained on 800 datapoints: %.3f "
      % score)
print(" * classifier trained on 600 datapoints and calibrated on "
      "200 datapoint: %.3f" % sig_score)

# Illustrate calibrator
plt.figure(1)
# generate grid over 2-simplex
p1d = np.linspace(0, 1, 20)
p0, p1 = np.meshgrid(p1d, p1d)
p2 = 1 - p0 - p1
p = np.c_[p0.ravel(), p1.ravel(), p2.ravel()]
p = p[p[:, 2] >= 0]

calibrated_classifier = sig_clf.calibrated_classifiers_[0]
prediction = np.vstack([calibrator.predict(this_p)
                        for calibrator, this_p in
                        zip(calibrated_classifier.calibrators_, p.T)]).T
prediction /= prediction.sum(axis=1)[:, None]

# Ploit modifications of calibrator
for i in range(prediction.shape[0]):
    plt.arrow(p[i, 0], p[i, 1],
              prediction[i, 0] - p[i, 0], prediction[i, 1] - p[i, 1],
              head_width=1e-2, color=colors[np.argmax(p[i])])
# Plot boundaries of unit simplex
plt.plot([0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], 'k', label="Simplex")

plt.grid("off")
for x in [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]:
    plt.plot([0, x], [x, 0], 'k', alpha=0.2)
    plt.plot([0, 0 + (1-x)/2], [x, x + (1-x)/2], 'k', alpha=0.2)
    plt.plot([x, x + (1-x)/2], [0, 0 + (1-x)/2], 'k', alpha=0.2)

plt.title("Illustration of sigmoid calibrator")
plt.xlabel("Probability class 1")
plt.ylabel("Probability class 2")
plt.xlim(-0.05, 1.05)
plt.ylim(-0.05, 1.05)

plt.show()


"""
=============================================
Feature Union with Heterogeneous Data Sources
=============================================

Datasets can often contain components of that require different feature
extraction and processing pipelines.  This scenario might occur when:

1. Your dataset consists of heterogeneous data types (e.g. raster images and
   text captions)
2. Your dataset is stored in a Pandas DataFrame and different columns
   require different processing pipelines.

This example demonstrates how to use
:class:`sklearn.feature_extraction.FeatureUnion` on a dataset containing
different types of features.  We use the 20-newsgroups dataset and compute
standard bag-of-words features for the subject line and body in separate
pipelines as well as ad hoc features on the body. We combine them (with
weights) using a FeatureUnion and finally train a classifier on the combined
set of features.

The choice of features is not particularly helpful, but serves to illustrate
the technique.
"""

# Author: Matt Terry <matt.terry@gmail.com>
#
# License: BSD 3 clause
from __future__ import print_function

import numpy as np

from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.datasets import fetch_20newsgroups
from sklearn.datasets.twenty_newsgroups import strip_newsgroup_footer
from sklearn.datasets.twenty_newsgroups import strip_newsgroup_quoting
from sklearn.decomposition import TruncatedSVD
from sklearn.feature_extraction import DictVectorizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import classification_report
from sklearn.pipeline import FeatureUnion
from sklearn.pipeline import Pipeline
from sklearn.svm import SVC

class ItemSelector(BaseEstimator, TransformerMixin):
    """For data grouped by feature, select subset of data at a provided key.

    The data is expected to be stored in a 2D data structure, where the first
    index is over features and the second is over samples.  i.e.

    >> len(data[key]) == n_samples

    Please note that this is the opposite convention to sklearn feature
    matrixes (where the first index corresponds to sample).

    ItemSelector only requires that the collection implement getitem
    (data[key]).  Examples include: a dict of lists, 2D numpy array, Pandas
    DataFrame, numpy record array, etc.

    >> data = {'a': [1, 5, 2, 5, 2, 8],
               'b': [9, 4, 1, 4, 1, 3]}
    >> ds = ItemSelector(key='a')
    >> data['a'] == ds.transform(data)

    ItemSelector is not designed to handle data grouped by sample.  (e.g. a
    list of dicts).  If your data is structured this way, consider a
    transformer along the lines of `sklearn.feature_extraction.DictVectorizer`.

    Parameters
    ----------
    key : hashable, required
        The key corresponding to the desired value in a mappable.
    """
    def __init__(self, key):
        self.key = key

    def fit(self, x, y=None):
        return self

    def transform(self, data_dict):
        return data_dict[self.key]

class TextStats(BaseEstimator, TransformerMixin):
    """Extract features from each document for DictVectorizer"""

    def fit(self, x, y=None):
        return self

    def transform(self, posts):
        return [{'length': len(text),
                 'num_sentences': text.count('.')}
                for text in posts]

class SubjectBodyExtractor(BaseEstimator, TransformerMixin):
    """Extract the subject & body from a usenet post in a single pass.

    Takes a sequence of strings and produces a dict of sequences.  Keys are
    `subject` and `body`.
    """
    def fit(self, x, y=None):
        return self

    def transform(self, posts):
        features = np.recarray(shape=(len(posts),),
                               dtype=[('subject', object), ('body', object)])
        for i, text in enumerate(posts):
            headers, _, bod = text.partition('\n\n')
            bod = strip_newsgroup_footer(bod)
            bod = strip_newsgroup_quoting(bod)
            features['body'][i] = bod

            prefix = 'Subject:'
            sub = ''
            for line in headers.split('\n'):
                if line.startswith(prefix):
                    sub = line[len(prefix):]
                    break
            features['subject'][i] = sub

        return features

pipeline = Pipeline([
    # Extract the subject & body
    ('subjectbody', SubjectBodyExtractor()),

    # Use FeatureUnion to combine the features from subject and body
    ('union', FeatureUnion(
        transformer_list=[

            # Pipeline for pulling features from the post's subject line
            ('subject', Pipeline([
                ('selector', ItemSelector(key='subject')),
                ('tfidf', TfidfVectorizer(min_df=50)),
            ])),

            # Pipeline for standard bag-of-words model for body
            ('body_bow', Pipeline([
                ('selector', ItemSelector(key='body')),
                ('tfidf', TfidfVectorizer()),
                ('best', TruncatedSVD(n_components=50)),
            ])),

            # Pipeline for pulling ad hoc features from post's body
            ('body_stats', Pipeline([
                ('selector', ItemSelector(key='body')),
                ('stats', TextStats()),  # returns a list of dicts
                ('vect', DictVectorizer()),  # list of dicts -> feature matrix
            ])),

        ],

        # weight components in FeatureUnion
        transformer_weights={
            'subject': 0.8,
            'body_bow': 0.5,
            'body_stats': 1.0,
        },
    )),

    # Use a SVC classifier on the combined features
    ('svc', SVC(kernel='linear')),
])

# limit the list of categories to make running this exmaple faster.
categories = ['alt.atheism', 'talk.religion.misc']
train = fetch_20newsgroups(random_state=1,
                           subset='train',
                           categories=categories,
                           )
test = fetch_20newsgroups(random_state=1,
                          subset='test',
                          categories=categories,
                          )

pipeline.fit(train.data, train.target)
y = pipeline.predict(test.data)
print(classification_report(y, test.target))


"""
============================================================
Single estimator versus bagging: bias-variance decomposition
============================================================

This example illustrates and compares the bias-variance decomposition of the
expected mean squared error of a single estimator against a bagging ensemble.

In regression, the expected mean squared error of an estimator can be
decomposed in terms of bias, variance and noise. On average over datasets of
the regression problem, the bias term measures the average amount by which the
predictions of the estimator differ from the predictions of the best possible
estimator for the problem (i.e., the Bayes model). The variance term measures
the variability of the predictions of the estimator when fit over different
instances LS of the problem. Finally, the noise measures the irreducible part
of the error which is due the variability in the data.

The upper left figure illustrates the predictions (in dark red) of a single
decision tree trained over a random dataset LS (the blue dots) of a toy 1d
regression problem. It also illustrates the predictions (in light red) of other
single decision trees trained over other (and different) randomly drawn
instances LS of the problem. Intuitively, the variance term here corresponds to
the width of the beam of predictions (in light red) of the individual
estimators. The larger the variance, the more sensitive are the predictions for
`x` to small changes in the training set. The bias term corresponds to the
difference between the average prediction of the estimator (in cyan) and the
best possible model (in dark blue). On this problem, we can thus observe that
the bias is quite low (both the cyan and the blue curves are close to each
other) while the variance is large (the red beam is rather wide).

The lower left figure plots the pointwise decomposition of the expected mean
squared error of a single decision tree. It confirms that the bias term (in
blue) is low while the variance is large (in green). It also illustrates the
noise part of the error which, as expected, appears to be constant and around
`0.01`.

The right figures correspond to the same plots but using instead a bagging
ensemble of decision trees. In both figures, we can observe that the bias term
is larger than in the previous case. In the upper right figure, the difference
between the average prediction (in cyan) and the best possible model is larger
(e.g., notice the offset around `x=2`). In the lower right figure, the bias
curve is also slightly higher than in the lower left figure. In terms of
variance however, the beam of predictions is narrower, which suggests that the
variance is lower. Indeed, as the lower right figure confirms, the variance
term (in green) is lower than for single decision trees. Overall, the bias-
variance decomposition is therefore no longer the same. The tradeoff is better
for bagging: averaging several decision trees fit on bootstrap copies of the
dataset slightly increases the bias term but allows for a larger reduction of
the variance, which results in a lower overall mean squared error (compare the
red curves int the lower figures). The script output also confirms this
intuition. The total error of the bagging ensemble is lower than the total
error of a single decision tree, and this difference indeed mainly stems from a
reduced variance.

For further details on bias-variance decomposition, see section 7.3 of [1]_.

References
----------

.. [1] T. Hastie, R. Tibshirani and J. Friedman,
       "Elements of Statistical Learning", Springer, 2009.

"""
print(__doc__)

# Author: Gilles Louppe <g.louppe@gmail.com>
# License: BSD 3 clause

import numpy as np
import matplotlib.pyplot as plt

from sklearn.ensemble import BaggingRegressor
from sklearn.tree import DecisionTreeRegressor

# Settings
n_repeat = 50       # Number of iterations for computing expectations
n_train = 50        # Size of the training set
n_test = 1000       # Size of the test set
noise = 0.1         # Standard deviation of the noise
np.random.seed(0)

# Change this for exploring the bias-variance decomposition of other
# estimators. This should work well for estimators with high variance (e.g.,
# decision trees or KNN), but poorly for estimators with low variance (e.g.,
# linear models).
estimators = [("Tree", DecisionTreeRegressor()),
              ("Bagging(Tree)", BaggingRegressor(DecisionTreeRegressor()))]

n_estimators = len(estimators)

# Generate data
def f(x):
    x = x.ravel()

    return np.exp(-x ** 2) + 1.5 * np.exp(-(x - 2) ** 2)

def generate(n_samples, noise, n_repeat=1):
    X = np.random.rand(n_samples) * 10 - 5
    X = np.sort(X)

    if n_repeat == 1:
        y = f(X) + np.random.normal(0.0, noise, n_samples)
    else:
        y = np.zeros((n_samples, n_repeat))

        for i in range(n_repeat):
            y[:, i] = f(X) + np.random.normal(0.0, noise, n_samples)

    X = X.reshape((n_samples, 1))

    return X, y

X_train = []
y_train = []

for i in range(n_repeat):
    X, y = generate(n_samples=n_train, noise=noise)
    X_train.append(X)
    y_train.append(y)

X_test, y_test = generate(n_samples=n_test, noise=noise, n_repeat=n_repeat)

# Loop over estimators to compare
for n, (name, estimator) in enumerate(estimators):
    # Compute predictions
    y_predict = np.zeros((n_test, n_repeat))

    for i in range(n_repeat):
        estimator.fit(X_train[i], y_train[i])
        y_predict[:, i] = estimator.predict(X_test)

    # Bias^2 + Variance + Noise decomposition of the mean squared error
    y_error = np.zeros(n_test)

    for i in range(n_repeat):
        for j in range(n_repeat):
            y_error += (y_test[:, j] - y_predict[:, i]) ** 2

    y_error /= (n_repeat * n_repeat)

    y_noise = np.var(y_test, axis=1)
    y_bias = (f(X_test) - np.mean(y_predict, axis=1)) ** 2
    y_var = np.var(y_predict, axis=1)

    print("{0}: {1:.4f} (error) = {2:.4f} (bias^2) "
          " + {3:.4f} (var) + {4:.4f} (noise)".format(name,
                                                      np.mean(y_error),
                                                      np.mean(y_bias),
                                                      np.mean(y_var),
                                                      np.mean(y_noise)))

    # Plot figures
    plt.subplot(2, n_estimators, n + 1)
    plt.plot(X_test, f(X_test), "b", label="$f(x)$")
    plt.plot(X_train[0], y_train[0], ".b", label="LS ~ $y = f(x)+noise$")

    for i in range(n_repeat):
        if i == 0:
            plt.plot(X_test, y_predict[:, i], "r", label="$\^y(x)$")
        else:
            plt.plot(X_test, y_predict[:, i], "r", alpha=0.05)

    plt.plot(X_test, np.mean(y_predict, axis=1), "c",
             label="$\mathbb{E}_{LS} \^y(x)$")

    plt.xlim([-5, 5])
    plt.title(name)

    if n == 0:
        plt.legend(loc="upper left", prop={"size": 11})

    plt.subplot(2, n_estimators, n_estimators + n + 1)
    plt.plot(X_test, y_error, "r", label="$error(x)$")
    plt.plot(X_test, y_bias, "b", label="$bias^2(x)$"),
    plt.plot(X_test, y_var, "g", label="$variance(x)$"),
    plt.plot(X_test, y_noise, "c", label="$noise(x)$")

    plt.xlim([-5, 5])
    plt.ylim([0, 0.1])

    if n == 0:
        plt.legend(loc="upper left", prop={"size": 11})

plt.show()


# coding=utf-8
# Copyright 2021 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Runner for transformer experiments.
"""
import os

from . import metrics
from . import plotting
from . import ring_dist
from . import ring_models
from . import train

from absl import app
from absl import flags
import jax
from jax.config import config
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as onp

flags.DEFINE_integer("num_encoders", 6,
                     "Number of encoder modules in the transformer.")
flags.DEFINE_integer("num_decoders", 6,
                     "Number of decoder modules in the transformer.")
flags.DEFINE_integer("num_heads", 8,
                     "Number of attention heads in the transformer.")
flags.DEFINE_integer("key_dim", 32,
                     "The dimension of the keys in the transformer.")
flags.DEFINE_integer("value_dim_per_head", 32,
                     "The dimension of the values in the transformer for each head.")
flags.DEFINE_integer("k", 2,
                     "The number of modes in the data.")
flags.DEFINE_integer("data_points_per_mode", 25,
                     "Number of data points to include per mode in the data.")
flags.DEFINE_boolean("parallel", True,
                     "If possible, train in parallel across devices.")
flags.DEFINE_integer("batch_size", 64,
                     "The batch size.")
flags.DEFINE_integer("eval_batch_size", 256,
                     "The batch size for evaluation.")
flags.DEFINE_integer("num_steps", int(1e6),
                     "The number of steps to train for.")
flags.DEFINE_float("lr", 1e-3,
                   "The learning rate for ADAM.")
flags.DEFINE_integer("summarize_every", 100,
                     "Number of steps between summaries.")
flags.DEFINE_integer("checkpoint_every", 5000,
                     "Number of steps between checkpoints.")
flags.DEFINE_boolean("clobber_checkpoint", False,
                     "If true, remove any existing summaries and checkpoints in logdir.")
flags.DEFINE_string("logdir", "/tmp/transformer",
                    "The directory to put summaries and checkpoints.")
flags.DEFINE_boolean("debug_nans", False,
                     "If true, run in debug mode and fail on nans.")

FLAGS = flags.FLAGS

def make_model(key,
               num_encoders=4,
               num_decoders=4,
               num_heads=8,
               value_dim=128,
               data_points_per_mode=25,
               k=10):

  model = ring_models.RingInferenceMachine(
      max_k=k,
      max_num_data_points=k*data_points_per_mode, num_heads=num_heads,
      num_encoders=num_encoders, num_decoders=num_decoders, qkv_dim=value_dim)
  params = model.init_params(key)

  return model, params

def sample_batch(key, batch_size, k, data_points_per_mode):
  keys = jax.random.split(key, num=batch_size)
  xs, cs, params = jax.vmap(
      ring_dist.sample_params_and_points,
      in_axes=(0, None, None, None, None, None, None, None, None,
               None))(keys, k * data_points_per_mode, k, 1., 0.5, 2, .02,
                      jnp.zeros([2]), jnp.eye(2), 0.1)
  return xs, cs, params

def make_loss(model,
              k=2,
              data_points_per_mode=25,
              batch_size=128):

  def sample_train_batch(key):
    xs, _, params = sample_batch(key, batch_size, k, data_points_per_mode)
    return xs, params

  def loss(params, key):
    key, subkey = jax.random.split(key)
    xs, ring_params = sample_train_batch(key)
    ks = jnp.full([batch_size], k)
    losses = model.loss(
        params, xs, ks*data_points_per_mode, ring_params, ks, subkey)
    return jnp.mean(losses)

  return jax.jit(loss)

def make_summarize(
    model,
    k=2,
    data_points_per_mode=25,
    eval_batch_size=256):

  def sample_eval_batch(key):
    return sample_batch(key, eval_batch_size, k, data_points_per_mode)

  sample_eval_batch = jax.jit(sample_eval_batch)

  def sample_single(key):
    xs, cs, params = sample_batch(key, 1, k, data_points_per_mode)
    return xs[0], cs[0], (params[0][0], params[1][0], params[2][0],
                          params[3][0])

  def model_classify(params, inputs, batch_size):
    return model.classify(params, inputs,
                          jnp.full([batch_size], k*data_points_per_mode),
                          jnp.full([batch_size], k))

  def sample_and_classify_eval_batch(key, params):
    xs, cs, true_ring_params = sample_eval_batch(key)
    tfmr_cs, tfmr_ring_params = model_classify(params, xs, eval_batch_size)
    return xs, cs, true_ring_params, tfmr_cs, tfmr_ring_params

  def sample_and_classify_single_mm(key, params):
    xs, cs, ring_params = sample_single(key)
    tfmr_cs, tfmr_ring_params = model_classify(params, xs[jnp.newaxis], 1)
    return xs, cs, ring_params, tfmr_cs, tfmr_ring_params

  sample_and_classify_eval_batch = jax.jit(sample_and_classify_eval_batch)

  sample_and_classify_single_mm= jax.jit(sample_and_classify_single_mm)

  def summarize_baselines(writer, step, key):
    key, subkey = jax.random.split(key)
    xs, cs, _ = sample_eval_batch(subkey)
    ks = onp.full([eval_batch_size], k)
    baseline_metrics = metrics.compute_masked_baseline_metrics(
        xs, cs, ks, ks*data_points_per_mode)
    for method_name, method_metrics in baseline_metrics.items():
      for metric_name, metric_val in method_metrics.items():
        writer.scalar("%s/%s" % (method_name, metric_name),
                      metric_val, step=step)
        print("%s %s: %0.3f" % (method_name, metric_name, metric_val))

  def plot_params(num_data_points, writer, step, params, key):
    outs = sample_and_classify_single_mm(key, params)
    xs, true_cs, true_params, pred_cs, pred_params = outs
    pred_cs = pred_cs[0]
    pred_params = (pred_params[0][0], pred_params[1][0],
                   pred_params[2][0], pred_params[3][0])
    fig = plotting.plot_rings(
        xs, k, true_cs, true_params, pred_cs, pred_params)
    plot_image = plotting.plot_to_numpy_image(plt)
    writer.image(
        "%d_modes_%d_points" % (k, num_data_points), plot_image, step=step)
    plt.close(fig)

  def comparison_inference(params):
    rings_inputs, true_cs = plotting.make_comparison_rings()
    rings_inputs = rings_inputs[jnp.newaxis, Ellipsis]
    new_model = ring_models.RingInferenceMachine(
        max_k=2, max_num_data_points=1500, num_heads=FLAGS.num_heads,
        num_encoders=FLAGS.num_encoders, num_decoders=FLAGS.num_decoders,
        qkv_dim=FLAGS.value_dim_per_head*FLAGS.num_heads)
    pred_cs, pred_params = new_model.classify(
        params, rings_inputs, jnp.array([1500]), jnp.array([2]))
    pred_cs = pred_cs[0]
    pred_params = (pred_params[0][0], pred_params[1][0],
                   pred_params[2][0], pred_params[3][0])
    return rings_inputs[0], true_cs, pred_cs, pred_params

  comparison_inference = jax.jit(comparison_inference)

  def plot_sklearn_comparison(writer, step, params):
    ring_xs, true_cs, pred_cs, pred_params = comparison_inference(params)
    fig = plotting.plot_comparison_rings(ring_xs, true_cs, pred_cs, pred_params)
    writer.image(
        "sklearn_comparison", plotting.plot_to_numpy_image(plt), step=step)
    plt.close(fig)

  def summarize(writer, step, params, key):
    k1, k2, k3 = jax.random.split(key, num=3)
    _, cs, _, tfmr_cs, _ = sample_and_classify_eval_batch(k1, params)
    ks = onp.full([eval_batch_size], k)
    tfmr_metrics = metrics.compute_masked_metrics(
        cs, tfmr_cs, ks, ks*data_points_per_mode,
        metrics=["pairwise_accuracy", "pairwise_f1",
                 "pairwise_macro_f1", "pairwise_micro_f1"])
    for metric_name, metric_val in tfmr_metrics.items():
      writer.scalar("transformer/%s" % metric_name,
                    metric_val, step=step)
      print("Transformer %s: %0.3f" % (metric_name, metric_val))

    plot_params(k*data_points_per_mode, writer, step, params, k2)
    plot_sklearn_comparison(writer, step, params)
    if step == 0:
      summarize_baselines(writer, step, k3)

  return summarize

def make_logdir(config):
  basedir = config.logdir
  exp_dir = (
      "ring_nheads_%d_nencoders_%d_ndecoders_%d_num_modes_%d"
      % (config.num_heads, config.num_encoders, config.num_decoders, config.k))
  return os.path.join(basedir, exp_dir)

def main(unused_argv):
  if FLAGS.debug_nans:
    config.update("jax_debug_nans", True)

  if FLAGS.parallel and train.can_train_parallel():
    assert FLAGS.batch_size % jax.local_device_count(
    ) == 0, "Device count must evenly divide batch_size"
    FLAGS.batch_size = int(FLAGS.batch_size / jax.local_device_count())

  key = jax.random.PRNGKey(0)
  key, subkey = jax.random.split(key)
  model, init_params = make_model(
      key,
      num_encoders=FLAGS.num_encoders,
      num_decoders=FLAGS.num_decoders,
      num_heads=FLAGS.num_heads,
      value_dim=FLAGS.value_dim_per_head*FLAGS.num_heads,
      data_points_per_mode=FLAGS.data_points_per_mode,
      k=FLAGS.k)
  loss_fn = make_loss(
      model,
      k=FLAGS.k,
      data_points_per_mode=FLAGS.data_points_per_mode,
      batch_size=FLAGS.batch_size)
  summarize_fn = make_summarize(
      model,
      k=FLAGS.k,
      data_points_per_mode=FLAGS.data_points_per_mode,
      eval_batch_size=FLAGS.eval_batch_size)
  train.train_loop(
      subkey,
      init_params,
      loss_fn,
      parallel=FLAGS.parallel,
      lr=FLAGS.lr,
      num_steps=FLAGS.num_steps,
      summarize_fn=summarize_fn,
      summarize_every=FLAGS.summarize_every,
      checkpoint_every=FLAGS.checkpoint_every,
      clobber_checkpoint=FLAGS.clobber_checkpoint,
      logdir=make_logdir(FLAGS))

if __name__ == "__main__":
  app.run(main)