# -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst

import numpy as np

from datetime import datetime

from astropy.time import Time
from astropy import units as u

__all__ = ['time_support']

__doctest_requires__ = {'time_support': ['matplotlib']}

UNSUPPORTED_FORMATS = ('datetime', 'datetime64')
YMDHMS_FORMATS = ('fits', 'iso', 'isot', 'yday')
STR_FORMATS = YMDHMS_FORMATS + ('byear_str', 'jyear_str')


def time_support(*, scale=None, format=None, simplify=True):
    """
    Enable support for plotting `astropy.time.Time` instances in
    matplotlib.

    May be (optionally) used with a ``with`` statement.

      >>> import matplotlib.pyplot as plt
      >>> from astropy import units as u
      >>> from astropy import visualization
      >>> with visualization.time_support():  # doctest: +IGNORE_OUTPUT
      ...     plt.figure()
      ...     plt.plot(Time(['2016-03-22T12:30:31', '2016-03-22T12:30:38', '2016-03-22T12:34:40']))
      ...     plt.draw()

    Parameters
    ----------
    scale : str, optional
        The time scale to use for the times on the axis. If not specified,
        the scale of the first Time object passed to Matplotlib is used.
    format : str, optional
        The time format to use for the times on the axis. If not specified,
        the format of the first Time object passed to Matplotlib is used.
    simplify : bool, optional
        If possible, simplify labels, e.g. by removing 00:00:00.000 times from
        ISO strings if all labels fall on that time.
    """

    import matplotlib.units as units
    from matplotlib.ticker import MaxNLocator, ScalarFormatter
    from astropy.visualization.wcsaxes.utils import select_step_hour, select_step_scalar

    class AstropyTimeLocator(MaxNLocator):

        # Note: we default to AutoLocator since many time formats
        # can just use this.

        def __init__(self, converter, *args, **kwargs):
            kwargs['nbins'] = 4
            super().__init__(*args, **kwargs)
            self._converter = converter

        def tick_values(self, vmin, vmax):

            # Where we put the ticks depends on the format we are using
            if self._converter.format in YMDHMS_FORMATS:

                # If we are here, we need to check what the range of values
                # is and decide how to find tick locations accordingly

                vrange = vmax - vmin

                if (self._converter.format != 'yday' and vrange > 31) or vrange > 366:  # greater than a month

                    # We need to be careful here since not all years and months have
                    # the same length

                    # Start off by converting the values from the range to
                    # datetime objects, so that we can easily extract the year and
                    # month.

                    tmin = Time(vmin, scale=self._converter.scale, format='mjd').datetime
                    tmax = Time(vmax, scale=self._converter.scale, format='mjd').datetime

                    # Find the range of years
                    ymin = tmin.year
                    ymax = tmax.year

                    if ymax > ymin + 1:  # greater than a year

                        # Find the step we want to use
                        ystep = int(select_step_scalar(max(1, (ymax - ymin) / 3)))

                        ymin = ystep * (ymin // ystep)

                        # Generate the years for these steps
                        times = []
                        for year in range(ymin, ymax + 1, ystep):
                            times.append(datetime(year=year, month=1, day=1))

                    else:  # greater than a month but less than a year

                        mmin = tmin.month
                        mmax = tmax.month + 12 * (ymax - ymin)

                        mstep = int(select_step_scalar(max(1, (mmax - mmin) / 3)))

                        mmin = mstep * max(1, mmin // mstep)

                        # Generate the months for these steps
                        times = []
                        for month in range(mmin, mmax + 1, mstep):
                            times.append(datetime(year=ymin + month // 12,
                                                  month=month % 12, day=1))

                    # Convert back to MJD
                    values = Time(times, scale=self._converter.scale).mjd

                elif vrange > 1:  # greater than a day

                    self.set_params(steps=[1, 2, 5, 10])
                    values = super().tick_values(vmin, vmax)

                else:

                    # Determine ideal step
                    dv = (vmax - vmin) / 3 * 24 << u.hourangle

                    # And round to nearest sensible value
                    dv = select_step_hour(dv).to_value(u.hourangle) / 24

                    # Determine tick locations
                    imin = np.ceil(vmin / dv)
                    imax = np.floor(vmax / dv)
                    values = np.arange(imin, imax + 1, dtype=np.int64) * dv

            else:

                values = super().tick_values(vmin, vmax)

            # Get rid of values outside of the input interval
            values = values[(values >= vmin) & (values <= vmax)]

            return values

        def __call__(self):
            vmin, vmax = self.axis.get_view_interval()
            return self.tick_values(vmin, vmax)

    class AstropyTimeFormatter(ScalarFormatter):

        def __init__(self, converter, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self._converter = converter
            self.set_useOffset(False)
            self.set_scientific(False)

        def __call__(self, value, pos=None):
            # Needed for Matplotlib <3.1
            if self._converter.format in STR_FORMATS:
                return self.format_ticks([value])[0]
            else:
                return super().__call__(value, pos=pos)

        def format_ticks(self, values):
            if len(values) == 0:
                return []
            if self._converter.format in YMDHMS_FORMATS:
                times = Time(values, format='mjd', scale=self._converter.scale)
                formatted = getattr(times, self._converter.format)
                if self._converter.simplify:
                    if self._converter.format in ('fits', 'iso', 'isot'):
                        if all([x.endswith('00:00:00.000') for x in formatted]):
                            split = ' ' if self._converter.format == 'iso' else 'T'
                            formatted = [x.split(split)[0] for x in formatted]
                    elif self._converter.format == 'yday':
                        if all([x.endswith(':001:00:00:00.000') for x in formatted]):
                            formatted = [x.split(':', 1)[0] for x in formatted]
                return formatted
            elif self._converter.format == 'byear_str':
                return Time(values, format='byear', scale=self._converter.scale).byear_str
            elif self._converter.format == 'jyear_str':
                return Time(values, format='jyear', scale=self._converter.scale).jyear_str
            else:
                return super().format_ticks(values)

    class MplTimeConverter(units.ConversionInterface):

        def __init__(self, scale=None, format=None, simplify=None):

            super().__init__()

            self.format = format
            self.scale = scale
            self.simplify = simplify

            # Keep track of original converter in case the context manager is
            # used in a nested way.
            self._original_converter = units.registry.get(Time)

            units.registry[Time] = self

        @property
        def format(self):
            return self._format

        @format.setter
        def format(self, value):
            if value in UNSUPPORTED_FORMATS:
                raise ValueError(f'time_support does not support format={value}')
            self._format = value

        def __enter__(self):
            return self

        def __exit__(self, type, value, tb):
            if self._original_converter is None:
                del units.registry[Time]
            else:
                units.registry[Time] = self._original_converter

        def default_units(self, x, axis):
            if isinstance(x, tuple):
                x = x[0]
            if self.format is None:
                self.format = x.format
            if self.scale is None:
                self.scale = x.scale
            return 'astropy_time'

        def convert(self, value, unit, axis):
            """
            Convert a Time value to a scalar or array.
            """

            # For Matplotlib < 2.2
            if not isinstance(value, Time):
                return value

            scaled = getattr(value, self.scale)
            if self.format in YMDHMS_FORMATS:
                return scaled.mjd
            elif self.format == 'byear_str':
                return scaled.byear
            elif self.format == 'jyear_str':
                return scaled.jyear
            else:
                return getattr(scaled, self.format)

        def axisinfo(self, unit, axis):
            """
            Return major and minor tick locators and formatters.
            """
            majloc = AstropyTimeLocator(self)
            majfmt = AstropyTimeFormatter(self)
            return units.AxisInfo(majfmt=majfmt,
                                  majloc=majloc,
                                  label=f'Time ({self.scale})')

    return MplTimeConverter(scale=scale, format=format, simplify=simplify)
