# Based on this article:
# https://medium.com/swlh/decoding-noaa-satellite-images-using-50-lines-of-code-3c5d1d0a08da

import scipy.io.wavfile as wav
import scipy.signal as signal
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from tqdm import tqdm

class NOAADecoder:
    def __init__(self, iq_data=None, filename=None) -> None:
        # Check if iq_data or filename are set correctly
        if iq_data and filename:
            raise ValueError("Both iq_data and filename are set.")
        if not iq_data and not filename:
            raise ValueError("Neither iq_data nor filename are set.")
        
        # Set iq_data or filename
        if filename:
            self.filename=filename
            self.sample_rate, self.data = wav.read(f'{self.filename}')
        if iq_data:
            self.sample_rate = 2048000
            self.data = self.iq_data

    def hilbert(self, data):
        # Hilbert transform
        analytical_signal = signal.hilbert(data)
        amplitude_envelope = np.abs(analytical_signal)
        return amplitude_envelope

    def decode(self, resample=4):

        # Resample data
        data = self.data[::resample]
        sample_rate = self.sample_rate//resample
        
        # Demodulate
        data_am = self.hilbert(data)
        frame_width = int(0.5*sample_rate)
        
        # Image processing
        w, h = frame_width, data_am.shape[0]//frame_width
        image = Image.new('L', (w, h))
        px, py = 0, 0
        
        for p in tqdm(range(data_am.shape[0])):
            if data_am[p][0] != data_am[p][1]:
                raise BufferError("Mismatch in data array.")
            lum = int(data_am[p][0]//32 - 32)
            if lum < 0: lum = 0
            if lum > 255: lum = 255
            image.putpixel((px, py),lum)
            px += 1
            if px >= w:
                px = 0
                py += 1
                if py >= h:
                    break
        image = image.resize((w, 4*h))
        plt.imshow(image, cmap='gray')
        # Save image
        image.save('noaa.png')
        return