from itertools import product
import dask.array as da
import numpy as np
import xarray as xr
import pylab as plt
import torch
from tqdm import tqdm

import warnings

# Suppress specific warning
warnings.filterwarnings("ignore", message=".*Consider loading the data with Dask directly.*")


def MFP_one_BH(ds,zrange=[0,100],freqrange=[100,200],xrange=[0,50],vrange=[3000,4000],dz=1,dx=1,dv=100):
    '''
    Perform Matched Field Processing (MFP) to compute beamforming power over a range of 
    spatial coordinates, velocities, and frequencies.

    :param ds: Input data as an xarray DataArray containing waveform data. 
               Must include 'distance' and 'time' coordinates.
    :type ds: xarray.DataArray

    :param zrange: Range of depths to compute (start and end), default is [0, 100].
    :type zrange: list of float or int

    :param freqrange: Frequency range for processing in Hz (start and end), default is [100, 200].
    :type freqrange: list of float

    :param xrange: Range of horizontal x-coordinates to search (start and end), default is [0, 50].
    :type xrange: list of float or int

    :param vrange: Velocity range to search in m/s (start and end), default is [1000, 6000].
    :type vrange: list of float or int

    :param dz: Spacing between grid points in the depth dimension (z), default is 1.
    :type dz: float or int

    :param dx: Spacing between grid points in the horizontal x dimension, default is 1.
    :type dx: float or int

    :param dv: Spacing between velocities, default is 100.
    :type dv: float or int

    :return: A 3D xarray DataArray containing the computed beampower values. The dimensions are
             'v' (velocity), 'z' (depth), and 'x' (horizontal position).
    :rtype: xarray.DataArray
    '''
    
    # Extract data from dataarray
    sampling_rate=((ds.time[1]-ds.time[0]).values.astype(float)*10**-9)
    # Compute FFT of all the receiver
    waveform_spectra=torch.fft.fft(torch.from_numpy(ds.values.T),axis=1).to(dtype=torch.complex128)
    freqs = torch.fft.fftfreq(len(ds.time),sampling_rate)
    omega = 2 * torch.pi * freqs
    # Build grid for search
    x_coords = torch.arange(xrange[0], xrange[1] + dx, dx)
    z_coords = torch.arange(zrange[0], zrange[1] + dz, dz)
    v_coords = torch.arange(vrange[0], vrange[1] + dv, dv)
    gridpoints = torch.tensor(list(product(x_coords, z_coords)))
    
    stations=torch.tensor(np.array([np.zeros(len(ds.distance)),ds.distance.values]).T)

    distances_to_all_gridpoints = torch.linalg.norm(gridpoints[:, None, :] - stations[None, :, :], axis=2)
    # limit to frequency band of interest for
    # a) speed# limit to frequency band of interest for
    # a) speed-up
    # b) focusing on specific frequencies
    freq_idx = torch.where((freqs > freqrange[0]) & (freqs < freqrange[1]))[0]
    omega_lim = omega[freq_idx]
    waveform_spectra_lim = waveform_spectra[:, freq_idx]

    bp=[]
    for vel in tqdm(v_coords):
        

        traveltimes = distances_to_all_gridpoints / vel
        
        # Green's functions between all stations and all grid points
        # within selected frequency band
        # G = exp(-iωt)
        greens_functions = torch.exp(-1j * omega_lim[None, None, :] * traveltimes[:, :, None])
        # move critical part to dask
        greens_functions_dask = da.from_array(greens_functions.numpy(), chunks='auto')

        S = (greens_functions_dask[:, :, None, :]*greens_functions_dask.conj()[:, None, :, :])

        # this assumes that K can be computed without memory issues.
        # If this is not the case, you can also use dask and chunk K.
        K = waveform_spectra_lim[:, None, :] * waveform_spectra_lim.conj()[None, :, :]

        # exclude autocorrelations by filling diagonal with complex zeros
        diag_idxs = torch.arange(K.shape[0])
        zero_spectra = torch.zeros(omega_lim.shape, dtype=torch.cdouble)
        K[diag_idxs, diag_idxs, :] = zero_spectra

        K = da.from_array(K.numpy())

        beampowers_d = da.einsum("gjiw, ijw -> g", S, K).real

        beampowers = beampowers_d.compute()
        bp.append(beampowers.reshape(len(x_coords),len(z_coords)).T)

    bp=np.stack(bp)
    bp /= (len(freq_idx)*len(ds.distance))

    res=xr.DataArray(np.stack(bp),dims=['v','z','x'])
    res['v']=v_coords
    res['x']=x_coords
    res['z']=z_coords
    return res

