from __future__ import division, print_function
import numpy
from os import path
from . import tls_constants
from .helpers import running_median, transit_mask


def FAP(SDE):
    """Returns FAP (False Alarm Probability) for a given SDE"""
    data = numpy.genfromtxt(
        path.join(tls_constants.resources_dir, "fap.csv"),
        dtype="f8, f8",
        delimiter=",",
        names=["FAP", "SDE"],
    )
    return data["FAP"][numpy.argmax(data["SDE"] > SDE)]


def rp_rs_from_depth(depth, law, params):
    """Takes the maximum transit depth, limb-darkening law and parameters
    Returns R_P / R_S (ratio of planetary to stellar radius)
    Source: Heller 2019, https://arxiv.org/abs/1901.01730"""

    # Validations:
    # - LD law must exist
    # - All parameters must be floats or ints
    # - All parameters must be given in the correct quanitity for the law

    if len(params) == 1:
        params = float(params[0])

    if not isinstance(params, (float, int)) and not all(
        isinstance(x, (float, int)) for x in params
    ):
        raise ValueError("All limb-darkening parameters must be numbers")

    laws = "linear, quadratic, squareroot, logarithmic, nonlinear"
    if law not in laws:
        raise ValueError("Please provide a supported limb-darkening law:", laws)

    if law == "linear" and not isinstance(params, float):
        raise ValueError("Please provide exactly one parameter")

    if law in "quadratic, logarithmic, squareroot" and len(params) != 2:
        raise ValueError("Please provide exactly two limb-darkening parameters")

    if law == "nonlinear" and len(params) != 4:
        raise ValueError("Please provide exactly four limb-darkening parameters")

    # Actual calculations of the return value
    if law == "linear":
        return (depth * (1 - params / 3)) ** (1 / 2)

    if law == "quadratic":
        return (depth * (1 - params[0] / 3 - params[1] / 6)) ** (1 / 2)

    if law == "squareroot":
        return (depth * (1 - params[0] / 3 - params[1] / 5)) ** (1 / 2)

    if law == "logarithmic":
        return (depth * (1 + 2 * params[1] / 9 - params[0] / 3)) ** (1 / 2)

    if law == "nonlinear":
        return (
            depth
            * (1 - params[0] / 5 - params[1] / 3 - 3 * params[2] / 7 - params[3] / 2)
        ) ** (1 / 2)


def pink_noise(data, width):
    std = 0
    datapoints = len(data) - width + 1
    for i in range(datapoints):
        std += numpy.std(data[i : i + width]) / width ** 0.5
    return std / datapoints


def period_uncertainty(periods, power):
    # Determine estimate for uncertainty in period
    # Method: Full width at half maximum
    try:
        # Upper limit
        index_highest_power = numpy.argmax(power)
        idx = index_highest_power
        while True:
            idx += 1
            if power[idx] <= 0.5 * power[index_highest_power]:
                idx_upper = idx
                break
        # Lower limit
        idx = index_highest_power
        while True:
            idx -= 1
            if power[idx] <= 0.5 * power[index_highest_power]:
                idx_lower = idx
                break
        period_uncertainty = 0.5 * (periods[idx_upper] - periods[idx_lower])
    except:
        period_uncertainty = float("inf")
    return period_uncertainty


def spectra(chi2, oversampling_factor, chi2_baseline):
    SR = numpy.min(chi2) / chi2
    SDE_raw = (1 - numpy.mean(SR)) / numpy.std(SR)

    # Scale SDE_power from 0 to SDE_raw
    power_raw = SR - numpy.mean(SR)  # shift down to the mean being zero
    scale = SDE_raw / numpy.max(power_raw)  # scale factor to touch max=SDE_raw
    power_raw = power_raw * scale

    # Detrended SDE, named "power"
    kernel = int(oversampling_factor * tls_constants.SDE_MEDIAN_KERNEL_SIZE)
    if kernel % 2 == 0:
        kernel = kernel + 1
    if len(power_raw) > 2 * kernel:
        my_median = running_median(power_raw, kernel)
        power = power_raw - my_median
        # Re-normalize to range between median = 0 and peak = SDE
        # shift down to the mean being zero
        power = power - numpy.mean(power)
        power = power_neighbour_correction(power, kernel, chi2, chi2_baseline)
        SDE = numpy.max(power / numpy.std(power))
        # scale factor to touch max=SDE
        scale = SDE / numpy.max(power)
        power = power * scale
    else:
        power = power_raw
        SDE = SDE_raw

    return SR, power_raw, power, SDE_raw, SDE

