#! /usr/bin/env python
"""
utils.py
Functions (and their associated functions) for utilities.
"""
import subprocess
import numpy as np
from astropy.table import Table
from astropy.table import vstack
import os
import gc
from popsycle import synthetic


def add_precision64(input_array, power):
    """
    Need more precision for kdtree to run properly. Convert inputs from
    float32 to float64, and add a random perturbation beginning in the
    nths decimal place.

    Parameters
    ----------
    input_array : float or array (float32)
        Thing that needs more precision.

    power : float
        To what place you want the perturbation.

    Return
    ------
    output_array : float or array (float64)
        Thing that has more precision.

    """
    input_array = np.atleast_1d(input_array)

    # Perturb.
    pert = 10 ** power * (np.random.rand(len(input_array)) - 0.5)

    # Convert to float64.
    output_array = np.atleast_1d(np.float64(input_array))

    # Add the perturbation.
    output_array = output_array + pert

    return output_array


def sample_spherical(npoints, speed, ndim=3):
    """
    Randomly sample points on a sphere.
    I found this code on stackexchange.

    Parameters
    ---------
    npoints : float
        The number of points you want to generate.

    speed : float
        The radius of the sphere (aka the magnitude of the vectors.)

    dim : float
        The dimension of the space in which the sphere is embedded
        (ndim = 3 samples points on a 2-sphere, aka a "normal" sphere)

    Return
    ------
    An array of the vectors.
    """
    # Check that the speed vector is either a float or an array of length npoints
    if type(speed) != int:
        if type(speed) != float:
            if len(speed) != npoints:
                raise ValueError("{speed} must be either an int, float "
                                 "or array of length {npoints}")

    vec = np.random.randn(ndim, npoints)
    vec /= np.linalg.norm(vec, axis=0)
    vec *= speed
    return vec


def return_nearest_gridpoint(grid, x_grid_arr, y_grid_arr, x_data, y_data):
    """
    Algorithm for finding the nearest grid cell on a 2D array given a
    datapoint that falls within the bounds of the 2D array.

    Parameters
    ----------
    grid : 2D numpy array
        2D array with size (len(y_grid_arr), len(x_grid_array))

    x_grid_arr : numpy array
        2D grid indices in the x-dimension

    y_grid_arr : numpy array
        2D grid indices in the y-dimension

    x_data : numpy array
        x-coordinate for data that will be located onto the grid

    y_data : numpy array
        y-coordinate for data that will be located onto the grid

    Output
    ------
    gridpoint_arr : numpy array
        list of nearest cell values on the grid at
        the location of (x_data, y_data)

    """
    # Convert x_data and y_data to array if single data point is received
    x_data = np.atleast_1d(x_data)
    y_data = np.atleast_1d(y_data)

    gridpoint_arr = []
    for x, y in zip(x_data, y_data):
        # Loop through x_data and y_data
        if np.isnan(x) or np.isnan(y):
            # If either x_data or y_data is nan, return nan
            gridpoint = np.nan
        else:
            # Find location on the grid where x_data and y_data are
            # closest to the grid indices
            x_idx = np.argmin(np.abs(x - x_grid_arr))
            y_idx = np.argmin(np.abs(y - y_grid_arr))
            gridpoint = grid[y_idx, x_idx]
        gridpoint_arr.append(gridpoint)

    # Convert gridpoint_arr into numpy array
    gridpoint_arr = np.array(gridpoint_arr)

    # If only a single data point was received, return a single value
    if len(gridpoint_arr) == 1:
        gridpoint_arr = gridpoint_arr[0]

    return gridpoint_arr


def execute(cmd, shell=False):
    """
    Executes a command line instruction, captures the stdout and stderr

    Parameters
    ----------
    cmd : str
        Command line instruction, including any executables and parameters

    Optional Parameters
    -------------------
    shell : bool
        Determines if the command is run through the shell. Default is False.

    Outputs
    -------
    stdout : str
        Contains the standard output of the executed process

    stderr : str
        Contains the standard error of the executed process

    """
    # Split the argument into a list suitable for Popen
    args = cmd.split()
    # subprocess.PIPE indicates that a pipe
    # to the standard stream should be opened.
    process = subprocess.Popen(args,
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE,
                               shell=shell)
    stdout, stderr = process.communicate()
    stdout = stdout.decode('utf-8')
    stderr = stderr.decode('utf-8')

    return stdout, stderr


def calc_normalized_counts(mag):
    """
    Parameters
    ----------
    mag: float
    
    Return
    ------
    counts: int
    
    CHANGE THIS CODE IN THE FUTURE TO TAKE IN DIFFERENT ZERO POINTS! NEEDS ADDITIONAL DOCUMENTATION
    Right now this is only applicable for OGLE I band
    See confluence wiki for where these values come from...
    """
    a = -0.67815949
    b = 15.37993393
    counts = 10 ** (b + a * mag)
    return counts

