# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/ripple/markRFI.ipynb (unless otherwise specified).

__all__ = ['find_t', 'mask_sf', 'mask_lf', 'mask_time_rfi', 'mask_nr', 'mask_freq_rfi', 'mask_thr_rfi']

# Cell
# author: Xu Chen 2021.06

import os

from astropy import log
from matplotlib import pyplot as plt

import numpy as np
import h5py
from copy import deepcopy
from tqdm import tqdm

from .markRFI import rms, real_rms, std, real_std, get_startend, find_edge_2sides

def load_hdf5_spec(name, frange = None):
    import h5py
    
    f = h5py.File(name,'r')
    fs = f['S']
    freq = fs['freq'][()]
    if frange is None:
        frange = [0, np.inf]
    use = (freq > frange[0]) & (freq < frange[1])
    if 'T' in fs.keys():
        infield = 'T'
    elif 'Ta' in fs.keys():
        infield = 'Ta'
    elif 'flux' in fs.keys():
        infield = 'flux'
    else:
        raise(ValueError('can not find spec'))
    s2p = fs[infield][()]
    if s2p.shape[0] == 2 or s2p.shape[0] == 1:
        from ..utils.io import PolarMjdChan_to_MjdChanPolar
        s2p = PolarMjdChan_to_MjdChanPolar(s2p)
    f.close()
    return s2p[:,use,:]
    
####################### time rfi ########################

def find_t(data,freq,frange,thr_type = 'input_med_times',
           mean_times = 10,diff_times = 20,rfi_width_lim = 20,ext_add = 0,plot = False,
           pdf = None,xylim = None,plot_norfi=False,axis = 'time'):
    """
    data: 2D arrays
    frange: like [a,b]
    mean_times: mean_thr = median value * mean_times
    diff_times: diff_thr = median value * diff_times
    rfi_width_lim: width limit
    ext_add: extend edge
    """
    if (freq is not None)&(frange is not None):
        if (len(frange) == 2):
            pattern_use = (freq>frange[0])&(freq<frange[1])
            pat_data = data[:,pattern_use]
        else:
            raise ValueError("frange should like [fmin,fmax].")
    else:
        pat_data = data

#     print(f"finding rfi along {axis} axis.")
    if axis == 'time':
        pat_mean = np.nanmean(pat_data,axis = 1)
        t = np.arange(data.shape[0])
        ylabel = 'mean along freq axis'
        xlabel = 'spec number (time axis)'
    elif axis == 'freq':
        pat_mean = np.nanmean(pat_data,axis = 0)
        t = deepcopy(freq)#np.arange(data.shape[1])
        ylabel = 'mean along time axis'
        xlabel = 'freq (MHz)'

    is_timerfi = np.zeros_like(t,dtype = 'bool')

    pat_med = np.nanmedian(pat_mean)

    if pat_med <= 0 and thr_type == 'input_med_times':
        log.warning(f"Median value <= 0. Switch to thr_type = 'input_absmed_times'")
        thr_type = 'input_absmed_times'

    if thr_type == 'input_med_times':
        pass
    else:
        if thr_type == 'input_absmed_times':   
            pat_mean = np.abs(pat_mean)
        elif thr_type == 'input_posimed_times':   
            pat_mean -= np.nanmin(pat_mean)
        pat_med = np.nanmedian(pat_mean)
        
    mean_thr = pat_med * mean_times
    diff_thr = pat_med * diff_times

    pat_mean[np.isnan(pat_mean)] = 0
    is_pat = pat_mean > mean_thr
    if (is_pat.any() == False) and (plot_norfi == True):
        if pdf is not None:
            plt.switch_backend('agg')
        fig,ax = plt.subplots(figsize = (15,3))
        ax.plot(t,pat_mean)
        ax.axhline(mean_thr,c = 'k',linestyle = '--',label='mean_thr')
        ax.grid()
        ax.set_ylabel(ylabel)
        ax.set_xlabel(xlabel)
        if frange is not None:
            ax.set_title(f'freq in {frange} MHz')
        if xylim is not None:
            ax.set_xlim(xylim[0],xylim[1])
            ax.set_ylim(xylim[2],xylim[3])
        if pdf is not None:
            pdf.savefig();plt.close()
        return  is_timerfi