def power_neighbour_correction(sde, kernel, chi2, chi2_baseline):
    """Returns scaled power spectra taking into account the neighbours used with the given kernel that are below the
    residuals baseline"""
    idx = numpy.arange(kernel) + numpy.arange(len(sde) - kernel + 1)[:, None]
    chi2_idx = chi2[idx]
    chi2_idx_gtbaseline = [len(numpy.argwhere(chi2_window < chi2_baseline).flatten()) for chi2_window in chi2_idx]
    # Append the first/last value at the beginning/end to match the length of
    # data and returned median
    first_values = chi2_idx_gtbaseline[0]
    last_values = chi2_idx_gtbaseline[-1]
    missing_values = len(sde) - len(chi2_idx_gtbaseline)
    values_front = int(missing_values * 0.5)
    values_end = missing_values - values_front
    chi2_idx_gtbaseline = numpy.append(numpy.full(values_front, first_values), chi2_idx_gtbaseline)
    chi2_idx_gtbaseline = numpy.append(chi2_idx_gtbaseline, numpy.full(values_end, last_values))
    sde = sde * chi2_idx_gtbaseline / kernel
    return sde


def model_lightcurve(transit_times, period, t, model_transit_single):
    """Creates the model light curve for the full unfolded dataset"""

    # Append one more transit after and before end of nominal time series
    # to fully cover beginning and end with out of transit calculations
    earlier_tt = transit_times[0] - period
    extended_transit_times = numpy.append(earlier_tt, transit_times)
    next_tt = transit_times[-1] + period
    extended_transit_times = numpy.append(extended_transit_times, next_tt)
    full_x_array = numpy.array([])
    full_y_array = numpy.array([])
    rounds = len(extended_transit_times)
    internal_samples = (
        int(len(t) / len(transit_times))
    ) * tls_constants.OVERSAMPLE_MODEL_LIGHT_CURVE

    # Append all periods
    for i in range(rounds):
        xmin = extended_transit_times[i] - period / 2
        xmax = extended_transit_times[i] + period / 2
        x_array = numpy.linspace(xmin, xmax, internal_samples)
        full_x_array = numpy.append(full_x_array, x_array)
        full_y_array = numpy.append(full_y_array, model_transit_single)

    if numpy.all(numpy.isnan(full_x_array)):
        return None, None
    else:  # Determine start and end of relevant time series, and crop it
        start_cadence = numpy.nanargmax(full_x_array > min(t))
        stop_cadence = numpy.nanargmax(full_x_array > max(t))
        full_x_array = full_x_array[start_cadence:stop_cadence]
        full_y_array = full_y_array[start_cadence:stop_cadence]
        model_lightcurve_model = full_y_array
        model_lightcurve_time = full_x_array
        return model_lightcurve_model, model_lightcurve_time


def all_transit_times(T0, t, period):
    """Return all mid-transit times within t"""

    if T0 < min(t):
        transit_times = [T0 + period]
    else:
        transit_times = [T0]
    previous_transit_time = transit_times[0]
    transit_number = 0
    while True:
        transit_number = transit_number + 1
        next_transit_time = previous_transit_time + period
        if next_transit_time < (numpy.min(t) + (numpy.max(t) - numpy.min(t))):
            transit_times.append(next_transit_time)
            previous_transit_time = next_transit_time
        else:
            break
    return transit_times


