"""Module for records related function.

A records is a dict of multiple sfg2d.SfgRecord objects that are treated
together.

"""
import os
import yaml
from . import plot
from . import fig as sfgfig
from .core import SfgRecord
import matplotlib.pyplot as plt
import logging
from numpy import linspace, array, concatenate
from scipy.stats import sem
import pandas as pd
import numpy as np

from sfg2d.utils.config import CONFIG
FRAME_AXIS_INDEX = CONFIG['FRAME_AXIS_INDEX']

## Dict of records.
#records = {}
## Dicto of model.
#models = {}

logger = logging.getLogger(__name__)

try:
    if not isinstance(records, dict):
        records = {}
except NameError:
    records = {}

try:
    if not isinstance(models, dict):
        models = {}
except NameError:
    models = {}


def records_agg(names, func):
    """Aggregation function for a selection of records."""
    ret = []
    for name in names:
        record = records[name]
        ret.append(func(record))
    return ret


def list_agg(object_list, func):
    """Aggregation function for a list of objects."""
    ret = []
    for elm in object_list:
        ret.append(func(elm))
    return ret


def records_list_agg(records, func):
    """Aggregation function for a list of records."""
    return list_agg(records, func)


def get_xyerr(
    record,
    kwargs_xdata=None,
    kwargs_ydata=None,
    kwargs_yerr=None,
    squeeze=False,
):
    """Get a triple of data from record.

    A triple of data can in general be xdata, ydata and yerr. Can be used for
    e.g. Pump-Vis, or to extract static spectra from time scan data.

    """
    if not kwargs_xdata:
        kwargs_xdata = {}
    if not kwargs_ydata:
        kwargs_ydata = {}
    if not kwargs_yerr:
        kwargs_yerr = {}

    kwargs_xdata.setdefault('prop', 'wavenumber')

    kwargs_ydata.setdefault('prop', 'normalized')
    kwargs_ydata.setdefault('frame_med', True)

    for key, value in kwargs_ydata.items():
        kwargs_yerr.setdefault(key, value)
    kwargs_yerr['frame_med'] = False

    xdata = record.select(**kwargs_xdata)
    ydata = record.select(**kwargs_ydata)
    yerr = record.sem(**kwargs_ydata)

    if squeeze:
        return xdata.squeeze(), ydata.squeeze(), yerr.squeeze()

    return xdata, ydata, yerr


def save_models(models, fpath):
    """Save models to fpath.

    models must be a dict with {'name': {'model': sfg2d.models.Model}}
    """

    s_dict = {}
    for name, value in models.items():
        model = value['model']
        s_dict[name] = {
            'fitarg': model.fitarg,
            'class': model.__class__,
            'xdata': model.xdata.tolist(),
            'ydata': model.ydata.tolist(),
            'sigma': model.sigma.tolist(),
        }

    with open(fpath, 'w') as ofile:
        yaml.dump(
            s_dict,
            ofile,
            default_flow_style=False
        )


def load_models(fpath):
    """Load models saved with save_models and transform into model objects."""
    with open(fpath, 'r') as stream:
        imp_models = yaml.load(stream)

    # Transform import into corresponding model objects
    for name, value in imp_models.items():
        value['xdata'] = np.array(value['xdata'])
        value['ydata'] = np.array(value['ydata'])
        value['sigma'] = np.array(value['sigma'])

        value['model'] = value['class'](
            xdata=value['xdata'],
            ydata=value['ydata'],
            sigma=value['sigma'],
            fitarg=value['fitarg'],
        )
    return imp_models


def plot_spectra(
        record_names, kwargs_xdata, kwargs_ydata, kwargs_plots=None,
        kwargs_xdata_record=None, kwargs_ydata_record=None, kwargs_plot=None):
    """High level function to generate plots.

    **Arguments**
      - **record_names**: List of keys for the `records` dict.
      - **kwargs_xdata**: dict with x kwargs for all records
      - **kwargs_ydata**: dict with y kwargs for all records
      - **kwargs_plots**: dict with kwargs for each record_name
      - **kwargs_xdata_record**: Dict with xdata kwargs per record
        specific kwargs overwrite general
      - **kwargs_ydata_record**: Dict with ydata kwargs per record
        specific kwargs overwrite general
      - **kwargs_plot**: Plot kwargs applied to all lines

    **Returns**:
    The xdatas and ydatas used during making the plot
    """
    global records
    if not kwargs_plots:
        kwargs_plots = {}
    if not kwargs_xdata_record:
        kwargs_xdata_record = {}
    if not kwargs_ydata_record:
        kwargs_ydata_record = {}
    if not kwargs_plot:
        kwargs_plot = {}

    xdatas = []
    ydatas = []
    for name in record_names:
        record = records.get(name)
        if not record:
            print('{} not found in records'.format(name))
            continue

        _kwargs_xdata = kwargs_ydata_record.get(name, {})
        _kwargs_ydata = kwargs_ydata_record.get(name, {})

        tkwx = {**kwargs_xdata, **_kwargs_xdata}
        tkwy = {**kwargs_ydata, **_kwargs_ydata}

        # Must combine global and specific kwarg dicts
        logger.debug('kwargs_plot is: {}'.format(kwargs_plot))
        _kwargs_plot = {**kwargs_plot, **kwargs_plots.get(name, {})}
        xdata = record.select(**tkwx)
        xdatas.append(xdata)