def MFP_BH(ds,stations_pos,freqrange=[100,200],xrange=[0,50],yrange=[0,50],zrange=[0,100],vrange=[3000,4000],dx=1,dy=1,dz=1,dv=100):
    '''
    Perform Matched Field Processing (MFP) to compute beamforming power over a range of 
    spatial coordinates, velocities, and frequencies.

    :param ds: Input data as an xarray DataArray containing waveform data. 
               Must include 'distance' and 'time' coordinates.
    :type ds: xarray.DataArray

    :param zrange: Range of depths to compute (start and end), default is [0, 100].
    :type zrange: list of float or int

    :param freqrange: Frequency range for processing in Hz (start and end), default is [100, 200].
    :type freqrange: list of float

    :param xrange: Range of horizontal x-coordinates to search (start and end), default is [0, 50].
    :type xrange: list of float or int

    :param vrange: Velocity range to search in m/s (start and end), default is [1000, 6000].
    :type vrange: list of float or int

    :param dz: Spacing between grid points in the depth dimension (z), default is 1.
    :type dz: float or int

    :param dx: Spacing between grid points in the horizontal x dimension, default is 1.
    :type dx: float or int

    :param dv: Spacing between velocities, default is 100.
    :type dv: float or int

    :return: A 3D xarray DataArray containing the computed beampower values. The dimensions are
             'v' (velocity), 'z' (depth), and 'x' (horizontal position).
    :rtype: xarray.DataArray
    '''
    
    # Extract data from dataarray
    sampling_rate=((ds.time[1]-ds.time[0]).values.astype(float)*10**-9)
    # Compute FFT of all the receiver
    waveform_spectra=torch.fft.fft(torch.from_numpy(ds.values.T),axis=1).to(dtype=torch.complex128)
    freqs = torch.fft.fftfreq(len(ds.time),sampling_rate)
    omega = 2 * torch.pi * freqs
    # Build grid for search
    x_coords = torch.arange(xrange[0], xrange[1] + dx, dx)
    y_coords = torch.arange(yrange[0], yrange[1] + dy, dy)
    z_coords = torch.arange(zrange[0], zrange[1] + dz, dz)
    v_coords = torch.arange(vrange[0], vrange[1] + dv, dv)
    gridpoints = torch.tensor(list(product(x_coords,y_coords,z_coords)))
    
    stations=torch.tensor(stations_pos)

    distances_to_all_gridpoints = torch.linalg.norm(gridpoints[:, None, :] - stations[None, :, :], axis=2)
    # limit to frequency band of interest for
    # a) speed# limit to frequency band of interest for
    # a) speed-up
    # b) focusing on specific frequencies
    freq_idx = torch.where((freqs > freqrange[0]) & (freqs < freqrange[1]))[0]
    omega_lim = omega[freq_idx]
    waveform_spectra_lim = waveform_spectra[:, freq_idx]

    bp=[]
    for vel in tqdm(v_coords):
        
        traveltimes = distances_to_all_gridpoints / vel
        
        # Green's functions between all stations and all grid points
        # within selected frequency band
        # G = exp(-iωt)
        greens_functions = torch.exp(-1j * omega_lim[None, None, :] * traveltimes[:, :, None])
        # move critical part to dask
        greens_functions_dask = da.from_array(greens_functions.numpy(), chunks='auto')

        S = (greens_functions_dask[:, :, None, :]*greens_functions_dask.conj()[:, None, :, :])

        # this assumes that K can be computed without memory issues.
        # If this is not the case, you can also use dask and chunk K.
        K = waveform_spectra_lim[:, None, :] * waveform_spectra_lim.conj()[None, :, :]

        # exclude autocorrelations by filling diagonal with complex zeros
        diag_idxs = torch.arange(K.shape[0])
        zero_spectra = torch.zeros(omega_lim.shape, dtype=torch.cdouble)
        K[diag_idxs, diag_idxs, :] = zero_spectra

        K = da.from_array(K.numpy())

        beampowers_d = da.einsum("gjiw, ijw -> g", S, K).real

        beampowers = beampowers_d.compute()
        bp.append(beampowers.reshape(len(x_coords),len(y_coords),len(z_coords)))

    bp=np.stack(bp)
    bp /= (len(freq_idx)*len(ds.distance))

    res=xr.DataArray(np.stack(bp),dims=['v','x','y','z'])
    res['x']=x_coords
    res['y']=y_coords
    res['z']=z_coords
    res['v']=v_coords

    return res