#!/usr/bin/env python3
# -*- coding: utf-8 -*-
__copyright__ = """ This code is licensed under the 3-clause BSD license.
Copyright ETH Zurich, Department of Chemistry and Applied Biosciences, Reiher Group.
See LICENSE.txt for details.
"""
"""
Utility functions for working with VTK arrays.
"""

import typing
import sys
import math
from itertools import islice, repeat, chain
import vtk


def iterable_to_vtk_array(
    data: typing.Iterable[float], number_of_values: int, default_value: float = 0.0
) -> vtk.vtkDoubleArray:
    """
    Stores the numbers generated by the iterable to a new
    vtkDoubleArray and returns this array.

    The parameter `number_of_values` is used to pre-allocate
    the correct size.
    """

    result = vtk.vtkDoubleArray()
    result.SetNumberOfValues(number_of_values)
    for index, value in enumerate(
        islice(chain(data, repeat(default_value)), number_of_values)
    ):
        result.SetValue(index, value)
    return result


def rescale_to_range(
    array: vtk.vtkDoubleArray,
    to_range: typing.Tuple[float, float],
    rel_tol: float = sys.float_info.epsilon,
    abs_tol: float = sys.float_info.epsilon,
) -> vtk.vtkDoubleArray:
    """
    Rescales the given array to the desired range.

    The tolerances are used to determine whether the input range has span zero.
    In this case a small nonzero span is used.
    """
    to_span = to_range[1] - to_range[0]
    assert to_span >= 0

    # The values need to span a positive range to be able to scale to `to_range`.
    # We use at least a small span derived from the tolerances.
    array_range = array.GetValueRange()
    array_span = array_range[1] - array_range[0]
    array_center = array_range[0] + array_span / 2
    from_range = (
        array_range
        if not math.isclose(array_span, 0.0, rel_tol=rel_tol, abs_tol=abs_tol)
        else (
            array_center - max(rel_tol * abs(array_center), abs_tol),
            array_center + max(rel_tol * abs(array_center), abs_tol),
        )
    )
    from_span = from_range[1] - from_range[0]

    assert not math.isclose(from_span, 0.0, rel_tol=rel_tol, abs_tol=abs_tol)
    factor = to_span / from_span

    result = vtk.vtkDoubleArray()
    result.SetNumberOfValues(array.GetNumberOfValues())
    for id in range(array.GetNumberOfValues()):
        result.InsertValue(
            id, to_range[0] + (array.GetValue(id) - from_range[0]) * factor
        )

    return result