def combine_refined_events(refined_events_filenames, overwrite=False,
                           output_file='default'):
    """
    Creates a combined refined_events out of multiple refined_events files

    Parameters
    ----------
    refined_events_filenames : list of strs
        Filenames of refined_events tables to be combined

    Optional Parameters
    -------------------

    overwrite : bool
        If set to True, overwrites output files. If set to False, exits the
        function if output files are already on disk.
        Default is False.

    output_file : str
        The name of the final refined_events file.
        If set to 'default', the format will be generated from the first
        filename in refined_events_filenames following:
            combined_refined_events_<photometric_system>_<filt>_<red_law>.fits
    """
    # Check to see that all refined_events tables are unique
    if len(refined_events_filenames) != len(set(refined_events_filenames)):
        raise Exception('Duplicate filename found in '
                        'refined_events_filenames. Exiting...')

    # Loop over filenames, checking that each one exists
    print('Creating combined refined_events')
    refined_events_arr = []
    for filename in refined_events_filenames:
        if not os.path.exists(filename):
            raise Exception(f'{filename} cannot be found. Skipping...')
        print(f'-- Loading {filename}')
        refined_events = Table.read(filename)
        refined_events['original_filename'] = filename
        refined_events_arr.append(refined_events)

    # Combine astropy tables
    combined_refined_events = vstack(refined_events_arr)

    # Create an output_filename following the first filename
    if output_file == 'default':
        base = refined_events_filenames[0]
        output_file = 'combined_'
        output_file += '_'.join(base.split('_')[-5:])

    # Overwrite any exiting file with the same name
    if overwrite and os.path.exists(output_file):
        os.remove(output_file)

    # Save combined table to disk
    print(f'Saving to {output_file}')
    Table(combined_refined_events).write(output_file)

    return output_file

def calc_delta_c_LL(fratio, u0, thetaE):
    """
    Calculate the maximum-ish centroid shift for a luminous
    lens, no neighbors

    Parameters
    ----------
    fratio : flux ratio of the lens to source, i.e. f_L/f_S

    u0 : impact parameter

    thetaE : Einstein radius

    """
    final_delta_array = np.zeros(len(u0))
    final_u_array = np.zeros(len(u0))
    for j in np.arange(len(u0)):
        f = fratio[j]
        u_array = np.linspace(0, u0[j], 20)
        delta_array = np.zeros(len(u_array))
        for i in np.arange(len(u_array)):
            u = u_array[i]
            delta_array[i] = (u - f * u ** 2 * np.sqrt(u ** 2 + 4)) / (2 + u ** 2 + f * u * np.sqrt(u ** 2 + 4)) + (u * f) / (1 + f)
        max_idx = np.argmax(delta_array)
        final_delta_array[j] = delta_array[max_idx]
        final_u_array[j] = u_array[max_idx]
    final_delta_array = final_delta_array * thetaE

    return final_delta_array, final_u_array

def max_delta_c(u0, thetaE):
    max_delta_c_array = np.zeros(len(u0))

    big_idx = np.where(u0 > np.sqrt(2))[0]
    small_idx = np.where(u0 <= np.sqrt(2))[0]

    max_delta_c_array[big_idx] = calc_delta_c(u0[big_idx], thetaE[big_idx])
    max_delta_c_array[small_idx] = calc_delta_c(np.ones(len(small_idx)) * np.sqrt(2),
                                                thetaE[small_idx])

    return max_delta_c_array

def get_u_from_t(u0, t0, tE, t):
    """
    Given the time and separation at closest approach of lens and source
    and the Einstein radius, calculate the separation as a function of time.

    NOTE 1: You need to be consistent with your units for t0, tE, and t,
    i.e. pick whatever you want (days, years, etc.) but be self consistent.

    NOTE 2: There is a positive and negative solution for u.
    We return the positive solution.

    Parameters
    ----------
    u0 : float or array
        Minimum separation of lens and source (normalized to Einstein radius)

    t0 : float or array
        Time of minimum separation of lens and source

    tE : float or array
        Einstein crossing time of microlensing event

    t : float or array
        Time at which you want to calculate the separation

    Return
    ------
    u : float or array
        Separation of lens and source (normalized to Einstein radius)
    """
    tau = (t - t0) / tE
    u = np.sqrt(u0 ** 2 + tau ** 2)
    return u


def get_t_from_u(u0, t0, tE, u):
    """
    Given the time and separation at closest approach of lens and source
    and the Einstein radius, calculate the time as a function of separation.

    NOTE 1: You need to be consistent with your units for t0, tE, and t,
    i.e. pick whatever you want (days, years, etc.) but be self consistent.

    NOTE 2: There is a positive and negative solution for t.
    We return the positive solution.

    Parameters
    ----------
    u0 : float or array
        Minimum separation of lens and source (normalized to Einstein radius)

    t0 : float or array
        Time of minimum separation of lens and source

    tE : float or array
        Einstein crossing time of microlensing event

    u : float or array
        Separation of lens and source (normalized to Einstein radius)

    Return
    ------
    t : float or array
        Time corresponding to the separation u
    """
    t = tE * np.sqrt(u ** 2 - u0 ** 2) + t0
    return t