def calculate_transit_duration_in_days(t, period, transit_times, duration):
    """Return estimate for transit duration in days"""

    # Difference between (time series duration / period) and epochs
    transit_duration_in_days_raw = (
        duration * calculate_stretch(t, period, transit_times) * period
    )

    # Correct the duration for gaps in the data
    transit_duration_in_days = transit_duration_in_days_raw * calculate_fill_factor(t)

    return transit_duration_in_days


def calculate_stretch(t, period, transit_times):
    """Return difference between (time series duration / period) and epochs
        Example: 
        - Time series duration = 100 days
        - Period = 40 days
        - Epochs = 2 at t0s = [30, 70] days
        ==> stretch = (100 / 40) / 2 = 1.25"""

    duration_timeseries = (numpy.max(t) - numpy.min(t)) / period
    epochs = len(transit_times)
    stretch = duration_timeseries / epochs
    return stretch


def calculate_fill_factor(t):
    """Return the fraction of existing cadences, assuming constant cadences"""

    average_cadence = numpy.median(numpy.diff(t))
    span = max(t) - min(t)
    theoretical_cadences = span / average_cadence
    fill_factor = (len(t) - 1) / theoretical_cadences
    return fill_factor


def count_stats(t, y, transit_times, transit_duration_in_days):
    """Return:
    * in_transit_count:     Number of data points in transit (phase-folded)
    * after_transit_count:  Number of data points in a bin of transit duration, 
                            after transit (phase-folded)
    * before_transit_count: Number of data points in a bin of transit duration, 
                            before transit (phase-folded)
    """
    in_transit_count = 0
    after_transit_count = 0
    before_transit_count = 0

    for mid_transit in transit_times:
        T0 = (
            mid_transit - 1.5 * transit_duration_in_days
        )  # start of 1 transit dur before ingress
        T1 = mid_transit - 0.5 * transit_duration_in_days  # start of ingress
        T4 = mid_transit + 0.5 * transit_duration_in_days  # end of egress
        T5 = (
            mid_transit + 1.5 * transit_duration_in_days
        )  # end of egress + 1 transit dur

        if T0 > min(t) and T5 < max(t):  # inside time
            idx_intransit = numpy.where(numpy.logical_and(t > T1, t < T4))
            idx_before_transit = numpy.where(numpy.logical_and(t > T0, t < T1))
            idx_after_transit = numpy.where(numpy.logical_and(t > T4, t < T5))
            points_in_this_in_transit = len(y[idx_intransit])
            points_in_this_before_transit = len(y[idx_before_transit])
            points_in_this_after_transit = len(y[idx_after_transit])

            in_transit_count += points_in_this_in_transit
            before_transit_count += points_in_this_before_transit
            after_transit_count += points_in_this_after_transit

    return in_transit_count, after_transit_count, before_transit_count


