#   Copyright notice
#   --------------------------------------------------------------------
#   Copyright (C) 2020 Deltares
#       Freek Scheel
#
#       freek.scheel@deltares.nl
#
#       P.O. Box 177
#       2600 MH Delft
#       The Netherlands
#
#   This library is free software: you can redistribute it and/or modify
#   it under the terms of the GNU Lesser General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This library 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 this library.  If not, see <http://www.gnu.org/licenses/>.
#   --------------------------------------------------------------------

from collections import OrderedDict
from collections.abc import Iterable
from datetime import datetime, timedelta
from itertools import count, takewhile

import numpy as np
from scipy.optimize import fsolve, leastsq

import cht_tide.constituent as constituent
from cht_tide.astro import astro

d2r, r2d = np.pi / 180.0, 180.0 / np.pi


class Tide(object):
    dtype = np.dtype([("constituent", object), ("amplitude", float), ("phase", float)])

    def __init__(
        self, constituents=None, amplitudes=None, phases=None, model=None, radians=False
    ):
        """
        Initialise a tidal model. Provide constituents, amplitudes and phases OR a model.
        Arguments:
        constituents -- list of constituents used in the model.
        amplitudes -- list of amplitudes corresponding to constituents
        phases -- list of phases corresponding to constituents
        model -- an ndarray of type Tide.dtype representing the constituents, amplitudes and phases.
        radians -- boolean representing whether phases are in radians (default False)
        """
        if None not in [constituents, amplitudes, phases]:
            if len(constituents) == len(amplitudes) == len(phases):
                model = np.zeros(len(phases), dtype=Tide.dtype)
                model["constituent"] = np.array(constituents)
                model["amplitude"] = np.array(amplitudes)
                model["phase"] = np.array(phases)
            else:
                raise ValueError(
                    "Constituents, amplitudes and phases should all be arrays of equal length."
                )
        elif model is not None:
            if not model.dtype == Tide.dtype:
                raise ValueError("Model must be a numpy array with dtype == Tide.dtype")
        else:
            raise ValueError(
                "Must be initialised with constituents, amplitudes and phases; or a model."
            )
        if radians:
            model["phase"] = r2d * model["phase"]
        self.model = model[:]
        self.normalize()

    def prepare(self, *args, **kwargs):
        return Tide._prepare(self.model["constituent"], *args, **kwargs)

    @staticmethod
    def _prepare(constituents, t0, t=None, radians=True):
        """
        Return constituent speed and equilibrium argument at a given time, and constituent node factors at given times.
        Arguments:
        constituents -- list of constituents to prepare
        t0 -- time at which to evaluate speed and equilibrium argument for each constituent
        t -- list of times at which to evaluate node factors for each constituent (default: t0)
        radians -- whether to return the angular arguments in radians or degrees (default: True)
        """
        # The equilibrium argument is constant and taken at the beginning of the
        # time series (t0).  The speed of the equilibrium argument changes very
        # slowly, so again we take it to be constant over any length of data. The
        # node factors change more rapidly.
        if isinstance(t0, Iterable):
            t0 = t0[0]
        if t is None:
            t = [t0]
        if not isinstance(t, Iterable):
            t = [t]
        a0 = astro(t0)
        a = [astro(t_i) for t_i in t]

        # For convenience give u, V0 (but not speed!) in [0, 360)
        V0 = np.array([c.V(a0) for c in constituents])[:, np.newaxis]
        speed = np.array([c.speed(a0) for c in constituents])[:, np.newaxis]
        u = [
            np.mod(np.array([c.u(a_i) for c in constituents])[:, np.newaxis], 360.0)
            for a_i in a
        ]
        f = [
            np.mod(np.array([c.f(a_i) for c in constituents])[:, np.newaxis], 360.0)
            for a_i in a
        ]

        if radians:
            speed = d2r * speed
            V0 = d2r * V0
            u = [d2r * each for each in u]
        return speed, u, f, V0

    def at(self, t):
        """
        Return the modelled tidal height at given times.
        Arguments:
        t -- array of times at which to evaluate the tidal height
        """
        t0 = t[0]
        hours = self._hours(t0, t)
        partition = 240.0
        t = self._partition(hours, partition)
        times = self._times(t0, [(i + 0.5) * partition for i in range(len(t))])
        speed, u, f, V0 = self.prepare(t0, times, radians=True)
        H = self.model["amplitude"][:, np.newaxis]
        p = d2r * self.model["phase"][:, np.newaxis]

        return np.concatenate(
            [
                Tide._tidal_series(t_i, H, p, speed, u_i, f_i, V0)
                for t_i, u_i, f_i in zip(t, u, f)
            ]
        )

    def highs(self, *args):
        """
        Generator yielding only the high tides.
        Arguments:
        see Tide.extrema()
        """
        for t in filter(lambda e: e[2] == "H", self.extrema(*args)):
            yield t

    def lows(self, *args):
        """
        Generator yielding only the low tides.
        Arguments:
        see Tide.extrema()
        """
        for t in filter(lambda e: e[2] == "L", self.extrema(*args)):
            yield t

    def form_number(self):
        """
        Returns the model's form number, a helpful heuristic for classifying tides.
        """
        k1, o1, m2, s2 = (
            np.extract(self.model["constituent"] == c, self.model["amplitude"])
            for c in [
                constituent._K1,
                constituent._O1,
                constituent._M2,
                constituent._S2,
            ]
        )
        return (k1 + o1) / (m2 + s2)

    def classify(self):
        """
        Classify the tide according to its form number
        """
        form = self.form_number()
        if 0 <= form <= 0.25:
            return "semidiurnal"
        elif 0.25 < form <= 1.5:
            return "mixed (semidiurnal)"
        elif 1.5 < form <= 3.0:
            return "mixed (diurnal)"
        else:
            return "diurnal"

    def extrema(self, t0, t1=None, partition=2400.0):
        """
        A generator for high and low tides.
        Arguments:
        t0 -- time after which extrema are sought
        t1 -- optional time before which extrema are sought (if not given, the generator is infinite)
        partition -- number of hours for which we consider the node factors to be constant (default: 2400.0)
        """
        if t1:
            # yield from in python 3.4
            for e in takewhile(lambda t: t[0] < t1, self.extrema(t0)):
                yield e
        else:
            # We assume that extrema are separated by at least delta hours
            delta = np.amin(
                [
                    90.0 / c.speed(astro(t0))
                    for c in self.model["constituent"]
                    if not c.speed(astro(t0)) == 0
                ]
            )
            # We search for stationary points from offset hours before t0 to
            # ensure we find any which might occur very soon after t0.
            offset = 24.0
            partitions = (Tide._times(t0, i * partition) for i in count()), (
                Tide._times(t0, i * partition) for i in count(1)
            )

            # We'll overestimate to be on the safe side;
            # values outside (start,end) won't get yielded.
            interval_count = int(np.ceil((partition + offset) / delta)) + 1
            amplitude = self.model["amplitude"][:, np.newaxis]
            phase = d2r * self.model["phase"][:, np.newaxis]

            for start, end in zip(*partitions):
                speed, [u], [f], V0 = self.prepare(
                    start, Tide._times(start, 0.5 * partition)
                )

                # These derivatives don't include the time dependence of u or f,
                # but these change slowly.
                def d(t):
                    return np.sum(
                        -speed * amplitude * f * np.sin(speed * t + (V0 + u) - phase),
                        axis=0,
                    )

                def d2(t):
                    return np.sum(
                        -(speed**2.0)
                        * amplitude
                        * f
                        * np.cos(speed * t + (V0 + u) - phase),
                        axis=0,
                    )

                # We'll overestimate to be on the safe side;
                # values outside (start,end) won't get yielded.
                intervals = (delta * i - offset for i in range(interval_count)), (
                    delta * (i + 1) - offset for i in range(interval_count)
                )
                for a, b in zip(*intervals):
                    if d(a) * d(b) < 0:
                        extrema = fsolve(d, (a + b) / 2.0, fprime=d2)[0]
                        time = Tide._times(start, extrema)
                        [height] = self.at([time])
                        hilo = "H" if d2(extrema) < 0 else "L"
                        if start < time < end:
                            yield (time, height, hilo)

    @staticmethod
    def _hours(t0, t):
        """
        Return the hourly offset(s) of a (list of) time from a given time.
        Arguments:
        t0 -- time from which offsets are sought
        t -- times to find hourly offsets from t0.
        """
        if not isinstance(t, Iterable):
            return Tide._hours(t0, [t])[0]
        elif isinstance(t[0], datetime):
            return np.array([(ti - t0).total_seconds() / 3600.0 for ti in t])
        else:
            return t

    @staticmethod
    def _partition(hours, partition=3600.0):
        """
        Partition a sorted list of numbers (or in this case hours).
        Arguments:
        hours -- sorted ndarray of hours.
        partition -- maximum partition length (default: 3600.0)
        """
        partition = float(partition)
        relative = hours - hours[0]
        total_partitions = np.ceil(
            relative[-1] / partition + 10 * np.finfo(np.float64).eps
        ).astype("int")
        return [
            hours[np.floor(np.divide(relative, partition)) == i]
            for i in range(total_partitions)
        ]

    @staticmethod
    def _times(t0, hours):
        """
        Return a (list of) datetime(s) given an initial time and an (list of) hourly offset(s).
        Arguments:
        t0 -- initial time
        hours -- hourly offsets from t0
        """
        if not isinstance(hours, Iterable):
            return Tide._times(t0, [hours])[0]
        elif not isinstance(hours[0], datetime):
            return np.array([t0 + timedelta(hours=h) for h in hours])
        else:
            return np.array(hours)

    @staticmethod
    def _tidal_series(t, amplitude, phase, speed, u, f, V0):
        return np.sum(amplitude * f * np.cos(speed * t + (V0 + u) - phase), axis=0)

    def normalize(self):
        """
        Adapt self.model so that amplitudes are positive and phases are in [0,360) as per convention
        """
        for i in range(len(self.model)):
            if self.model["amplitude"][i] < 0:
                self.model["amplitude"][i] = -self.model["amplitude"][i]
                self.model["phase"][i] = self.model["phase"][i] + 180.0
            self.model["phase"][i] = np.mod(self.model["phase"][i], 360.0)

    @classmethod
    def decompose(
        cls,
        heights,
        t=None,
        t0=None,
        interval=None,
        constituents=constituent.noaa,
        initial=None,
        n_period=2,
        callback=None,
        full_output=False,
    ):
        """
        Return an instance of Tide which has been fitted to a series of tidal observations.
        Arguments:
        It is not necessary to provide t0 or interval if t is provided.
        heights -- ndarray of tidal observation heights
        t -- ndarray of tidal observation times
        t0 -- datetime representing the time at which heights[0] was recorded
        interval -- hourly interval between readings
        constituents -- list of constituents to use in the fit (default: constituent.noaa)
        initial -- optional Tide instance to use as first guess for least squares solver
        n_period -- only include constituents which complete at least this many periods (default: 2)
        callback -- optional function to be called at each iteration of the solver
        full_output -- whether to return the output of scipy's leastsq solver (default: False)
        """
        if t is not None:
            if isinstance(t[0], datetime):
                hours = Tide._hours(t[0], t)
                t0 = t[0]
            elif t0 is not None:
                hours = t
            else:
                raise ValueError(
                    "t can be an array of datetimes, or an array "
                    "of hours since t0 in which case t0 must be "
                    "specified."
                )
        elif None not in [t0, interval]:
            hours = np.arange(len(heights)) * interval
        else:
            raise ValueError(
                "Must provide t(datetimes), or t(hours) and "
                "t0(datetime), or interval(hours) and t0(datetime) "
                "so that each height can be identified with an "
                "instant in time."
            )

        # Remove duplicate constituents (those which travel at exactly the same
        # speed, irrespective of phase)
        constituents = list(OrderedDict.fromkeys(constituents))

        # No need for least squares to find the mean water level constituent z0,
        # work relative to mean
        constituents = [c for c in constituents if not c == constituent._Z0]
        z0 = np.mean(heights)
        heights = heights - z0

        # Only analyse frequencies which complete at least n_period cycles over
        # the data period.
        constituents = [
            c for c in constituents if 360.0 * n_period < hours[-1] * c.speed(astro(t0))
        ]
        n = len(constituents)

        sort = np.argsort(hours)
        hours = hours[sort]
        heights = heights[sort]

        # We partition our time/height data into intervals over which we consider
        # the values of u and f to assume a constant value (that is, their true
        # value at the midpoint of the interval).  Constituent
        # speeds change much more slowly than the node factors, so we will
        # consider these constant and equal to their speed at t0, regardless of
        # the length of the time series.

        partition = 240.0

        t = Tide._partition(hours, partition)
        times = Tide._times(t0, [(i + 0.5) * partition for i in range(len(t))])

        speed, u, f, V0 = Tide._prepare(constituents, t0, times, radians=True)

        # Residual to be minimised by variation of parameters (amplitudes, phases)
        def residual(hp):
            H, p = hp[:n, np.newaxis], hp[n:, np.newaxis]
            s = np.concatenate(
                [
                    Tide._tidal_series(t_i, H, p, speed, u_i, f_i, V0)
                    for t_i, u_i, f_i in zip(t, u, f)
                ]
            )
            res = heights - s
            if callback:
                callback(res)
            return res

        # Analytic Jacobian of the residual - this makes solving significantly
        # faster than just using gradient approximation, especially with many
        # measurements / constituents.
        def D_residual(hp):
            H, p = hp[:n, np.newaxis], hp[n:, np.newaxis]
            ds_dH = np.concatenate(
                [
                    f_i * np.cos(speed * t_i + u_i + V0 - p)
                    for t_i, u_i, f_i in zip(t, u, f)
                ],
                axis=1,
            )

            ds_dp = np.concatenate(
                [
                    H * f_i * np.sin(speed * t_i + u_i + V0 - p)
                    for t_i, u_i, f_i in zip(t, u, f)
                ],
                axis=1,
            )

            return np.append(-ds_dH, -ds_dp, axis=0)

        # Initial guess for solver, haven't done any analysis on this since the
        # solver seems to converge well regardless of the initial guess We do
        # however scale the initial amplitude guess with some measure of the
        # variation
        amplitudes = np.ones(n) * (np.sqrt(np.dot(heights, heights)) / len(heights))
        phases = np.ones(n)

        if initial:
            for c0, amplitude, phase in initial.model:
                for i, c in enumerate(constituents):
                    if c0 == c:
                        amplitudes[i] = amplitude
                        phases[i] = d2r * phase

        initial = np.append(amplitudes, phases)

        lsq = leastsq(residual, initial, Dfun=D_residual, col_deriv=True, ftol=1e-7)

        model = np.zeros(1 + n, dtype=cls.dtype)
        model[0] = (constituent._Z0, z0, 0)
        model[1:]["constituent"] = constituents[:]
        model[1:]["amplitude"] = lsq[0][:n]
        model[1:]["phase"] = lsq[0][n:]

        if full_output:
            return cls(model=model, radians=True), lsq
        return cls(model=model, radians=True)