def calc_new_position(l0, b0, mu_lcosb, mu_b, t):
    """
    Given an initial position and proper motions in l and b,
    calculate the new position at some later time.

    Parameters
    ----------
    l0 : float or array
        Initial longitude, in DEGREES

    b0 : float or array
        Initial latitude, in DEGREES

    mu_lcosb : float or array
        Longitudinal proper motion l * cos(b), in MAS/YEAR

    mu_b : float or array
        Latitudinal roper motion, in MAS/YEAR

    t : float or array
        Time, in DAYS

    Return
    ------
    l : float or array
        Latitude, in DEGREES

    b : float or array
        Longitude, in DEGREES
    """

    cos_b0 = np.cos(np.radians(b0))

    l = l0 + t * mu_lcosb * masyr_to_degday / cos_b0
    b = b0 + t * mu_b * masyr_to_degday

    return l, b


def calc_centroid_shift(glat_S, glon_S, glat_N, glon_N, f_L, f_S, f_N, u):
    """
    Calculate the centroid (astrometric) shift
    for a luminous lens and neighbors

    Parameters
    ----------
    glat_x : float or array
        Longitude of x (L = lens, S = source, N = neighbor centroid)

    glon_x : float or array
        Latitude of x (L = lens, S = source, N = neighbor centroid)

    f_x : float or array
        Flux of x (L = lens, S = source, N = all neighbors)

    u : float or array
        Dimensionless separation

    Return
    ------
    delta_c_obs : float or array
        Magnitude of observed astrometric shift, in mas
    """
    ##########
    # Calculating the centroid shift in the frame of the lens
    ##########

    t1a1_t2a2 = (u ** 2 + 3) / (u * np.sqrt(u ** 2 + 4))
    a1_a2 = calc_magnification(u)

    glat_c_lensed = (t1a1_t2a2 * f_S * glat_S + glat_N * f_N) / (
            a1_a2 * f_S + f_L + f_N)
    glon_c_lensed = (t1a1_t2a2 * f_S * glon_S + glon_N * f_N) / (
            a1_a2 * f_S + f_L + f_N)

    glat_c_unlensed = (glat_S * f_S + glat_N * f_N) / (f_S + f_L + f_N)
    glon_c_unlensed = (glon_S * f_S + glon_N * f_N) / (f_S + f_L + f_N)

    glat_delta_c = glat_c_lensed - glat_c_unlensed
    glon_delta_c = glon_c_lensed - glon_c_unlensed

    delta_c_obs = np.sqrt(glat_delta_c ** 2 + glon_delta_c ** 2 * np.cos(
        np.radians(glat_S)) ** 2)

    # Convert from degrees to mas
    delta_c_obs *= 3.6 * 10 ** 6

    return delta_c_obs


def calc_f(lambda_eff):
    """
    Calculate that coefficient f that multiples E(B-V) to get the
    extinction in magnitudes
    """
    B = get_Alambda_AKs('Damineli16', 0.445)
    V = get_Alambda_AKs('Damineli16', 0.551)
    L = get_Alambda_AKs('Damineli16', lambda_eff)

    f = L * (B - V) ** -1

    return f


def generate_all_isochrones(iso_dir, include_ztf=True):
    """
    Generates all solar metallicity isochrones needed for PopSyCLE simulations

    Parameters
    ----------
    iso_dir : filepath
        Where are the isochrones stored (for PopStar)

    Optional Parameters
    -------------------
    include_ztf : bool
        Determines if ztf filters should be included
        in the generated isochrones. Default is True.
    """
    from popsycle.synthetic import all_filt_list
#     from popstar import evolution, synthetic 
    from spisea import evolution, synthetic
    my_filt_list = all_filt_list
    # Add ztf filters to generated isochrones
    if include_ztf:
        my_filt_list += ['ztf,g', 'ztf,r', 'ztf,i']

    # All possibles ages to two digit precision
    log_age_arr = np.arange(5.01, 10.291, 0.01)
    print('Generating %i isochrones for '
          '%.2f < log_age < %.2f' % (len(log_age_arr),
                                     log_age_arr[0],
                                     log_age_arr[-1]))
    for i, log_age in enumerate(log_age_arr):
        print('Generating log_age = %.2f (%03d / %03d)' % (log_age, i+1, len(log_age_arr)))
        my_iso = synthetic.IsochronePhot(log_age, 0, 10,
                                         evo_model=evolution.MISTv1(),
                                         filters=my_filt_list,
                                         iso_dir=iso_dir)

        # Check that the isochrone has all of the filters in my_filt_list
        # If not, force recreating the isochrone with recomp=True
        my_iso_filters = [f for f in my_iso.points.colnames if 'm_' in f]
        my_filt_list_fmt = ['m_%s' % f.replace(',', '_') for f in my_filt_list]
        if len(set(my_filt_list_fmt) - set(my_iso_filters)) > 0:
            _ = synthetic.IsochronePhot(log_age, 0, 10,
                                        evo_model=evolution.MISTv1(),
                                        filters=my_filt_list,
                                        iso_dir=iso_dir,
                                        recomp=True)
       
            
    
