"""
Classes for easy interpolation of trajectories and Curves.
Requires Scipy installed.
"""

import numpy as np


class Interpolator:
    """ Poorman's linear interpolator, doesn't require Scipy. """

    def __init__(self, tt=None, ss=None, ttss=None, left=None, right=None):

        if ttss is not None:
            tt, ss = zip(*ttss)

        self.tt = 1.0 * np.array(tt)
        self.ss = 1.0 * np.array(ss)
        self.left = left
        self.right = right
        self.tmin, self.tmax = min(tt), max(tt)

    def __call__(self, t):
        return np.interp(t, self.tt, self.ss, self.left, self.right)


class Trajectory:
    def __init__(self, tt, xx, yy):

        self.tt = 1.0 * np.array(tt)
        self.xx = np.array(xx)
        self.yy = np.array(yy)
        self.update_interpolators()

    def __call__(self, t):
        return np.array([self.xi(t), self.yi(t)])

    def addx(self, x):
        return Trajectory(self.tt, self.xx + x, self.yy)

    def addy(self, y):
        return Trajectory(self.tt, self.xx, self.yy + y)

    def update_interpolators(self):
        self.xi = Interpolator(self.tt, self.xx)
        self.yi = Interpolator(self.tt, self.yy)

    def txy(self, tms=False):
        return zip((1000 if tms else 1) * self.tt, self.xx, self.yy)

    def to_file(self, filename):
        np.savetxt(filename, np.array(self.txy(tms=True)), fmt="%d", delimiter="\t")

    @staticmethod
    def from_file(filename):
        arr = np.loadtxt(filename, delimiter="\t")
        tt, xx, yy = arr.T
        return Trajectory(1.0 * tt / 1000, xx, yy)

    @staticmethod
    def save_list(trajs, filename):
        N = len(trajs)
        arr = np.hstack([np.array(list(t.txy(tms=True))) for t in trajs])
        np.savetxt(
            filename,
            arr,
            fmt="%d",
            delimiter="\t",
            header="\t".join(N * ["t(ms)", "x", "y"]),
        )

    @staticmethod
    def load_list(filename):
        arr = np.loadtxt(filename, delimiter="\t").T
        Nlines = arr.shape[0]
        return [
            Trajectory(tt=1.0 * a[0] / 1000, xx=a[1], yy=a[2])
            for a in np.split(arr, Nlines / 3)
        ]