def intransit_stats(t, y, transit_times, transit_duration_in_days):
    """Return all intransit odd and even flux points"""

    all_flux_intransit_odd = numpy.array([])
    all_flux_intransit_even = numpy.array([])
    all_flux_intransit = numpy.array([])
    all_idx_intransit = numpy.array([])
    per_transit_count = numpy.zeros([len(transit_times)])
    transit_depths = numpy.zeros([len(transit_times)])
    transit_depths_uncertainties = numpy.zeros([len(transit_times)])

    for i in range(len(transit_times)):

        depth_mean_odd = numpy.nan
        depth_mean_even = numpy.nan
        depth_mean_odd_std = numpy.nan
        depth_mean_even_std = numpy.nan

        mid_transit = transit_times[i]
        tmin = mid_transit - 0.5 * transit_duration_in_days
        tmax = mid_transit + 0.5 * transit_duration_in_days
        if numpy.isnan(tmin) or numpy.isnan(tmax):
            idx_intransit = []
            flux_intransit = []
            mean_flux = numpy.nan
        else:
            idx_intransit = numpy.where(numpy.logical_and(t > tmin, t < tmax))
            flux_intransit = y[idx_intransit]
            if len(y[idx_intransit]) > 0:
                mean_flux = numpy.mean(y[idx_intransit])
            else:
                mean_flux = numpy.nan
        intransit_points = numpy.size(y[idx_intransit])
        transit_depths[i] = mean_flux
        if len(y[idx_intransit] > 0):
            transit_depths_uncertainties[i] = numpy.std(y[idx_intransit]) / numpy.sqrt(
                intransit_points
            )
        else:
            transit_depths_uncertainties[i] = numpy.nan
        per_transit_count[i] = intransit_points

        # Check if transit odd/even to collect the flux for the mean calculations
        if i % 2 == 0:  # even
            all_flux_intransit_even = numpy.append(
                all_flux_intransit_even, flux_intransit
            )
        else:  # odd
            all_flux_intransit_odd = numpy.append(
                all_flux_intransit_odd, flux_intransit
            )
        if len(all_flux_intransit_odd) > 0:
            depth_mean_odd = numpy.mean(all_flux_intransit_odd)

            depth_mean_odd_std = numpy.std(all_flux_intransit_odd) / numpy.sum(
                len(all_flux_intransit_odd)
            ) ** (0.5)
        if len(all_flux_intransit_even) > 0:
            depth_mean_even = numpy.mean(all_flux_intransit_even)
            depth_mean_even_std = numpy.std(all_flux_intransit_even) / numpy.sum(
                len(all_flux_intransit_even)
            ) ** (0.5)

    return (
        depth_mean_odd,
        depth_mean_even,
        depth_mean_odd_std,
        depth_mean_even_std,
        all_flux_intransit_odd,
        all_flux_intransit_even,
        per_transit_count,
        transit_depths,
        transit_depths_uncertainties,
    )


def snr_stats(
    t,
    y,
    period,
    duration,
    T0,
    transit_times,
    transit_duration_in_days,
    per_transit_count,
    intransit=None
):
    """Return snr_per_transit and snr_pink_per_transit"""

    snr_per_transit = numpy.zeros([len(transit_times)])
    snr_pink_per_transit = numpy.zeros([len(transit_times)])
    if intransit is None:
        intransit = transit_mask(t, period, 2 * duration, T0)
    flux_ootr = y[~intransit]

    try:
        pinknoise = pink_noise(flux_ootr, int(numpy.mean(per_transit_count)))
    except:
        pinknoise = numpy.nan

    # Estimate SNR and pink SNR
    # Second run because now the out of transit points are known
    if len(flux_ootr) > 0:
        std = numpy.std(flux_ootr)
    else:
        std = numpy.nan
    for i in range(len(transit_times)):
        mid_transit = transit_times[i]
        tmin = mid_transit - 0.5 * transit_duration_in_days
        tmax = mid_transit + 0.5 * transit_duration_in_days
        if numpy.isnan(tmin) or numpy.isnan(tmax):
            idx_intransit = []
            mean_flux = numpy.nan
        else:
            idx_intransit = numpy.where(numpy.logical_and(t > tmin, t < tmax))
            if len(y[idx_intransit]) > 0:
                mean_flux = numpy.mean(y[idx_intransit])
            else:
                mean_flux = numpy.nan

        intransit_points = numpy.size(y[idx_intransit])
        try:
            snr_pink_per_transit[i] = (1 - mean_flux) / pinknoise if mean_flux < 1 else (mean_flux - 1) / pinknoise
            if intransit_points > 0 and not numpy.isnan(std):
                std_binned = std / intransit_points ** 0.5
                snr_per_transit[i] = (1 - mean_flux) / std_binned if mean_flux < 1 else (mean_flux - 1) / std_binned
            else:
                snr_per_transit[i] = 0
                snr_pink_per_transit[i] = 0
        except:
            snr_per_transit[i] = 0
            snr_pink_per_transit[i] = 0

    return snr_per_transit, snr_pink_per_transit