#     try:
    start,end = get_startend(is_pat,rfi_width_lim = rfi_width_lim,ext_sec=0)
#     except ValueError as V:
#         #import traceback
#         #traceback.print_exc()
#         #print(f"{V}.Not found.")
#         return  is_timerfi

    pat_diff = np.abs(np.diff(pat_mean, prepend=0, append=0))
    s,e = get_startend(is_timerfi, exclude = False)
    if len(start) > 0:
        print(f"rfi starts at tn = {start}, ends in tn = {end}")
        for s,e in zip(start,end):
            cond = (pat_diff[s] > diff_thr) & (pat_diff[e] > diff_thr)
            if cond:
                is_timerfi[s:e] = True

        if ext_add > 0 and is_timerfi.any():
            from ..utils.misc import extend_Trues
            is_timerfi = extend_Trues(is_timerfi,axis = 0,ext_frac = 0,ext_add = ext_add)
            s,e = get_startend(is_timerfi, exclude = False)
            print(f"After extension, rfi starts at tn = {s}, ends in tn = {e}")
#     else:
#         print("No True meets width condition.")
    if axis == 'time':
        if is_timerfi[1] == True: is_timerfi[0] = True
        if is_timerfi[-2] == True: is_timerfi[-1] = True

    if plot:
        if pdf is not None:
            plt.switch_backend('agg')
        print(f"Median value is {pat_med}. mean_thr = {mean_thr}. diff_thr = {diff_thr}")
        fig,ax = plt.subplots(figsize = (15,3))
        ax.plot(t,pat_mean)
        ax.axhline(mean_thr,c = 'k',linestyle = '--',label='mean_thr')
        ax.plot(t,is_timerfi*np.max(pat_mean),label='is_timeRFI')
        offset = np.max(pat_mean)
        ax.axhline( - offset,c = 'k',linestyle = '-')
        ax.plot(t,pat_diff[:-1] - offset,label='abs(diff)')
        ax.axhline(diff_thr - offset,c = 'k',linestyle = '--',label='diff_thr')

        ax.grid()
        ax.plot(t[s],pat_mean[s],'r.',label = 'start')
        ax.plot(t[e],pat_mean[e],'g.',label = 'end')
        ax.legend()
        ax.set_ylabel(ylabel)
        ax.set_xlabel(xlabel)
        if frange is not None:
            ax.set_title(f'freq in {frange} MHz')
        if xylim is not None:
            ax.set_xlim(xylim[0],xylim[1])
            ax.set_ylim(xylim[2],xylim[3])
        if pdf is not None:
            pdf.savefig();plt.close()
    return  is_timerfi

