import numpy as np
import scipy.interpolate as interp
from scipy.integrate import odeint, solve_ivp

def calc_null_curves_t(times, spatial_coords, gXX, gXt, gtt):
    """
    """

    gXX_interp      = interp.RectBivariateSpline(times, spatial_coords, gXX)
    gXt_interp      = interp.RectBivariateSpline(times, spatial_coords, gXt)
    gtt_interp      = interp.RectBivariateSpline(times, spatial_coords, gtt)

    def _dXdt(coord, t, sign,
              coord_min = spatial_coords[0], coord_max = spatial_coords[-1],
              gXX = gXX_interp, gtt = gtt_interp, gXt = gXt_interp):
        if coord <= coord_min or coord >= coord_max:
            return float('nan')
        gXt_val = gXt(t, coord)
        gXX_val = gXX(t, coord)
        gtt_val = gtt(t, coord)
        return ((-gXt_val + sign * np.sqrt((gXt_val ** 2) - gtt_val * gXX_val)) / gXX_val).flatten()[0]

    rays = []
    for i, t in enumerate(times):
        t_back = times[:i + 1][::-1]
        t_fwd  = times[i:]

        ray_fwd  = odeint(_dXdt, 0., t_fwd,  (1.,))[:, 0]
        ray_back = odeint(_dXdt, 0., t_back, (1.,))[:, 0]

        ray_fwd  = np.where(np.isnan(ray_fwd),  spatial_coords[-1], ray_fwd)
        ray_back = np.where(np.isnan(ray_back), spatial_coords[0],  ray_back)

        rays.append(np.concatenate((ray_back[::-1][:-1], ray_fwd)))

    return np.array(rays)

def calc_null_curves(times, spatial_coords, gXX, gXt, gtt, reverse = False):
    """
    Compute null curves along a given axis.

    Shoot a null ray from each point in spatial_coords, in the positive and
    negative direction and compute its trajectory.

    :param array_like times: 1D array of coordinate times at which the spacetime
                             curvature is provided
    :param array_like spatial_coords: 1D array of spatial coordinates
    :param array_like gXX: 2D array containing the values of the XX component
                           of the spacetime metric, where X is the spatial
                           coordinate along which the rays are traced.
                           gXX[i, j] is the value at spacetime point
                           (t=times[i], X=spatial_coords[j]).
    :param array_like gXt: same as gXX, but for the Xt component of the metric
    :param array_like gtt: same as gXX, but for the tt component of the metric
    :return: tuple containing two 2D arrays with, respectively, the positive and
             negative-direction rays. rays[i, j] contains the X-coordinate of
             the ray shot from (t=times[0], X=spatial_coords[i]) at time
             t=times[j]
    """

    gXX_interp      = interp.RectBivariateSpline(times, spatial_coords, gXX)
    gXt_interp      = interp.RectBivariateSpline(times, spatial_coords, gXt)
    gtt_interp      = interp.RectBivariateSpline(times, spatial_coords, gtt)

    if reverse:
        ray_times = times[::-1]
    else:
        ray_times = times

    def event_x_bound_upper(t, x, sign):
        return x[0] - spatial_coords[-1]
    event_x_bound_upper.terminal  = True
    event_x_bound_upper.direction = 1.0

    def event_x_bound_lower(t, x, sign):
        return x[0] - spatial_coords[0]
    event_x_bound_lower.terminal  = True
    event_x_bound_lower.direction = -1.0

    events = [event_x_bound_upper, event_x_bound_lower]

    def dXdt(t, coord, sign, gXX = gXX_interp, gtt = gtt_interp, gXt = gXt_interp):
        gXt_val = gXt(t, coord)
        gXX_val = np.maximum(gXX(t, coord), 1e-6)
        gtt_val = np.minimum(gtt(t, coord), -1e-6)
        return ((-gXt_val + sign * np.sqrt((gXt_val ** 2) - gtt_val * gXX_val)) / gXX_val).flatten()[0]

    rays_pos = np.empty((spatial_coords.shape[0], times.shape[0]))
    rays_neg = np.empty_like(rays_pos)
    for j, X0 in enumerate(spatial_coords):
        for tgt, sign in ((rays_pos, 1.0), (rays_neg, -1.0)):
            ret = solve_ivp(dXdt, (ray_times[0], ray_times[-1]), (X0,),
                            method = 'RK45', t_eval = ray_times, args = (sign,),
                            dense_output = True, events = events, rtol = 1e-6, atol = 1e-8)

            t, x = ret.t, ret.y[0]

            if len(t) < len(ray_times):
                x_ext = np.empty_like(ray_times)
                x_ext[:x.shape[0]] = x
                x_ext[x.shape[0]:] = x[-1] + sign * (ray_times[x.shape[0]:] - t[-1])

                x = x_ext

            tgt[j] = x

    return (ray_times, rays_pos, rays_neg)

def calc_null_coordinates(times, spatial_coords, u_rays, v_rays,
                          gXX, gXt, gtt):
    """
    Compute double-null coordinates (u, v) as functions of
    position and time.

    :param array_like times: 1D array of coordinate times at which the spacetime
                             curvature is provided
    :param array_like spatial_coords: 1D array of spatial coordinates
    :param array_like u_rays: 1D array assigning the values of u on the initial
                              time slice. u_rays[i] is the value of u at
                              X=spatial_coords[i].
    :param array_like v_rays: same as u_rays, but for v.
    :param array_like gXX: 2D array containing the values of the XX component
                           of the spacetime metric, where X is the spatial
                           coordinate along which the rays are traced.
                           gXX[i, j] is the value at spacetime point
                           (t=times[i], X=spatial_coords[j]).
    :param array_like gXt: same as gXX, but for the Xt component of the metric
    :param array_like gtt: same as gXX, but for the tt component of the metric
    :return: tuple containing two 2D arrays with, respectively, values of u and
             v as functions of t and X. u/v[i, j] is the value of u/v at
             t=times[i], X=spatial_coords[j].
    """
    _, X_of_ut, X_of_vt = calc_null_curves(times, spatial_coords, gXX, gXt, gtt)

    u_of_tx = np.empty((times.shape[0], spatial_coords.shape[0]))
    v_of_tx = np.empty_like(u_of_tx)
    for i, t in enumerate(times):
        Xu = X_of_ut[:, i]
        Xv = X_of_vt[:, i]
        u_of_tx[i] = interp.interp1d(Xu, u_rays, fill_value = 'extrapolate')(spatial_coords)
        v_of_tx[i] = interp.interp1d(Xv, v_rays, fill_value = 'extrapolate')(spatial_coords)

    return (u_of_tx, v_of_tx)
