from rtlsdr import RtlSdr
from contextlib import closing
from matplotlib import pyplot as plt
import numpy as np
from scipy.signal import spectrogram, windows
from scipy import signal
from skimage.io import imsave, imread
from datetime import datetime
import json
import os
from tqdm import tqdm
import time
from queue import Queue
import asyncio
from pathlib import Path
import warnings

for cat in [RuntimeWarning, UserWarning, FutureWarning]:
    warnings.filterwarnings("ignore", category=cat) 


def split_images(dir="sdr_captures/specs_raw"):
    jpgs = list(Path(dir).rglob('*.jpg'))
    pngs = list(Path(dir).rglob('*.png'))

    img_files = pngs + jpgs
    img_files = list(filter(lambda x: 'chk' not in str(x), img_files))
    for img_file in tqdm(img_files, desc="splitting images"):
        im = imread(img_file)
        shp = list(im.shape)
        shp = list(filter(lambda x: x != 1, shp))
        shp = np.array(shp)
        dim_to_slide_over = shp.argmax()
        win_size = shp[shp.argmin()]
        im_size = shp[dim_to_slide_over]
        for start in range(0, im_size, win_size):
            stop = start + win_size
            if stop >= im_size:
                break
            if dim_to_slide_over == 0:
                chunk = im[start:stop, :]
            elif dim_to_slide_over == 1:
                chunk = im[:, start:stop]
            file_out = str(
                Path(img_file).with_suffix(f".chk_{start}_{stop}.png"))
            imsave(file_out, chunk)


# y -- spectrogram, nf by nt array
# dbf -- Dynamic range of the spectrum

def adjust_dyn_range(x, mx=3, mn=10, rel_to=np.median):
    r = rel_to(x)
    zmax = r+mx
    zmin = r-mn
    x[x<zmin] = zmin
    x[x>zmax] = zmax
    return x


def to_spec(y, fs, fc, NFFT=1024, dbf=60, nperseg=128, normalize=True):

    #w = windows.hamming(nperseg)
    #window = signal.kaiser(nperseg, beta=14)
    
    f, t, y = spectrogram(y, detrend=None, noverlap=int(nperseg/2), nfft=NFFT, fs=fs)

    y = np.fft.fftshift(y, axes=0)
    if normalize:
        #y = norm_spectrum(y)
        y = np.sqrt(np.power(y.real, 2) + np.power(y.imag, 2))
        y = 20 * np.log10(np.abs(y)/ np.abs(y).max())
        y = np.abs(y)
        y = y / y.max()

    
    return y

from sklearn.preprocessing import MinMaxScaler, StandardScaler

