#  Copyright 2017 Carlos Pascual-Izarra <cpascual@users.sourceforge.net>
#
#  This file is part of PeakEvo.
#
#  PeakEvo is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  PeakEvo is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Foobar.  If not, see <http://www.gnu.org/licenses/>.

import logging
from collections import OrderedDict

import numpy


class LimitedSizeDict(OrderedDict):
    """
    limited-size dictionary for caches.
    From https://stackoverflow.com/a/2437645
    """

    def __init__(self, *args, **kwargs):
        self.size_limit = kwargs.pop("size_limit", None)
        OrderedDict.__init__(self, *args, **kwargs)
        self._check_size_limit()

    def __setitem__(self, *args, **kwargs):
        OrderedDict.__setitem__(self, *args, **kwargs)
        self._check_size_limit()

    def _check_size_limit(self):
        if self.size_limit is not None:
            while len(self) > self.size_limit:
                self.popitem(last=False)


class CacheDict(LimitedSizeDict):
    """A limited-size dict + default dict that generates and returns
    key->CacheValue(key) entry if key is not already cached
    """

    def __getitem__(self, k):
        if not self.__contains__(k):
            ret = CacheValue(k)
            self.__setitem__(k, ret)
            return ret
        return LimitedSizeDict.__getitem__(self, k)


class CacheValue(object):
    def __init__(
        self, name, spectrum=None, roi=None, peak=None, acq_time=None, datetime=None
    ):
        self.name = name
        self._spectrum = spectrum
        self._roi = roi
        self._peak = peak
        self._filt = None
        self.acq_time = acq_time
        self.timestamp = datetime

    def load(self):
        if self.name.endswith(".txt"):
            return self.loadtxt()
        _file = numpy.load(self.name)
        self._spectrum = _file["x"], _file["y"]
        self.acq_time = int(_file["acq_time"])
        self.timestamp = float(_file["timestamp"])

    def savez(self):
        x, y = self.spectrum
        numpy.savez(
            self.name, x=x, y=y, acq_time=self.acq_time, timestamp=self.timestamp
        )

    def loadtxt(self):
        import datetime
        import time

        with open(self.name) as f:
            isotime = f.readline().strip("#").strip()
            self._spectrum = numpy.loadtxt(f, unpack=True)
            dt = datetime.datetime.strptime(isotime, "%Y-%m-%dT%H:%M:%S.%f")
            self.timestamp = time.mktime(dt.timetuple()) + dt.microsecond * 1e-6

    @property
    def spectrum(self):
        if self._spectrum is None:
            self.load()
        return self._spectrum

    @spectrum.setter
    def spectrum(self, value):
        self._spectrum = value
        # invalidate spectrum-dependent cached values
        self._roi = self._peak = self._filt = None

    @property
    def roi(self):
        if self._roi is None:
            x = self.spectrum[0]
            self.roi = x[0], x[-1]
        return self._roi

    @roi.setter
    def roi(self, value):
        if value == self._roi:
            return
        self._roi = value
        # invalidate roi-dependent cached values
        self._peak = self._filt = None

    @property
    def filt(self):
        if self._filt is None:
            x, y = self.spectrum
            # crop data to roi and smooth it
            roi_min, roi_max = self.roi
            mask = (x > roi_min) * (x < roi_max)
            self._filt = x[mask], self.smooth(y[mask])
        return self._filt

    @property
    def peak(self):
        if self._peak is None:
            self._peak = self.findPeak(*self.filt)
        return self._peak

    @staticmethod
    def smooth_savgol(masked_y):
        from scipy.signal import savgol_filter

        # use half the roi (+1 if it is even)
        window = (masked_y.size // 4) * 2 + 1
        try:
            return savgol_filter(masked_y, window, 2)
        except Exception as e:
            logging.warning("Problem smoothing %r", e)
            return masked_y

    @staticmethod
    def smooth_poly2(masked_y):
        x = numpy.arange(len(masked_y))
        try:
            p = numpy.polyfit(x, masked_y, 2)
            return numpy.polyval(p, x)
        except Exception as e:
            logging.warning("Problem smoothing %r", e)
            return masked_y

    @staticmethod
    def findPeak_dummy(x, y):
        # TODO. This is a dummy implementation.
        argmax = y.argmax()
        return x[argmax], y[argmax]

    @staticmethod
    def findPeak_poly2(x, y):
        try:
            p = numpy.polyfit(x, y, 2)
        except Exception as e:
            logging.warning("Problem finding peak: %r", e)
            return CacheValue.findPeak_dummy(x, y)
        xmax = -p[1] * 0.5 / p[0]
        ymax = numpy.polyval(p, xmax)
        return xmax, ymax

    @staticmethod
    def findPeak_parab(x, y):
        try:
            # Use a parabolic approx around the max for interpolation
            argmax = y.argmax()
            if argmax == 0 or argmax == len(y) - 1:
                return x[argmax], y[argmax]

            x1, x2, x3 = x[argmax - 1 : argmax + 2] * 1.0
            y1, y2, y3 = y[argmax - 1 : argmax + 2] * 1.0

            denom = (x1 - x2) * (x1 - x3) * (x2 - x3)
            a = (x3 * (y2 - y1) + x2 * (y1 - y3) + x1 * (y3 - y2)) / denom
            b = (
                x3 * x3 * (y1 - y2) + x2 * x2 * (y3 - y1) + x1 * x1 * (y2 - y3)
            ) / denom
            c = (
                x2 * x3 * (x2 - x3) * y1
                + x3 * x1 * (x3 - x1) * y2
                + x1 * x2 * (x1 - x2) * y3
            ) / denom
            xv = -b / (2 * a)
            yv = c - b**2 / (4 * a)
        except Exception as e:
            logging.warning("Cannot find precise position of peak. Reason: %s", e)
            xv, yv = x[argmax], y[argmax]
        return xv, yv  # coordinates of the vertex

    # Select the desired findPeak implementation
    # findPeak = findPeak_parab
    # findPeak = findPeak_dummy
    findPeak = findPeak_poly2

    # Select the desired smooth implementation
    # smooth = smooth_savgol
    # smooth = smooth_dummy
    smooth = smooth_poly2
