# Copyright (C) 2020 SuperDARN Canada, University of Saskatchewan
# Author: Marina Schmidt
# This code is improvement based on acf.py in the DaVitpy library
# https://github.com/vtsuperdarn/davitpy/blob/master/davitpy

import copy
import matplotlib.pyplot as plt
import numpy as np
import warnings
import matplotlib.ticker as ticker

from datetime import datetime
from typing import List

from pydarn import (dmap2dict, DmapArray, DmapScalar,
                    plot_exceptions, SuperDARNRadars,
                    standard_warning_format, time2datetime,
                    check_data_type)

warnings.formatwarning = standard_warning_format


class ACF():
    """
    ACF class plots SuperDARN RAWACF data

    Class pattern design: Builder
    This class inherits matplotlib.pyplot to inherit plotting features as well
    build off their builder design pattern.

    Notes
    -----
    This class extends from matplotlib to generate the figures

    Methods
    -------
    plot
    plot_amplitude (not implemented yet)
    """

    def __str__(self):
        return "This class is static class that provides"\
                " the following methods: \n"\
                "   - plot_acfs()\n"\


    @classmethod
    def plot_acfs(cls, dmap_data: List[dict], beam_num: int = 0,
                  gate_num: int = 15, parameter: str = 'acfd',
                  scan_num: int = 0, start_time: datetime = None, ax=None,
                  normalized: bool = True, real_color: str = 'red',
                  plot_blank: bool = True, blank_marker: str = 'o',
                  imaginary_color: str = 'blue', legend: bool = True,
                  **kwargs):
        """
        plots the parameter ACF/XCF field from SuperDARN file,
        typically RAWACF format for a given beam and gate number

        Parameters
        ----------
        dmap_data : list[dict]
            records from a dmap file
        beam_num : int
            the beam number to plot, default: 0
        gate_num : int
            the gate number to plot, default: 15
        parameter : str
            the parameter to plot, default: acfd
        scan_number: int
            will plot the scan number for the provided beam
        start_time: datetime object
            datetime object containing the record to plot for
            the given beam
        ax: matplotlib.axes
            axes object for another way of plotting
            Default: None
        normalized: bool
            normalizes the parameter data with the associated
            power 0 value for the given gate number
            default: True
        real_color: str
            line color of the real part of the paramter
            default: red
        plot_blanked: bool
            boolean to determine if blanked lags should be plotted
            default: False
        blank_marker: str
            the marker symbol of blanked lags
            default: o - dot
        imaginary_color: str
            line color of the imaginary part of the parameter
            default: blue
        legend: bool
            produces a standard legend
            default: True
        kwargs: dict
            are applied to the real and imaginary plots


        Notes
        -----
        plot blanked lags as hollow circles

        Raises
        ------
        UnkownParameterError
        IncorrectPlotMethodError
        OutOfRangeGateError

        Returns
        -------

        """
        # If an axes object is not passed in then store
        # the equivalent object in matplotlib. This allows
        # for variant matplotlib plotting styles.
        if not ax:
            ax = plt.gca()
        # Determine if a DmapRecord was passed in, instead of a list
        try:
            # because of partial records we need to find the first
            # record that has that parameter
            index_first_match = next(i for i, d in enumerate(dmap_data)
                                     if parameter in d)
            if isinstance(dmap_data[index_first_match][parameter],
                          DmapArray) or\
               isinstance(dmap_data[index_first_match][parameter],
                          DmapScalar):
                dmap_data = dmap2dict(dmap_data)
        except StopIteration:
            raise plot_exceptions.UnknownParameterError(parameter)

        cls.dmap_data = dmap_data
        check_data_type(cls.dmap_data, parameter, 'array', index_first_match)

        # search over the records to find the correct beam, and scan/time
        # to plot the corresponding ACF/XCF plot
        scan_count = 0
        re = []
        im = []
        for record in cls.dmap_data:
            if record['bmnum'] == beam_num:
                time = time2datetime(record)
                if start_time is not None:
                    if time.day != start_time.day or \
                       time.month != start_time.month or \
                       time.year != start_time.year:
                        raise plot_exceptions.IncorrectDateError(time,
                                                                 start_time)

                if (scan_count == scan_num and start_time is None) or\
                   start_time < time:
                    if gate_num >= record['nrang'] or gate_num < 0:
                        raise plot_exceptions.\
                                OutOfRangeGateError(parameter, gate_num,
                                                    record['nrang'])
                    # get the difference to get the lag number
                    lags = [lag[1] - lag[0] for lag in record['ltab']]
                    # remove the last lag as it is the zeroth lag
                    lags = lags[:-1]

                    # grab corresponding real and imaginary parts
                    # of the ACF/XCF field for a given gate number
                    re = [x[0] for x in record[parameter][gate_num]]
                    im = [x[1] for x in record[parameter][gate_num]]

                    # calculate the blank lags for this record and gate
                    blanked_lags = cls.__blanked_lags(record, lags, gate_num)
                    # Need to deep clone so we keep the values
                    blank_re = copy.deepcopy(re)
                    blank_im = copy.deepcopy(im)
                    lag_num = 0
                    lag_idx = 0
                    lags_len = len(lags)
                    # Search for missing lags
                    # Note: had to use while loop do to insert method
                    while lag_idx < lags_len:
                        if lag_num in blanked_lags:
                            # to remove lines going through the points
                            re[lag_idx] = np.nan
                            im[lag_idx] = np.nan
                        if lags[lag_idx] != lag_num:
                            lags.insert(lag_num, lag_num)
                            # increase length by one due to insert
                            lags_len += 1
                            # masking the values so it can create
                            # gaps in the data
                            re.insert(lag_num, np.nan)
                            im.insert(lag_num, np.nan)
                            # ensures the points line up on the plot
                            blank_re.insert(lag_num, np.nan)
                            blank_im.insert(lag_num, np.nan)

                        lag_idx += 1
                        if lag_idx < lags_len:
                            if lag_num != lags[lag_idx]:
                                if lag_num > lags[lag_idx]:
                                    lag_num = lags[lag_idx]
                                else:
                                    lag_num += 1
                    # once we got the data break free!!
                    break
                scan_count += 1
        if record['cp'] == 503:
            warnings.warn("Please note this data is from Tauscan which has"
                          "different lag properties to other control problems"
                          "this ACF plot may not be correct."
                          "Please contact the PI of the radar to"
                          "confirm if the data looks correct.")
        if re == [] or im == []:
            if gate_num > 0 and gate_num < record['nrang']:
                time = time2datetime(record)
                raise plot_exceptions.NoDataFoundError(parameter, beam_num,
                                                       None, time,
                                                       record['bmnum'])
            else:
                raise plot_exceptions.OutOfRangeGateError(parameter, gate_num,
                                                          record['nrang'])

        if normalized:
            blank_re /= record['pwr0'][gate_num]
            blank_im /= record['pwr0'][gate_num]

            re /= record['pwr0'][gate_num]
            im /= record['pwr0'][gate_num]

        # generates gaps where there are nan's
        masked_re = np.ma.array(re)
        masked_re = np.ma.masked_where(np.isnan(masked_re), masked_re)

        masked_im = np.ma.array(im)
        masked_im = np.ma.masked_where(np.isnan(masked_im), masked_im)

        # plot real and imaginary
        ax.plot(lags, masked_im, marker='o', color=imaginary_color,
                label='Imaginary', **kwargs)
        ax.plot(lags, masked_re, marker='o', color=real_color,
                label='Real', **kwargs)

        # plot blanked lags
        if plot_blank:
            for blank in blanked_lags:
                # I use scatter here to make points not lines
                # also shows up in the legend nicer
                line_re = ax.scatter(blank, blank_re[lags.index(blank)],
                                     edgecolors=real_color, facecolors='white',
                                     marker=blank_marker)
                line_im = ax.scatter(blank, blank_im[lags.index(blank)],
                                     edgecolors=imaginary_color,
                                     facecolors='white', marker=blank_marker)

            # generate generic legend
            if legend and blanked_lags != []:
                line_re.set_label('Real Blanked')
                line_im.set_label('Imaginary Blanked')
        ax.legend()
        ax.set_ylabel(parameter)
        ax.set_xlabel('Lag Number')
        ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True))
        radar_name = SuperDARNRadars.radars[cls.dmap_data[0]['stid']].name
        title = "{date} UT {radar} Beam {beam}, Gate {gate}, Control "\
                "Program: {cpid}"\
                "".format(radar=radar_name, beam=beam_num, gate=gate_num,
                          date=time.strftime("%Y %b %d %H:%M"),
                          cpid=record['cp'])
        ax.set_title(title)

    @classmethod
    def __blanked_lags(cls, record: dict, lags: list, gate: int):
        """
        Determines the blanked lags in the data record

        Parameters
        ----------
        record : dict
            data record containing the lag and ACF/SCF data
        lags : list
            the lag number list where elments are: ltab[i][1] - ltab[i][0]
        gate : int
            gate number
        """
        sorted_lags = lags
        lags_pair = [lag for lags in record['ltab'] for lag in lags]
        sorted_lags.sort()
        lags_pair.sort()
        tau_per_txpl = record['mpinc'] / record['txpl']
        blanked_samples = [lag_pair * tau_per_txpl for lag_pair in lags_pair]
        blanked_lags = []
        for blank in blanked_samples:
            blanked_lags.extend([blank, blank + 1])

        blanked_lag_numbers = []
        for s_lag in sorted_lags:
            for pair in record['ltab']:
                # find which lags didn't get changed
                # in the sorting
                if s_lag == (pair[1] - pair[0]):
                    sample1 = tau_per_txpl * pair[0] + gate + 1 * \
                        record['lagfr'] / record['txpl']
                    sample2 = tau_per_txpl * pair[1] + gate + 1 * \
                        record['lagfr'] / record['txpl']
                    if sample1 in blanked_lags or sample2 in blanked_lags:
                        blanked_lag_numbers.append(s_lag)

        return blanked_lag_numbers
