#!/usr/bin/env python
# encoding: utf-8
"""
Analyze multiple periods and starspot evolution
"""
from __future__ import print_function, division, absolute_import

import datetime
import os
import sys

from scipy.ndimage import gaussian_filter
from scipy import signal
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import h5py

from . import eclipsing_binary, config, utils


def find_subpeak_range(peak_freqs, peak_heights, height_ratio):
    """
    Find the highest subpeak and range of significant subpeak frequencies.

    Parameters
    ----------
    peak_freqs : numpy.array_like
        Frequencies of subpeaks
    peak_heights : numpy.array_like
        Heights of subpeaks
    height_ratio : float
        Significant peaks have `peak_heights` > `height_ratio` * highest peak

    Returns
    -------
    freq_range, height_range : None, numpy.ndarray, or 2-tuple of numpy.ndarray
        The frequencies and heights of the subpeaks.
        None if no subpeaks are found,
        2-tuple of length 1 array if one subpeak is found,
        3-tuple of length 1 arrays if two subpeaks are found.
    """
    if len(peak_freqs) == 0:
        return None, None

    # Sort in descending order of peak height
    height_sort = np.argsort(peak_heights)[::-1]
    peak_heights_sorted = peak_heights[height_sort]
    peak_freqs_sorted = peak_freqs[height_sort]

    sig_indices = peak_heights_sorted > height_ratio * peak_heights_sorted[0]

    if np.sum(sig_indices) <= 1:
        height_range = (peak_heights_sorted[sig_indices],
                        peak_heights_sorted[sig_indices])
        freq_range = (peak_freqs_sorted[sig_indices],
                      peak_freqs_sorted[sig_indices])
    else:
        sig_heights = peak_heights_sorted[sig_indices]
        sig_freqs = peak_freqs_sorted[sig_indices]

        subpeak_index_1 = np.argmin(sig_freqs)
        subpeak_index_2 = np.argmax(sig_freqs)

        height_range = (sig_heights[0],
                        sig_heights[subpeak_index_1],
                        sig_heights[subpeak_index_2])
        freq_range = (sig_freqs[0],
                      sig_freqs[subpeak_index_1],
                      sig_freqs[subpeak_index_2])

    return freq_range, height_range


