"""Classes for probing values of Variables, for example, along a line."""
from __future__ import absolute_import
import numpy as nm
import numpy.linalg as nla

from sfepy.base.base import get_default, basestr, Struct
from sfepy.linalg import make_axis_rotation_matrix, norm_l2_along_axis
import six

def write_results(filename, probe, results):
    """
    Write probing results into a file.

    Parameters
    ----------
    filename : str or file object
        The output file name.
    probe : Probe subclass instance
        The probe used to obtain the results.
    results : dict
        The dictionary of probing results. Keys are data names, values are
        the probed values.
    """
    fd = open(filename, 'w') if isinstance(filename, basestr) else filename

    fd.write('\n'.join(probe.report()) + '\n')
    for key, result in six.iteritems(results):
        pars, vals = result
        fd.write('\n# %s %d\n' % (key, vals.shape[-1]))

        if vals.ndim == 1:
            aux = nm.hstack((pars[:,None], vals[:,None]))

        else:
            aux = nm.hstack((pars[:,None], vals))

        nm.savetxt(fd, aux)

    if isinstance(filename, basestr):
        fd.close()

def read_results(filename, only_names=None):
    """
    Read probing results from a file.

    Parameters
    ----------
    filename : str or file object
        The probe results file name.

    Returns
    -------
    header : Struct instance
        The probe data header.
    results : dict
        The dictionary of probing results. Keys are data names, values are
        the probed values.
    """
    from sfepy.base.ioutils import read_array

    only_names = get_default(only_names, [])

    fd = open(filename, 'r') if isinstance(filename, basestr) else filename

    header = read_header(fd)
    results = {}
    for name, nc in get_data_name(fd):
        if name not in only_names: continue

        result = read_array(fd, header.n_point, nc + 1, nm.float64)
        results[name] = result

    return header, results

def read_header(fd):
    """
    Read the probe data header from file descriptor fd.

    Returns
    -------
    header : Struct instance
        The probe data header.
    """
    header = Struct(name='probe_data_header')
    header.probe_class = fd.readline().strip()

    aux = fd.readline().strip().split(':')[1]
    header.n_point = int(aux.strip().split()[0])

    details = []
    while 1:
        line = fd.readline().strip()

        if line == '-----':
            break
        else:
            details.append(line)
    header.details = '\n'.join(details)

    return header

def get_data_name(fd):
    """
    Try to read next data name in file fd.

    Returns
    -------
    name : str
        The data name.
    nc : int
        The number of data columns.
    """
    name = None
    while 1:
        try:
            line = fd.readline()
            if (len(line) == 0): break
            if len(line) == 1: continue
        except:
            raise StopIteration

        line = line.strip().split()
        if (len(line) == 3) and (line[0] == '#'):
            name = line[1]
            nc = int(line[2])

            yield name, nc