def mask_sf(data,freq,frange = None,rms_frange = None,ext_times=1,mask_rms_times=3, is_timerfi = None, **kwargs):
    """
    Other parameters are the same as previous.
    """
    ret = np.zeros_like(data,dtype = 'bool')

    f_use = (freq>frange[0])&(freq<frange[1])
    if is_timerfi is None:
        is_timerfi = find_t(data,freq,frange = frange,**kwargs)

    if is_timerfi.any() == False:
        return ret
    log.info(f"Looking for short-freq time RFI in {frange} ...")

    start,end = get_startend(is_timerfi,ext_sec = 1)
    from ..utils.misc import smooth1d
    for s,e in zip(start,end):
        mspec = np.nanmean(data[s:e,:],axis = 0)
        sm = smooth1d(mspec[f_use],method='gaussian',sigma=10)
        MAX = max(sm)
        fmax = freq[f_use][np.nanargmax(sm)]
        f50 = freq[f_use][sm > MAX / 2]
        if len(f50) > 1:
            w50 = f50[-1] - f50[0]
            rms_thresh = rms(mspec,freq,rms_frange) * mask_rms_times

            mask_frange = find_edge_2sides(mspec,freq,peak_position=fmax-w50/2,step=w50,
                     rms_thresh=rms_thresh,Print=False,ext_times=ext_times,small_rfi_times=0)
            print(f"tn = {[s,e]} mask frange:",mask_frange)

            mask_use =  (freq>mask_frange[0])&(freq<mask_frange[1])
            ret[s:e,mask_use] = True
        else:
            ret[s:e,:] = True

    ext_add = kwargs['ext_add']
    if ext_add is not None:
        from ..utils.misc import extend_Trues
        ret = extend_Trues(ret,axis = -1,ext_frac = 0,ext_add = ext_add)

    print("Found :D")
    return ret

def mask_lf(data,freq, frange = None, rms_frange = None, mask_rms_times = -1,**kwargs):
    """
    Other parameters are the same as previous.
    """
    ret = np.zeros_like(data,dtype = 'bool')

    is_timerfi = find_t(data,freq,frange =frange, **kwargs)
    if is_timerfi.any() == False:
        print("No long-freq time rfi is found.")
        return ret
    log.info(f"Looking for long-freq time RFI ...")

    if mask_rms_times == -1:
        ret[is_timerfi,:] = True
        print("Mask the whole spec with rfi :(")
    elif mask_rms_times == 0:
        mask_use = (freq>frange[0])&(freq<frange[1])
        ret[is_timerfi,:] = True
        ret[:,~mask_use] = False
        print(f"Mask spec with rfi in input freq range {frange}")
    elif mask_rms_times > 0:
        # compute rms
        data2 = deepcopy(data)
        data[is_timerfi] = np.nan
        tmin = np.nanargmin(np.nanmean(data2, axis = 1))
        spec = data[tmin]
        
        rms_thresh = rms(spec,freq,rms_frange) * mask_rms_times
        mask_use = data[is_timerfi,:] > rms_thresh
        ret[is_timerfi,:] = mask_use
        print(f"Mask spec with rfi above rms thr {rms_thresh}")
    else:
        raise ValueError("mask_rms_times should be -1, 0, or >0 !")

    ext_add = kwargs['ext_add']
    if ext_add is not None:
        from ..utils.misc import extend_Trues
        ret = extend_Trues(ret,axis = -1,ext_frac = 0,ext_add = ext_add)

    print("Found :D")
    return ret