def detect_multiple_periods(results_file, pgram_file, kernel=30,
                            height_ratio=0.3, fix=False,
                            restrict_file='restrict_freqs.csv',
                            plot_all=False, plot_example=False):
    """
    Look for multiple periods and constrain differential rotation.

    Parameters
    ----------
    results_file : str
        Name of the CSV file containing results.
    pgram_file : str
        Name of the HDF5 file containing the periodograms
    kernel : float, optional
        Standard deviation of Gaussian smoothing kernel
    height_ratio : float, optional
        Significant peaks have `peak_heights` > `height_ratio` * highest peak
    fix : bool, optional
        Set to True to only run on EBs that need fixing.
    restrict_file : str, optional
        CSV file containing KIC numbers and period range restrictions.
    plot_all : bool, optional
        Set to True to save plots of periodogram and peaks as PNGs.
    plot_example : bool, optional
        Set to True to plot results for a single example.
    """
    # Load inspection data
    df = pd.read_csv('{}/{}'.format(config.repo_data_dir, results_file))

    # Load periodograms
    h5 = h5py.File('{}/{}'.format(config.data_dir, pgram_file), 'r')

    # Only analyze likely starspot EBs
    sp_mask = df['class'].values == 'sp'

    if fix:
        sp_mask &= df['p_acf_ok'] == 'f'

    kics = df['KIC'][sp_mask].values
    p_rot = df['p_acf'][sp_mask].values

    # -1 filled arrays to hold results
    # First column for highest peak, next two columns for subpeaks
    peak_freqs_1 = np.zeros((len(kics), 3), dtype=np.float64) - 1.
    peak_heights_1 = np.zeros_like(peak_freqs_1, dtype=np.float64) - 1.
    peak_freqs_2 = np.zeros_like(peak_freqs_1, dtype=np.float64) - 1.
    peak_heights_2 = np.zeros_like(peak_freqs_1, dtype=np.float64) - 1.

    total_systems = len(kics)
    print('Finding peaks in {} periodograms...'.format(total_systems))

    if plot_all:
        if not os.path.exists('{}/png'.format(config.data_dir)):
            os.mkdir('{}/png'.format(config.data_dir))

    if plot_example:
        save_results = False
        # HACK: Set first rotation period to rotation period of example
        p_rot[0] = p_rot[np.where(kics == 4751083)[0][0]]
        kics = kics[kics == 4751083]
    else:
        save_results = True

    df_rf = pd.read_csv('{}/{}'.format(config.repo_data_dir, restrict_file))

    for ii in range(len(kics)):

        # Compute periodogram on frequency grid
        freqs = h5['{}/freqs'.format(kics[ii])][:]
        powers = h5['{}/powers'.format(kics[ii])][:]

        # Only look for peaks around p_rot from ACF
        if p_rot[ii] > 45:
            # Expand search area
            f_min = 1 / 200
        else:
            f_min = 1 / 45

        if kics[ii] in df_rf['KIC'].values:
            mask = df_rf['KIC'].values == kics[ii]
            f_min = 1 / df_rf['pmax'].values[mask]
            f_max = 1 / df_rf['pmin'].values[mask]
        else:
            f_max = 1.5 / p_rot[ii]

        search_area = (freqs > f_min) & (freqs < f_max)
        search_freqs = freqs[search_area]
        search_powers = powers[search_area]

        # Smooth periodogram
        smoothed_power = gaussian_filter(search_powers, sigma=kernel)

        # Detect peaks and minima in smoothed periodogram
        relative_maxes = signal.argrelmax(smoothed_power)[0]
        relative_mins = signal.argrelmin(smoothed_power)[0]

        # Sort  peaks by height
        height_sort = np.argsort(smoothed_power[relative_maxes])[::-1]
        powers_sorted = smoothed_power[relative_maxes][height_sort]

        if len(powers_sorted) == 0:
            # No peaks found
            peak_freqs_1[ii] = (-2., -2., -2.)
            peak_heights_1[ii] = (-2., -2., -2.)
            peak_freqs_2[ii] = (-2., -2., -2.)
            peak_heights_2[ii] = (-2., -2., 2.)
            continue

        # Only consider peaks higher than height_ratio * highest peak
        height_cutoff = powers_sorted > height_ratio * powers_sorted[0]

        n_significant_peaks = np.sum(height_cutoff)

        if n_significant_peaks > 1:
            # Take the two highest peaks
            peak_index_1 = relative_maxes[height_sort][height_cutoff][0]
            peak_index_2 = relative_maxes[height_sort][height_cutoff][1]
            peak_indicies = (peak_index_1, peak_index_2)
            store_arrs = [[peak_freqs_1, peak_heights_1],
                          [peak_freqs_2, peak_heights_2]]
        else:
            # There is only one significant subpeak
            peak_indicies = relative_maxes[height_sort][height_cutoff]
            store_arrs = [[peak_freqs_1, peak_heights_1]]

        # Find peaks in unsmoothed, oversampled periodogram
        orig_peaks = signal.argrelmax(search_powers)[0]

        for index, store in zip(peak_indicies, store_arrs):

            # Get indices of adjacent minima
            lt_max = relative_mins < index

            if np.sum(lt_max) == 0:
                # There is no minimum to the left
                index_left = index - 1
            else:
                index_left = relative_mins[np.nonzero(lt_max)[0][-1]]

            if np.sum(~lt_max) == 0:
                # There is no minimum to the right
                index_right = index + 1
            else:
                index_right = relative_mins[np.nonzero(~lt_max)[0][0]]

            # Peak indices within adjacent minima
            in_range = (orig_peaks > index_left) & (orig_peaks < index_right)

            if np.sum(in_range) == 0:
                # No adjacent minima, search all frequencies
                xx = search_freqs[orig_peaks]
                yy = search_powers[orig_peaks]
            else:
                xx = search_freqs[orig_peaks][in_range]
                yy = search_powers[orig_peaks][in_range]

            freq_range, height_range = find_subpeak_range(xx, yy, height_ratio)

            if freq_range is None:
                # No subpeaks found
                continue

            # Store results to arrays
            peak_freqs, peak_heights = store

            if len(freq_range) == 2:
                peak_freqs[ii, :2] = freq_range
                peak_heights[ii, :2] = height_range
            elif len(freq_range) == 3:
                peak_freqs[ii] = freq_range
                peak_heights[ii] = height_range

        sys.stdout.write('\r{:.1f}% complete'.format((ii + 1) * 100 /
                                                     total_systems))
        sys.stdout.flush()

        if plot_example:

            fig, ax1 = plt.subplots()

            ax1.plot(search_freqs, search_powers, lw=0.75, color='k')
            ax1.plot(search_freqs, smoothed_power, lw=1.5, ls=':',
                     color='purple')

            for jj in [1, 2]:
                xx = peak_freqs_1[ii, jj]
                yy = peak_heights_1[ii, jj]
                ax1.scatter(xx, yy, color='k', s=30, zorder=4, marker='x')

                xx = peak_freqs_2[ii, jj]
                yy = peak_heights_2[ii, jj]
                ax1.scatter(xx, yy, color='k', s=30, zorder=4, marker='x')

            ax1.set_xlabel('Frequency (day$^{-1}$)')
            ax1.set_ylabel('Power')
            ax1.set_title('KIC {}'.format(kics[ii]), fontsize=20, y=1.24)
            ax1.minorticks_on()

            # Set x limits
            x_vals = np.concatenate((peak_freqs_1[ii], peak_freqs_2[ii]))
            x_vals = x_vals[x_vals > 0]

            if len(x_vals) > 1:
                xmin = x_vals.min()
                xmax = x_vals.max()
                xdiff = xmax - xmin
                ax1.set_xlim(xmin - 0.5 * xdiff, xmax + 0.5 * xdiff)
            else:
                ax1.set_xlim(search_freqs.min(), search_freqs.max())

            ax2 = ax1.twiny()
            ax2_ticks = ax1.get_xticks()

            def tick_function(x):
                per = 1 / x
                return ['{:.1f}'.format(p) for p in per]

            ax2.set_xticks(ax2_ticks)
            ax2.set_xbound(ax1.get_xbound())
            ax2.set_xticklabels(tick_function(ax2_ticks))
            ax2.tick_params(axis='x', which='major', pad=5)
            ax2.set_xlabel('Period (days)')

            plt.savefig('{}/multipeak_example.pdf'.format(config.data_dir))
            plt.close()

        if plot_all:

            plt.plot(search_freqs, search_powers, lw=0.5)
            plt.plot(search_freqs, smoothed_power, lw=0.5, ls=':')
            plt.axvline(1 / p_rot[ii], lw=0.5, linestyle='--', color='k')

            colors = ['r', 'b']
            for jj in [0, 1, 2]:
                xx = peak_freqs_1[ii, jj]
                yy = peak_heights_1[ii, jj]
                plt.scatter(xx, yy, color=colors[0], s=5, zorder=4)

                xx = peak_freqs_2[ii, jj]
                yy = peak_heights_2[ii, jj]
                plt.scatter(xx, yy, color=colors[1], s=5, zorder=4)

            plt.xlabel('Frequency (day$^{-1}$)')
            plt.ylabel('Power')
            plt.title('KIC {}'.format(kics[ii]))

            plt.xlim(search_freqs.min(), search_freqs.max())
            plt.ylim(0, 1.1 * search_powers.max())

            plt.savefig('{}/png/pgram_KIC{:09d}.png'.format(config.data_dir,
                                                            kics[ii]))
            plt.close()

    print()

    if save_results:
        # Save results to CSV
        df.loc[sp_mask, 'freq_1_1'] = peak_freqs_1[:, 1]
        df.loc[sp_mask, 'freq_1_2'] = peak_freqs_1[:, 2]
        df.loc[sp_mask, 'freq_2_1'] = peak_freqs_2[:, 1]
        df.loc[sp_mask, 'freq_2_2'] = peak_freqs_2[:, 2]
        df.loc[sp_mask, 'height_1_1'] = peak_heights_1[:, 1]
        df.loc[sp_mask, 'height_1_2'] = peak_heights_1[:, 2]
        df.loc[sp_mask, 'height_2_1'] = peak_heights_2[:, 1]
        df.loc[sp_mask, 'height_2_2'] = peak_heights_2[:, 2]

        df.to_csv('{}/{}'.format(config.repo_data_dir, results_file),
                  index=False)