class Probe(Struct):
    """
    Base class for all point probes. Enforces two points minimum.
    """
    cache = Struct(name='probe_shared_evaluate_cache')
    is_cyclic = False

    def __init__(self, name, share_geometry=True, n_point=None, **kwargs):
        """
        Parameters
        ----------
        name : str
            The probe name, set automatically by the subclasses.
        share_geometry : bool
            Set to True to indicate that all the probes will work on the same
            domain. Certain data are then computed only for the first probe and
            cached.
        n_point : int
           The (fixed) number of probe points, when positive. When non-positive,
           the number of points is adaptively increased starting from -n_point,
           until the neighboring point distance is less than the diameter of the
           elements enclosing the points. When None, it is set to -10.

        For additional parameters see the __init__() docstrings of the
        subclasses.
        """
        Struct.__init__(self, name=name, share_geometry=share_geometry,
                        **kwargs)

        self.set_n_point(n_point)

        self.options = Struct(close_limit=0.1, size_hint=None)
        self.cache = Struct(name='probe_local_evaluate_cache')

        self.is_refined = False

    def get_evaluate_cache(self):
        """
        Return the evaluate cache for domain-related data given by
        `self.share_geometry`.
        """
        return Probe.cache if self.share_geometry else self.cache

    def set_n_point(self, n_point):
        """
        Set the number of probe points.

        Parameters
        ----------
        n_point : int
           The (fixed) number of probe points, when positive. When non-positive,
           the number of points is adaptively increased starting from -n_point,
           until the neighboring point distance is less than the diameter of the
           elements enclosing the points. When None, it is set to -10.
        """
        if n_point is None:
            n_point = -10

        if n_point <= 0:
            n_point = max(-n_point, 2)
            self.n_point_required = -1

        else:
            n_point = max(n_point, 2)
            self.n_point_required = n_point

        self.n_point0 = self.n_point = n_point

    def set_options(self, close_limit=None, size_hint=None):
        """
        Set the probe options.

        Parameters
        ----------
        close_limit : float
            The maximum limit distance of a point from the closest
            element allowed for extrapolation.
        size_hint : float
            Element size hint for the refinement of probe parametrization.
        """
        if close_limit is not None:
            self.options.close_limit = close_limit

        if size_hint is not None:
            self.options.size_hint = size_hint

    def report(self):
        """Report the probe parameters."""
        out = [self.__class__.__name__]

        if self.n_point_required == -1:
            aux = 'adaptive'
        else:
            aux = 'fixed'
        out.append('number of points: %s (%s)' % (self.n_point, aux))

        return out

    def __call__(self, variable, **kwargs):
        """
        Probe the given variable. The actual implementation is in self.probe(),
        so that it can be overridden in subclasses.

        Parameters
        ----------
        variable : Variable instance
            The variable to be sampled along the probe.
        **kwargs : additional arguments
            See :func:`Probe.probe()`.
        """
        return self.probe(variable, **kwargs)

    def probe(self, variable, mode='val', ret_points=False):
        """
        Probe the given variable.

        Parameters
        ----------
        variable : Variable instance
            The variable to be sampled along the probe.
        mode : {'val', 'grad'}, optional
            The evaluation mode: the variable value (default) or the
            variable value gradient.
        ret_points : bool
            If True, return also the probe points.

        Returns
        -------
        pars : array
            The parametrization of the probe points.
        points : array, optional
            If `ret_points` is True, the coordinates of points corresponding to
            `pars`, where the `variable` is evaluated.
        vals : array
            The probed values.
        """
        refine_flag = None

        ev = variable.evaluate_at
        field = variable.field

        cache = field.get_evaluate_cache(cache=self.get_evaluate_cache(),
                                         share_geometry=self.share_geometry)
        self.reset_refinement()

        while True:
            pars, points = self.get_points(refine_flag)
            if not nm.isfinite(points).all():
                raise ValueError('Inf/nan in probe points!')

            vals, cells = ev(points, mode=mode, strategy='general',
                             close_limit=self.options.close_limit,
                             cache=cache, ret_cells=True)

            if self.is_refined:
                break

            else:
                refine_flag = self.refine_points(variable, points, cells)
                if (refine_flag == False).all():
                    break

        self.is_refined = True

        if ret_points:
            return pars, points, vals

        else:
            return pars, vals

    def reset_refinement(self):
        """
        Reset the probe refinement state.
        """
        self.is_refined = False
        self.n_point = self.n_point0

    def refine_points(self, variable, points, cells):
        """
        Mark intervals between points for a refinement, based on element
        sizes at those points. Assumes the points to be ordered.

        Returns
        -------
        refine_flag : bool array
            True at places corresponding to intervals between subsequent points
            that need to be refined.
        """
        if self.n_point_required == self.n_point:
            refine_flag = nm.array([False])

        else:
            if self.options.size_hint is None:
                ed = variable.get_element_diameters(cells, 0)
                pd = 0.5 * (ed[1:] + ed[:-1])

            else:
                pd = self.options.size_hint

            dist = 
norm_l2_along_axis(points[1:] - points[:-1])