def spectrogram(x, fs, fc, m=None, dbf=60):

    if not m:
        m = 1024
    
    isreal_bool = np.isreal(x).all()
    
    lx = len(x);
    nt = (lx + m - 1) // m
    x = np.append(x,np.zeros(-lx+nt*m))
    x = x.reshape((int(m/2),nt*2), order='F')
    x = np.concatenate((x,x),axis=0)
    x = x.reshape((m*nt*2,1),order='F')
    x = x[np.r_[m//2:len(x),np.ones(m//2)*(len(x)-1)].astype(int)].reshape((m,nt*2),order='F')
    xmw = x * windows.hamming(m)[:,None]
    t_range = [0.0, lx / fs]
    if isreal_bool:
        f_range = [ fc, fs / 2.0 + fc]
        xmf = np.fft.fft(xmw,len(xmw),axis=0)
        xmf = xmf[0:m/2,:]
    else:
        f_range = [-fs / 2.0 + fc, fs / 2.0 + fc]
        xmf = np.fft.fftshift( np.fft.fft( xmw ,len(xmw),axis=0), axes=0 )
    f_range = np.linspace(*f_range, xmf.shape[0])
    t_range = np.linspace(*t_range, xmf.shape[1])

    h = xmf.shape[0]
    each = int(h*.10)
    xmf = xmf[each:-each, :]

    xmf = np.sqrt(np.power(xmf.real, 2) + np.power(xmf.imag, 2))
    xmf = np.abs(xmf)
    
    xmf /= xmf.max()

    #throw away sides

    xmf = 20 * np.log10(xmf)
    xmf = np.clip(xmf, -dbf, 0)
    xmf = MinMaxScaler().fit_transform(StandardScaler(with_mean=True, with_std=True).fit_transform(xmf))
    xmf = np.abs(xmf)
    #xmf-=np.median(xmf)
    xmf/=xmf.max()

    
    print(xmf.min(), xmf.max())
    return f_range, t_range, xmf


def append_json(data, path):
    with open(path, 'a') as f:
        f.write(json.dumps(data) + '\n')


async def stream(sdr, N):
    samples_buffer = Queue()
    total = 0
    with tqdm(total=N, desc='sampling') as pbar:
        #for i in range(10):
        #    time.sleep(0.1)

        async for samples in sdr.stream():
            # do something with samples
            # ...

            samples_buffer.put(samples)
            #print(f'put {len(samples)} into buffer')
            total += len(samples)
            pbar.update(len(samples))
            if total >= N:
                break
        # to stop streaming:
        await sdr.stop()

    # done
    sdr.close()
    return samples_buffer


def capture(fc=94.3e6,
            fs=int(1e6),
            gain='auto',
            seconds_dwell=.4
            #offset_dc=5e4
            ):

    
    N = int(seconds_dwell * fs)
    with closing(RtlSdr()) as sdr:
        sdr.sample_rate = fs
        sdr.center_freq = fc# + int(offset_dc)
        sdr.gain = gain
        t = datetime.now()
        stamp = datetime.timestamp(t)

        loop = asyncio.get_event_loop()
        samples_buffer = loop.run_until_complete(stream(sdr, N))

    iq_samples = np.hstack(np.array(list(samples_buffer.queue)))[:N].astype("complex64")
    #iq_samples = shift_mix(iq_samples, -offset_dc, fs)
    #path = os.path.join(out_dir, f'{stamp}.png')
    meta = dict(
        fs=fs,
        fc=fc,
        gain=gain,
        seconds_dwell=seconds_dwell,
        dt_start=stamp
    )



    return iq_samples, meta

def shift_mix(x, hz, fs):
    return x*np.exp(1j*2*np.pi*hz/fs*np.arange(len(x)))


def save_capture(path, spec_img, meta, meta_path):
    imsave(path, spec_img.T)
    append_json(meta, meta_path)


def scan(
        low=80e6,
        high=1000e6,
        repeats=10,
        target_hpb=300,
    ):

    out_dir="sdr_captures/specs_raw"
    meta_path="sdr_captures/dataset.json"

    os.makedirs(out_dir, exist_ok=True)
    for repeat in tqdm(range(repeats), desc='repeats'):
        for fs in [int(3.2e6)]:#list(map(int, (3.2e6, 2e6, 1e6))):
            #for NFFT in [1024, 2048, 2048 * 2]:

            fcs = []
            fc = low
            while fc < high:
                fc += int((fs * (1/3.)))
                fcs.append(fc)
            fcs = np.array(fcs)
            print(f'scanning {len(fcs)} total frequencies...')

            for fc in tqdm(fcs, desc='fcs'):
                try:
                    iq, meta = capture(fc=fc, fs=fs)
                    meta['NFFT'] = closest_power_of_two(fs / target_hpb)
                    meta['hpb'] = fs/meta['NFFT']
                    ff, tt, spec_img = spectrogram(iq, fs, fc, m=meta['NFFT'])
                    
                    img_path = os.path.join(out_dir, f"{meta['dt_start']}.png")
                    save_capture(img_path, spec_img, meta, meta_path)
            


                except Exception as e:
                    print(e)
                    time.sleep(1)
                    pass


def get_optimal_fs(max_fs=3e6):
    fss = np.array([np.power(2,i) for i in range(30)])
    fss = fss[fss<=max_fs][-1]
    return fss


def optimal_scan(
        min_freq=80e6,
        max_freq=107e6,
        fs=3e6,
        hpb_target=4096

    ):
    fs2 = get_optimal_fs(fs)
    if fs2!=fs:
        print(f'optimal fs found: {fs2}, original: {fs}')
        fs = fs2
        del fs2

    n_bins = closest_power_of_two(fs / hpb_target)
    print(f'given hz per bin target: {hpb_target} -> nfft bins per sweep: {n_bins}')
    assert fs == hpb_target * n_bins
    print(f'{fs} = {hpb_target} * {n_bins}')
    
    diff_bw = max_freq-min_freq
    sweeps = np.ceil(diff_bw/fs) + 1
    sweep_bw = sweeps * fs
    delta_bw = sweep_bw - diff_bw
    adjusted_min_freq = min_freq - int(delta_bw//2)
    adjusted_max_freq = max_freq + int(delta_bw//2)
    assert (adjusted_max_freq-adjusted_min_freq) == sweep_bw
    print(f'optimal min/max frequecies: {adjusted_min_freq}/{adjusted_max_freq}')
    min_freq = adjusted_min_freq
    max_freq = adjusted_max_freq

    freq_bins = np.arange(n_bins*sweeps)
    fz = np.arange(min_freq, max_freq, hpb_target).astype(int)
    return freq_bins, fz

def closest_power_of_two(number):
    # Returns next power of two following 'number'
    n = np.ceil(np.log2(number))
    a = np.array([np.power(2, n - 1), np.power(2, n), np.power(2, n + 1)])
    return int(a[np.argmin(np.abs(a - number))])


def norm_spectrum(spec_img):
    spec_img = 20 * np.log10(np.abs(spec_img) / np.max(np.abs(spec_img)))

    mid = np.median(spec_img)
    # high = mid + 30
    # low = mid - 30
    # spec_img[spec_img < low] = low
    # spec_img[spec_img > high] = high

    spec_img = np.abs(spec_img)
    spec_img /= spec_img.max()
    print('spec max:', spec_img.max(), 'spec min:', spec_img.min())
    return spec_img

def parse_measure(s):
    s = s.lower()
    if s[-1].isalpha():
        h, mod = float(s[:-1]), s[-1]
        if mod == 'm':
            h*=1e6
        elif mod == 'k':
            h*=1e3
    else:
        h = int(s)
    return h

def string_to_linspace(s, delim=':'):
    return np.arange(*list(map(parse_measure, s.split(delim))))

#string_to_linspace('24M:28M:3M')


def plot_one(fc=94.3 * 1e6, fs=3e6, target_hpb=300, seconds_dwell=.2):

    NFFT = closest_power_of_two(fs / target_hpb)
    iq_samples, meta = capture(fc=fc, fs=fs, seconds_dwell=seconds_dwell)
    spec_img = to_spec(iq_samples, fs, fc, NFFT=NFFT)
    #spec_img = norm_spectrum(spec_img)
    #spec_img = np.abs(spec_img)
    #spec_img /= spec_img.max()

    #spec_img = 1 - spec_img
    print('img shape:', spec_img.shape)
    fig, ax = plt.subplots(1, 1, figsize=(14, 4))
    ax.matshow(spec_img.T[:NFFT], cmap=plt.get_cmap('viridis'))
    print(spec_img.T.shape)
    #Wplt.plot(spec_img.T[0, :])
    plt.show()


if __name__ == "__main__":
    #split_images()
    #plot_one()
    scan(repeats=3, target_hpb=1500)
    split_images()
    #plot_one()
[((19, 4, 19, 51), 'warnings.filterwarnings', 'warnings.filterwarnings', (), '', False, 'import warnings\n'), ((28, 20, 28, 60), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm\n'), ((68, 8, 68, 34), 'numpy.fft.fftshift', 'np.fft.fftshift', (), '', True, 'import numpy as np\n'), ((92, 8, 92, 36), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((104, 14, 104, 49), 'numpy.linspace', 'np.linspace', ({(104, 26, 104, 34): '*f_range', (104, 36, 104, 48): 'xmf.shape[0]'}, {}), '(*f_range, xmf.shape[0])', True, 'import numpy as np\n'), ((105, 14, 105, 49), 'numpy.linspace', 'np.linspace', ({(105, 26, 105, 34): '*t_range', (105, 36, 105, 48): 'xmf.shape[1]'}, {}), '(*t_range, xmf.shape[1])', True, 'import numpy as np\n'), ((112, 10, 112, 21), 'numpy.abs', 'np.abs', ({(112, 17, 112, 20): 'xmf'}, {}), '(xmf)', True, 'import numpy as np\n'), ((119, 10, 119, 31), 'numpy.clip', 'np.clip', ({(119, 18, 119, 21): 'xmf', (119, 23, 119, 27): '-dbf', (119, 29, 119, 30): '0'}, {}), '(xmf, -dbf, 0)', True, 'import numpy as np\n'), ((121, 10, 121, 21), 'numpy.abs', 'np.abs', ({(121, 17, 121, 20): 'xmf'}, {}), '(xmf)', True, 'import numpy as np\n'), ((136, 21, 136, 28), 'queue.Queue', 'Queue', ({}, {}), '()', False, 'from queue import Queue\n'), ((199, 4, 199, 28), 'skimage.io.imsave', 'imsave', ({(199, 11, 199, 15): 'path', (199, 17, 199, 27): 'spec_img.T'}, {}), '(path, spec_img.T)', False, 'from skimage.io import imsave, imread\n'), ((213, 4, 213, 39), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((279, 16, 279, 40), 'numpy.arange', 'np.arange', ({(279, 26, 279, 39): 'n_bins * sweeps'}, {}), '(n_bins * sweeps)', True, 'import numpy as np\n'), ((293, 10, 293, 29), 'numpy.median', 'np.median', ({(293, 20, 293, 28): 'spec_img'}, {}), '(spec_img)', True, 'import numpy as np\n'), ((299, 15, 299, 31), 'numpy.abs', 'np.abs', ({(299, 22, 299, 30): 'spec_img'}, {}), '(spec_img)', True, 'import numpy as np\n'), ((333, 14, 333, 49), 'matplotlib.pyplot.subplots', 'plt.subplots', (), '', True, 'from matplotlib import pyplot as plt\n'), ((337, 4, 337, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'from matplotlib import pyplot as plt\n'), ((29, 13, 29, 29), 'skimage.io.imread', 'imread', ({(29, 20, 29, 28): 'img_file'}, {}), '(img_file)', False, 'from skimage.io import imsave, imread\n'), ((32, 14, 32, 27), 'numpy.array', 'np.array', ({(32, 23, 32, 26): 'shp'}, {}), '(shp)', True, 'import numpy as np\n'), ((73, 12, 73, 21), 'numpy.abs', 'np.abs', ({(73, 19, 73, 20): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((90, 20, 90, 38), 'numpy.zeros', 'np.zeros', ({(90, 29, 90, 37): '-lx + nt * m'}, {}), '(-lx + nt * m)', True, 'import numpy as np\n'), ((118, 15, 118, 28), 'numpy.log10', 'np.log10', ({(118, 24, 118, 27): 'xmf'}, {}), '(xmf)', True, 'import numpy as np\n'), ((138, 9, 138, 39), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm\n'), ((173, 12, 173, 26), 'datetime.datetime.now', 'datetime.now', ({}, {}), '()', False, 'from datetime import datetime\n'), ((174, 16, 174, 37), 'datetime.datetime.timestamp', 'datetime.timestamp', ({(174, 35, 174, 36): 't'}, {}), '(t)', False, 'from datetime import datetime\n'), ((176, 15, 176, 39), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ({}, {}), '()', False, 'import asyncio\n'), ((269, 13, 269, 32), 'numpy.ceil', 'np.ceil', ({(269, 21, 269, 31): '(diff_bw / fs)'}, {}), '(diff_bw / fs)', True, 'import numpy as np\n'), ((285, 16, 285, 31), 'numpy.log2', 'np.log2', ({(285, 24, 285, 30): 'number'}, {}), '(number)', True, 'import numpy as np\n'), ((46, 12, 46, 35), 'skimage.io.imsave', 'imsave', ({(46, 19, 46, 27): 'file_out', (46, 29, 46, 34): 'chunk'}, {}), '(file_out, chunk)', False, 'from skimage.io import imsave, imread\n'), ((86, 18, 86, 30), 'numpy.isreal', 'np.isreal', ({(86, 28, 86, 29): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((95, 14, 95, 32), 'scipy.signal.windows.hamming', 'windows.hamming', ({(95, 30, 95, 31): 'm'}, {}), '(m)', False, 'from scipy.signal import spectrogram, windows\n'), ((111, 18, 111, 39), 'numpy.power', 'np.power', ({(111, 27, 111, 35): 'xmf.real', (111, 37, 111, 38): '2'}, {}), '(xmf.real, 2)', True, 'import numpy as np\n'), ((111, 42, 111, 63), 'numpy.power', 'np.power', ({(111, 51, 111, 59): 'xmf.imag', (111, 61, 111, 62): '2'}, {}), '(xmf.imag, 2)', True, 'import numpy as np\n'), ((120, 10, 120, 24), 'sklearn.preprocessing.MinMaxScaler', 'MinMaxScaler', ({}, {}), '()', False, 'from sklearn.preprocessing import MinMaxScaler, StandardScaler\n'), ((169, 17, 169, 25), 'rtlsdr.RtlSdr', 'RtlSdr', ({}, {}), '()', False, 'from rtlsdr import RtlSdr\n'), ((223, 18, 223, 31), 'numpy.array', 'np.array', ({(223, 27, 223, 30): 'fcs'}, {}), '(fcs)', True, 'import numpy as np\n'), ((226, 22, 226, 43), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm\n'), ((245, 20, 245, 33), 'numpy.power', 'np.power', ({(245, 29, 245, 30): '2', (245, 31, 245, 32): 'i'}, {}), '(2, i)', True, 'import numpy as np\n'), ((280, 9, 280, 50), 'numpy.arange', 'np.arange', ({(280, 19, 280, 27): 'min_freq', (280, 29, 280, 37): 'max_freq', (280, 39, 280, 49): 'hpb_target'}, {}), '(min_freq, max_freq, hpb_target)', True, 'import numpy as np\n'), ((286, 18, 286, 36), 'numpy.power', 'np.power', ({(286, 27, 286, 28): '2', (286, 30, 286, 35): 'n - 1'}, {}), '(2, n - 1)', True, 'import numpy as np\n'), ((286, 38, 286, 52), 'numpy.power', 'np.power', ({(286, 47, 286, 48): '2', (286, 50, 286, 51): 'n'}, {}), '(2, n)', True, 'import numpy as np\n'), ((286, 54, 286, 72), 'numpy.power', 'np.power', ({(286, 63, 286, 64): '2', (286, 66, 286, 71): 'n + 1'}, {}), '(2, n + 1)', True, 'import numpy as np\n'), ((334, 39, 334, 62), 'matplotlib.pyplot.get_cmap', 'plt.get_cmap', ({(334, 52, 334, 61): '"""viridis"""'}, {}), "('viridis')", True, 'from matplotlib import pyplot as plt\n'), ((23, 16, 23, 25), 'pathlib.Path', 'Path', ({(23, 21, 23, 24): 'dir'}, {}), '(dir)', False, 'from pathlib import Path\n'), ((24, 16, 24, 25), 'pathlib.Path', 'Path', ({(24, 21, 24, 24): 'dir'}, {}), '(dir)', False, 'from pathlib import Path\n'), ((71, 20, 71, 39), 'numpy.power', 'np.power', ({(71, 29, 71, 35): 'y.real', (71, 37, 71, 38): '2'}, {}), '(y.real, 2)', True, 'import numpy as np\n'), ((71, 42, 71, 61), 'numpy.power', 'np.power', ({(71, 51, 71, 57): 'y.imag', (71, 59, 71, 60): '2'}, {}), '(y.imag, 2)', True, 'import numpy as np\n'), ((120, 39, 120, 84), 'sklearn.preprocessing.StandardScaler', 'StandardScaler', (), '', False, 'from sklearn.preprocessing import MinMaxScaler, StandardScaler\n'), ((132, 16, 132, 32), 'json.dumps', 'json.dumps', ({(132, 27, 132, 31): 'data'}, {}), '(data)', False, 'import json\n'), ((291, 29, 291, 45), 'numpy.abs', 'np.abs', ({(291, 36, 291, 44): 'spec_img'}, {}), '(spec_img)', True, 'import numpy as np\n'), ((72, 26, 72, 35), 'numpy.abs', 'np.abs', ({(72, 33, 72, 34): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((231, 39, 231, 78), 'scipy.signal.spectrogram', 'spectrogram', (), '', False, 'from scipy.signal import spectrogram, windows\n'), ((233, 31, 233, 79), 'os.path.join', 'os.path.join', ({(233, 44, 233, 51): 'out_dir', (233, 53, 233, 78): 'f"""{meta[\'dt_start\']}.png"""'}, {}), '(out_dir, f"{meta[\'dt_start\']}.png")', False, 'import os\n'), ((287, 27, 287, 45), 'numpy.abs', 'np.abs', ({(287, 34, 287, 44): '(a - number)'}, {}), '(a - number)', True, 'import numpy as np\n'), ((291, 55, 291, 71), 'numpy.abs', 'np.abs', ({(291, 62, 291, 70): 'spec_img'}, {}), '(spec_img)', True, 'import numpy as np\n'), ((45, 16, 45, 30), 'pathlib.Path', 'Path', ({(45, 21, 45, 29): 'img_file'}, {}), '(img_file)', False, 'from pathlib import Path\n'), ((240, 20, 240, 33), 'time.sleep', 'time.sleep', ({(240, 31, 240, 32): '(1)'}, {}), '(1)', False, 'import time\n'), ((72, 37, 72, 46), 'numpy.abs', 'np.abs', ({(72, 44, 72, 45): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((94, 28, 94, 41), 'numpy.ones', 'np.ones', ({(94, 36, 94, 40): 'm // 2'}, {}), '(m // 2)', True, 'import numpy as np\n')]