def generate_frequency_periodograms(sp_class='sp', window=1.5,
                                    output_file=None, kic_list=None,
                                    sap_list=None, detrend_list=None):
    """
    Generate periodograms on a uniform frequency grid.
    
    Parameters
    ----------
    sp_class : {'sp', 'spx'}
        Likely or possible starspots
    window : float, optional
        Window around the eclipse to interpolate over
    output_file : str, optional
        Specify an alternate output filename. Default uses today's date.
    kic_list : list, optional
        Only run on these KIC IDs
    detrend_list : list, optional
        These KIC IDs will be detrended with a low order polynomial.
    sap_list : list, optional
        These KIC IDs will use the SAP flux instead of PDC.
    """
    df = utils.collect_results()

    if output_file is None:
        today = '{:%Y%m%d}'.format(datetime.date.today())
        output_file = '{}/grid_pgrams.{}.h5'.format(config.data_dir, today)
    else:
        output_file = '{}/{}'.format(config.data_dir, output_file)

    h5 = h5py.File(output_file, 'w')
    h5.attrs['window'] = window

    kics = df['KIC'].values[df['class'].values == sp_class]

    if kic_list is not None:
        kics = kics[np.in1d(kics, kic_list)]

    total_systems = len(kics)
    print('Computing grid periodograms for {} systems...'.format(total_systems))

    for ii, kic in enumerate(kics):

        if sap_list is not None and kic in sap_list:
            eb = eclipsing_binary.EclipsingBinary.from_kic(kic, use_pdc=False)
        else:
            eb = eclipsing_binary.EclipsingBinary.from_kic(kic)

        if detrend_list is not None and kic in detrend_list:
            eb.normalize(detrend=True)
        else:
            eb.normalize()

        eb.interpolate_over_eclipse(window=window)

        freqs, powers = eb.frequency_periodogram()

        group = h5.create_group(str(kic))

        group.create_dataset('freqs', data=freqs)
        group.create_dataset('powers', data=powers)

        del eb

        sys.stdout.write('\r{:.1f}% complete'.format((ii + 1) * 100 / total_systems))
        sys.stdout.flush()

    print()