# Cell
def mask_time_rfi(data,freq,rtype = 'short-freq',frange = None, file=None,
                  frange_step = None, plot_norfi = False, **kwargs):
    """
    rtype: 'short-freq','long-freq'
    frange: freq range
    file: freq range exists short-freq time rfi npy filename
    frange_step: if frange is None and file is None, cycle in whole freq band
    Other parameters are the same as previous.
    """

    ret = np.zeros_like(data,dtype = 'bool')

    if rtype == 'short-freq':
        if file is not None or frange_step is not None:
            if file is not None:
                if not os.path.exists(file):
                    raise ValueError(f"can not find the input file path: {file}")
                franges = np.load(file)
                if franges.shape[1] != 2:
                    raise ValueError("time RFI freq shape must like (n,2)")
            else:
                frq1 = np.arange(freq[0],freq[-1],frange_step//2)[:-1]
                frq2 = np.hstack((np.arange(freq[0]+frange_step,freq[-1],frange_step//2),freq[-1]))
                franges = np.vstack((frq1,frq2)).T

            for nf in tqdm(range(franges.shape[0])):
                ret |= mask_sf(data,freq,frange = franges[nf],**kwargs)

        elif len(frange) == 2 and (frange[1] - frange[0]) > 0:
            ret = mask_sf(data,freq,frange,**kwargs)
        else:
            raise ValueError("frange should like [fmin,fmax] or a string (RFI npy filepath) or a certain frange_step.")

        if ret.any() == False:
            print("No short-freq time rfi is found.")

    elif rtype == 'long-freq':
        if len(frange) == 2:
            ret = mask_lf(data,freq, frange,  **kwargs)
        else:
            raise ValueError("frange should like [fmin,fmax].")

    print("Finish")
    return ret

def mask_nr(data,freq, rms_frange = None, mask_rms_times = 0, **kwargs):
    """
    Other parameters are the same as previous.
    """
    ret = np.zeros_like(data,dtype = 'bool')

    log.info(f"Looking for long-time narrowband freq RFI ...")
    is_freqrfi = find_t(data,freq,axis = 'freq', **kwargs)
    if is_freqrfi.any() == False:
        print("No long-time freq rfi is found.")
        return ret

    if mask_rms_times == 0:
        ret[:,is_freqrfi] = True
        print("Mask the whole chanels with rfi :P")
    elif mask_rms_times > 0:
        # compute rms
        tmin = np.nanargmin(np.nanmean(data, axis = 1))
        spec = data[tmin]
        rms_thresh = rms(spec,freq,rms_frange) * mask_rms_times
        mask_use = data[:,is_freqrfi] > rms_thresh
        ret[:,is_freqrfi] = mask_use
        print(f"Mask channels with rfi above rms thr {rms_thresh}")

    ext_add = kwargs['ext_add']
    if ext_add is not None:
        from ..utils.misc import extend_Trues
        ret = extend_Trues(ret,axis = 0,ext_frac = 0,ext_add = ext_add)

    print("Found :D")
    return ret

# Cell
def mask_freq_rfi(data,freq,rtype = 'long-time',**kwargs):
    if rtype == 'long-time':
        ret = mask_nr(data,freq, **kwargs)

    print("Finish")
    return ret

def mask_thr_rfi(T2p, freq, s_sigma = 5, rms_frange = None, times=6, ext_add=0, ext_frac=0):
    """
    Mask RFI based on threshold detection after Gaussian smoothing.
    
    Parameters
    ----------
    T2p : array
        Array of shape (time, freq, 2); T or flux with two polarizations sorted by time.
    freq : array
        Frequency array corresponding to the second axis of T2p.
    s_sigma : float, optional
        Sigma for Gaussian smoothing along frequency axis. Default is 5.
    rms_frange : list or None, optional
        Frequency range [fmin, fmax] for RMS calculation. If None, use all frequencies.
    times : float, optional
        Threshold multiplier. RFI is identified where T > RMS * times. Default is 6.
    ext_add : int, optional
        Number of additional channels to extend on each side of detected RFI. Default is 0.
    ext_frac : float, optional
        Fractional extension of RFI regions. Default is 0.
    
    Returns
    -------
    is_rfi : ndarray
        Boolean array of shape (time, freq) indicating RFI positions.
    
    """
    from scipy import ndimage
    from ..utils.misc import extend_Trues
    
    # import time  # Add this at the top of the file if not already imported
    # start_time = time.time()  # Start timing

    T = np.nanmean(T2p, axis = 2) # shape (time, freq)
    
    from scipy import ndimage
    T = ndimage.gaussian_filter1d(T, s_sigma, axis=1) # smooth along freq axis

    RMS = rms(T,freq,rms_frange)[:, None]

    # find rfi
    is_rfi = (T > RMS * times) 
    
    if ext_add > 0 or ext_frac > 0:
        is_rfi = extend_Trues(is_rfi, axis=1, ext_add=ext_add, ext_frac=ext_frac)

    # end_time = time.time()  # End timing
    # print(f"mask_thr_rfi execution time: {end_time - start_time:.2f} seconds")

    return is_rfi
